Maven/internal/store/backfill.go

package store

import (
	"context"
	"encoding/json"
	"fmt"
	"time"
)

// EmbedFunc embeds one piece of stored text. The caller passes
// router.EmbedPassage — the STORE side of the query/passage asymmetry, which is
// the side every vector in the DB was written with. (Passing the query side
// would put the stored vectors in the wrong half of the space and quietly halve
// recall.) A func instead of an interface keeps this package free of any
// dependency on internal/router.
type EmbedFunc func(ctx context.Context, text string) ([]float32, error)

// BackfillResult is what the re-embed run did, for logging.
type BackfillResult struct {
	Skipped  bool // marker already matched — nothing to do
	Notes    int  // rows rewritten in the notes table
	Facts    int  // fact rows rewritten in memory_vectors
	MemNotes int  // note rows rewritten in memory_vectors
	NoText   int  // memory_vectors rows with no text in their meta, left alone
	Took     time.Duration
}

// ReembedAll rewrites every stored vector with the currently configured
// embedder and then records that embedder as the one that owns the DB.
//
// Both places a vector lives are rewritten in the same pass: the `notes` table
// `embedding` column and the `memory_vectors` rows (notes AND facts). Doing
// only one would leave the two indexes disagreeing, which is worse than leaving
// both stale.
//
// Safe to re-run: if the marker already names the current embedder there is
// nothing to fix, so it returns immediately with Skipped set.
//
// Crash safety: everything — every vector and the marker — happens inside one
// transaction. If anything fails or the process dies partway, the transaction
// rolls back: no vectors changed and no marker written, so the next run does
// the whole job again. The marker is never set unless the full rewrite
// committed.
func (s *Store) ReembedAll(ctx context.Context, currentID string, embed EmbedFunc) (BackfillResult, error) {
	start := time.Now()
	var res BackfillResult

	stored, err := s.Meta(ctx, metaKeyEmbedderID)
	if err != nil {
		return res, err
	}
	if stored == currentID {
		res.Skipped = true
		res.Took = time.Since(start)
		return res, nil
	}

	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return res, fmt.Errorf("reembed: begin: %w", err)
	}
	defer tx.Rollback() // no-op once committed

	// ----- notes table -----
	type noteRow struct {
		id   int64
		text string
	}
	var notes []noteRow
	rows, err := tx.QueryContext(ctx, `SELECT id, text FROM notes WHERE text != ''`)
	if err != nil {
		return res, fmt.Errorf("reembed: read notes: %w", err)
	}
	for rows.Next() {
		var n noteRow
		if err := rows.Scan(&n.id, &n.text); err != nil {
			rows.Close()
			return res, fmt.Errorf("reembed: note row: %w", err)
		}
		notes = append(notes, n)
	}
	rows.Close()
	if err := rows.Err(); err != nil {
		return res, fmt.Errorf("reembed: notes: %w", err)
	}

	for _, n := range notes {
		vec, err := embed(ctx, n.text)
		if err != nil {
			return res, fmt.Errorf("reembed: embed note %d: %w", n.id, err)
		}
		if _, err := tx.ExecContext(ctx,
			`UPDATE notes SET embedding = ? WHERE id = ?`, floatsToBlob(vec), n.id); err != nil {
			return res, fmt.Errorf("reembed: write note %d: %w", n.id, err)
		}
		res.Notes++
	}

	// ----- memory_vectors (the unified index: notes AND facts) -----
	// The text to re-embed is the one carried in the row's meta blob, which is
	// exactly the text that was embedded when the row was written.
	type vecRow struct {
		id, text, kind string
	}
	var vecs []vecRow
	rows, err = tx.QueryContext(ctx, `SELECT id, meta FROM memory_vectors`)
	if err != nil {
		return res, fmt.Errorf("reembed: read memory vectors: %w", err)
	}
	for rows.Next() {
		var id, metaJSON string
		if err := rows.Scan(&id, &metaJSON); err != nil {
			rows.Close()
			return res, fmt.Errorf("reembed: memory row: %w", err)
		}
		meta := map[string]string{}
		if err := json.Unmarshal([]byte(metaJSON), &meta); err != nil {
			rows.Close()
			return res, fmt.Errorf("reembed: meta for %q: %w", id, err)
		}
		if meta["text"] == "" {
			res.NoText++
			continue
		}
		vecs = append(vecs, vecRow{id: id, text: meta["text"], kind: meta["type"]})
	}
	rows.Close()
	if err := rows.Err(); err != nil {
		return res, fmt.Errorf("reembed: memory vectors: %w", err)
	}

	for _, v := range vecs {
		vec, err := embed(ctx, v.text)
		if err != nil {
			return res, fmt.Errorf("reembed: embed %q: %w", v.id, err)
		}
		if _, err := tx.ExecContext(ctx,
			`UPDATE memory_vectors SET vec = ? WHERE id = ?`, encodeVec(vec), v.id); err != nil {
			return res, fmt.Errorf("reembed: write %q: %w", v.id, err)
		}
		if v.kind == "fact" {
			res.Facts++
		} else {
			res.MemNotes++
		}
	}

	// Same transaction as the rewrite, on purpose: the marker can only exist if
	// every vector above was written.
	if _, err := tx.ExecContext(ctx,
		`INSERT INTO meta (key, value) VALUES (?,?)
		 ON CONFLICT(key) DO UPDATE SET value = excluded.value`,
		metaKeyEmbedderID, currentID); err != nil {
		return res, fmt.Errorf("reembed: write marker: %w", err)
	}
	if err := tx.Commit(); err != nil {
		return res, fmt.Errorf("reembed: commit: %w", err)
	}
	res.Took = time.Since(start)
	return res, nil
}