story/internal/memoryindex/processor.go

package memoryindex

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"

	"data-terminal-cloud/internal/embeddingprovider"
	"data-terminal-cloud/internal/outbox"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

var ErrInvalidIndexEvent = errors.New("invalid memory index event")
var ErrMemoryChangedDuringIndex = errors.New("memory changed while embedding was generated")

type Processor struct {
	pool     *pgxpool.Pool
	provider embeddingprovider.Provider
}

type indexSnapshot struct {
	content, status, currentVersionID, sensitivity string
	consentEffective                               bool
}

func NewProcessor(pool *pgxpool.Pool, providers ...embeddingprovider.Provider) (*Processor, error) {
	if pool == nil {
		return nil, errors.New("database pool is required")
	}
	var provider embeddingprovider.Provider = embeddingprovider.Local{}
	if len(providers) > 1 || (len(providers) == 1 && providers[0] == nil) {
		return nil, errors.New("one valid embedding provider is required")
	}
	if len(providers) == 1 {
		provider = providers[0]
	}
	return &Processor{pool: pool, provider: provider}, nil
}

func (p *Processor) Handle(ctx context.Context, event outbox.Event) error {
	var payload struct {
		MemoryVersionID string `json:"memory_version_id"`
	}
	if err := json.Unmarshal(event.Payload, &payload); err != nil || strings.TrimSpace(payload.MemoryVersionID) == "" {
		return ErrInvalidIndexEvent
	}

	before, err := loadIndexSnapshot(ctx, p.pool, payload.MemoryVersionID, false)
	if err != nil {
		return err
	}
	if !eligible(before, payload.MemoryVersionID) {
		return p.markStale(ctx, payload.MemoryVersionID, "memory is not currently eligible")
	}
	result, err := p.provider.Embed(ctx, before.content)
	if err != nil {
		return err
	}
	literal, err := embeddingprovider.Literal(result.Vector)
	if err != nil {
		return err
	}

	// Embedding is generated without an open transaction. Only the short
	// revalidation and vector update below hold memory/index row locks.
	tx, err := p.pool.BeginTx(ctx, pgx.TxOptions{})
	if err != nil {
		return err
	}
	defer tx.Rollback(context.Background())
	current, err := loadIndexSnapshot(ctx, tx, payload.MemoryVersionID, true)
	if err != nil {
		return err
	}
	if !eligible(current, payload.MemoryVersionID) {
		if _, err = tx.Exec(ctx, `UPDATE memory_index_entry SET embedding=NULL,status='STALE',last_attempt_at=now(),last_error='memory is not currently eligible',updated_at=now() WHERE memory_version_id=$1::uuid AND index_type='pgvector'`, payload.MemoryVersionID); err != nil {
			return err
		}
		return tx.Commit(ctx)
	}
	if current.content != before.content {
		return outbox.Retryable(ErrMemoryChangedDuringIndex)
	}
	if _, err = tx.Exec(ctx, `
		UPDATE memory_index_entry SET embedding=$2::vector,status='ACTIVE',provider_ref=$3,
			indexed_at=now(),last_attempt_at=now(),last_error=NULL,updated_at=now()
		WHERE memory_version_id=$1::uuid AND index_type='pgvector'`,
		payload.MemoryVersionID, literal, result.ProviderRef); err != nil {
		return fmt.Errorf("activate memory index entry: %w", err)
	}
	return tx.Commit(ctx)
}

type indexQueryer interface {
	QueryRow(context.Context, string, ...any) pgx.Row
}

func loadIndexSnapshot(ctx context.Context, q indexQueryer, memoryVersionID string, lock bool) (indexSnapshot, error) {
	lockClause := ""
	if lock {
		lockClause = " FOR UPDATE OF mi,cm,sp"
	}
	var snapshot indexSnapshot
	err := q.QueryRow(ctx, `
		SELECT mv.content,cm.status,COALESCE(cm.current_version_id::text,''),cm.sensitivity,
			EXISTS(SELECT 1 FROM consent_record cr
				WHERE cr.tenant_id=cm.tenant_id AND cr.subject_id=cm.subject_id
				AND cr.consent_type='personalized_memory' AND cr.withdrawn_at IS NULL
				AND (cr.expires_at IS NULL OR cr.expires_at>now()))
		FROM memory_index_entry mi
		JOIN memory_version mv ON mv.id=mi.memory_version_id
		JOIN curated_memory cm ON cm.id=mi.memory_id
		JOIN subject_profile sp ON sp.id=cm.subject_id AND sp.tenant_id=cm.tenant_id AND sp.status IN ('active','suspended')
		WHERE mi.memory_version_id=$1::uuid AND mi.index_type='pgvector'`+lockClause, memoryVersionID).
		Scan(&snapshot.content, &snapshot.status, &snapshot.currentVersionID, &snapshot.sensitivity, &snapshot.consentEffective)
	if errors.Is(err, pgx.ErrNoRows) {
		return indexSnapshot{}, ErrInvalidIndexEvent
	}
	if err != nil {
		return indexSnapshot{}, fmt.Errorf("load memory index entry: %w", err)
	}
	return snapshot, nil
}

func eligible(snapshot indexSnapshot, memoryVersionID string) bool {
	return snapshot.status == "ACTIVE" && snapshot.currentVersionID == memoryVersionID && snapshot.consentEffective && snapshot.sensitivity == "normal"
}

func (p *Processor) markStale(ctx context.Context, memoryVersionID, reason string) error {
	_, err := p.pool.Exec(ctx, `UPDATE memory_index_entry SET embedding=NULL,status='STALE',last_attempt_at=now(),last_error=$2,updated_at=now() WHERE memory_version_id=$1::uuid AND index_type='pgvector'`, memoryVersionID, reason)
	return err
}