From 13ac4b753c3dbf600a084a24a45d7cbec293e3f8 Mon Sep 17 00:00:00 2001 From: Quick <31828688+Quick104@users.noreply.github.com> Date: Fri, 26 Jun 2026 00:02:20 -0400 Subject: [PATCH] perf(recommendations): split embedding backfill into cheap and text-stale passes EmbedAll previously ran the full text-staleness CTE on every page, paying five item_people LATERAL joins per eligible row just to detect whether an item's canonical text had drifted - even when the real work was embedding brand-new (missing) items during active backfill. Restructure EmbedAll into two passes: - Pass 1 (cheap): drain missing/model-stale items via the repurposed ItemsNeedingEmbedding query (single LEFT JOIN, no LATERAL), paged by a content_id cursor so a failed/skipped item is retried next run instead of stalling the page. - Pass 2 (expensive): only once Pass 1 fully drains, run one bounded ListEmbeddingTextCandidates scan (LIMIT embeddingTextStaleQuotaPerRun=200) to re-embed text-drifted items. Re-embedding refreshes canonical_text, so handled rows drop out next run - no Pass 2 cursor needed. Coverage-first tradeoff (documented in EmbedAll): under steady state Pass 1 drains every run so text-stale items stay fresh; only under pathological continuous heavy ingest does Pass 2 get skipped, deliberately prioritizing covering new items over re-embedding changed ones. Supporting changes: - ItemsNeedingEmbedding gains an afterID cursor + ORDER BY; SQL extracted to buildItemsNeedingEmbeddingSQL for a cheap-shape unit test (no item_people / LATERAL). - Add an `embedder` interface seam (Engine.embClient) so EmbedAll is testable with a fake; *embeddings.Client still satisfies it. - Extract embedBatch to DRY both passes, preserving quota/billing early-return, single-item fallback, ensureEmbeddingLock-before-upsert, and skip-on-store-error. Co-Authored-By: Claude Opus 4.8 --- internal/recommendations/embed_all_test.go | 223 +++++++++++++++++++++ internal/recommendations/engine.go | 10 +- internal/recommendations/repo.go | 34 +++- internal/recommendations/repo_test.go | 126 ++++++++++++ internal/recommendations/similar.go | 165 ++++++++++++--- 5 files changed, 521 insertions(+), 37 deletions(-) create mode 100644 internal/recommendations/embed_all_test.go diff --git a/internal/recommendations/embed_all_test.go b/internal/recommendations/embed_all_test.go new file mode 100644 index 00000000..628f6105 --- /dev/null +++ b/internal/recommendations/embed_all_test.go @@ -0,0 +1,223 @@ +package recommendations + +import ( + "context" + "strings" + "testing" + + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/Silo-Server/silo-server/internal/catalog" + "github.com/Silo-Server/silo-server/internal/config" + "github.com/Silo-Server/silo-server/internal/models" + "github.com/Silo-Server/silo-server/internal/recommendations/embeddings" +) + +// recordingEmbedder is a fake embedder seam used to drive EmbedAll without a +// real embedding API. It records, in order, every text it is asked to embed so +// tests can assert the cheap pass drains before the expensive text-staleness +// pass runs. Returned vectors are deterministic and canonical-width so the +// embedding lock validates without a real model. +type recordingEmbedder struct { + calls [][]string // one entry per Embed call (a batch) + texts []string // flattened, in embed order +} + +func (r *recordingEmbedder) Embed(_ context.Context, texts []string) ([][]float32, error) { + batch := make([]string, len(texts)) + copy(batch, texts) + r.calls = append(r.calls, batch) + r.texts = append(r.texts, batch...) + + out := make([][]float32, len(texts)) + for i := range texts { + vec := make([]float32, CanonicalEmbeddingDimensions) + // A non-zero leading component keeps the vector distinguishable from the + // seeded zero vectors; the rest stay zero. + vec[0] = 0.5 + out[i] = vec + } + return out, nil +} + +var _ embedder = (*recordingEmbedder)(nil) + +// firstIndexContaining returns the index of the first recorded text containing +// needle, or -1. +func firstIndexContaining(texts []string, needle string) int { + for i, txt := range texts { + if strings.Contains(txt, needle) { + return i + } + } + return -1 +} + +// lastIndexContaining returns the index of the last recorded text containing +// needle, or -1. +func lastIndexContaining(texts []string, needle string) int { + idx := -1 + for i, txt := range texts { + if strings.Contains(txt, needle) { + idx = i + } + } + return idx +} + +func newEmbedAllTestEngine(t *testing.T, pool *pgxpool.Pool, fake embedder, model string) *Engine { + t.Helper() + return &Engine{ + repo: NewRepo(pool), + itemRepo: catalog.NewItemRepository(pool), + personRepo: catalog.NewPersonRepository(pool), + embClient: fake, + cfg: config.RecommendationsConfig{ + EmbeddingModel: model, + EmbeddingBaseURL: "http://embed.test", + }, + pool: pool, + } +} + +// TestEmbedAllCheapPassDrainsBeforeTextStalePass verifies the coverage-first +// ordering: items with no embedding or a stale model (cheap Pass 1) are all +// embedded before the expensive text-staleness CTE (Pass 2) re-embeds an item +// whose only change is its canonical text. +func TestEmbedAllCheapPassDrainsBeforeTextStalePass(t *testing.T) { + pool := newEngineTestPool(t) + snapshotEmbeddingLock(t, pool) + ctx := context.Background() + + const ( + prefix = "t7embed-" + currentModel = "model-current" + oldModel = "model-old" + ) + cleanupRecoMediaItems(t, pool, prefix) + // Start from a clean lock so the first embed establishes it for currentModel. + if _, err := pool.Exec(ctx, `DELETE FROM server_settings WHERE key = $1`, embeddingLockSettingKey); err != nil { + t.Fatalf("clear embedding lock: %v", err) + } + + // Two cheap rows: one missing entirely, one under an old model. Distinct + // titles let us find their generated text in the embedder's record. + missingID := prefix + "1-missing" + modelStaleID := prefix + "2-modelstale" + seedRecoMediaItemTitled(t, pool, missingID, "movie", "matched", "CHEAPMISS Origins") + seedRecoMediaItemTitled(t, pool, modelStaleID, "movie", "matched", "CHEAPMODEL Returns") + seedRecoEmbedding(t, pool, modelStaleID, oldModel, "old text") + + // One text-stale-only row: embedded under the CURRENT model, but its stored + // canonical_text does not match what BuildEmbeddingText now produces. Only + // the expensive CTE can detect this. + textStaleID := prefix + "3-textstale" + seedRecoMediaItemTitled(t, pool, textStaleID, "movie", "matched", "TEXTSTALE Reloaded") + seedRecoEmbedding(t, pool, textStaleID, currentModel, "deliberately stale canonical text") + + fake := &recordingEmbedder{} + e := newEmbedAllTestEngine(t, pool, fake, currentModel) + + n, err := e.EmbedAll(ctx) + if err != nil { + t.Fatalf("EmbedAll: %v", err) + } + + // All three should have been (re)embedded: 2 cheap + 1 text-stale. + if n < 3 { + t.Fatalf("EmbedAll embedded %d items, want >= 3 (2 cheap + 1 text-stale)", n) + } + + // End state: every row now carries the current model and matching text. + assertEmbeddingModel(t, pool, missingID, currentModel) + assertEmbeddingModel(t, pool, modelStaleID, currentModel) + assertEmbeddingModel(t, pool, textStaleID, currentModel) + + // Ordering: both cheap items must be embedded before the text-stale item, + // proving Pass 2 (the expensive CTE) only ran after Pass 1 drained. + lastCheap := lastIndexContaining(fake.texts, "CHEAP") + firstTextStale := firstIndexContaining(fake.texts, "TEXTSTALE") + if lastCheap < 0 { + t.Fatalf("cheap items were never embedded; recorded texts: %v", fake.texts) + } + if firstTextStale < 0 { + t.Fatalf("text-stale item was never embedded; recorded texts: %v", fake.texts) + } + if firstTextStale <= lastCheap { + t.Fatalf("text-stale item embedded at %d before cheap drained at %d; coverage-first ordering violated:\n%v", + firstTextStale, lastCheap, fake.texts) + } +} + +// TestEmbedAllSkipsTextStaleWhenAlreadyCurrent confirms Pass 2 does not waste an +// embed call on an item whose stored canonical_text already matches the freshly +// generated text (it is neither missing, model-stale, nor text-stale). +func TestEmbedAllSkipsAlreadyCurrentItems(t *testing.T) { + pool := newEngineTestPool(t) + snapshotEmbeddingLock(t, pool) + ctx := context.Background() + + const ( + prefix = "t7current-" + currentModel = "model-current" + ) + cleanupRecoMediaItems(t, pool, prefix) + if _, err := pool.Exec(ctx, `DELETE FROM server_settings WHERE key = $1`, embeddingLockSettingKey); err != nil { + t.Fatalf("clear embedding lock: %v", err) + } + + // Seed an item, then store the embedding using the SAME canonical text the + // engine will generate, so it is fully up to date. + currentID := prefix + "1-current" + seedRecoMediaItemTitled(t, pool, currentID, "movie", "matched", " alreadyCurrent") + items, err := catalog.NewItemRepository(pool).GetByIDs(ctx, []string{currentID}) + if err != nil || len(items) == 0 { + t.Fatalf("load seeded item: %v", err) + } + wantText := embeddings.BuildEmbeddingText(items[0]) + seedRecoEmbedding(t, pool, currentID, currentModel, wantText) + + fake := &recordingEmbedder{} + e := newEmbedAllTestEngine(t, pool, fake, currentModel) + + n, err := e.EmbedAll(ctx) + if err != nil { + t.Fatalf("EmbedAll: %v", err) + } + if n != 0 { + t.Fatalf("EmbedAll re-embedded %d up-to-date items, want 0", n) + } + if len(fake.calls) != 0 { + t.Fatalf("embedder was called %d times for an up-to-date item, want 0: %v", len(fake.calls), fake.texts) + } +} + +// seedRecoMediaItemTitled inserts an embed-eligible media item via +// ItemRepository.Upsert. Using the repository (rather than a hand-written +// INSERT) guarantees every non-nullable column the scan path reads is +// populated, so the round-trip through GetByIDs in EmbedAll succeeds. +func seedRecoMediaItemTitled(t *testing.T, pool *pgxpool.Pool, contentID, mediaType, status, title string) { + t.Helper() + repo := catalog.NewItemRepository(pool) + if err := repo.Upsert(context.Background(), &models.MediaItem{ + ContentID: contentID, + Type: mediaType, + Title: title, + Status: status, + }); err != nil { + t.Fatalf("seed media item %s: %v", contentID, err) + } +} + +func assertEmbeddingModel(t *testing.T, pool *pgxpool.Pool, contentID, wantModel string) { + t.Helper() + var model string + err := pool.QueryRow(context.Background(), + `SELECT model FROM media_item_embeddings WHERE media_item_id = $1`, contentID).Scan(&model) + if err != nil { + t.Fatalf("read embedding model for %s: %v", contentID, err) + } + if model != wantModel { + t.Fatalf("embedding for %s has model %q, want %q", contentID, model, wantModel) + } +} diff --git a/internal/recommendations/engine.go b/internal/recommendations/engine.go index d2e138a0..3221b777 100644 --- a/internal/recommendations/engine.go +++ b/internal/recommendations/engine.go @@ -11,6 +11,14 @@ import ( "github.com/Silo-Server/silo-server/internal/userstore" ) +// embedder is the minimal embedding-client seam the Engine depends on. The +// concrete *embeddings.Client satisfies it; tests substitute a fake so the +// backfill loop (EmbedAll) and query-vector path can run without a real +// embedding API. +type embedder interface { + Embed(ctx context.Context, texts []string) ([][]float32, error) +} + // Engine implements the Recommender interface. type Engine struct { repo *Repo @@ -19,7 +27,7 @@ type Engine struct { personRepo *catalog.PersonRepository storeProvider userstore.UserStoreProvider signals *SignalReader - embClient *embeddings.Client + embClient embedder cfg config.RecommendationsConfig pool *pgxpool.Pool } diff --git a/internal/recommendations/repo.go b/internal/recommendations/repo.go index 86b37503..f1fd4499 100644 --- a/internal/recommendations/repo.go +++ b/internal/recommendations/repo.go @@ -463,20 +463,38 @@ func (r *Repo) findTasteProfileCandidates( return items, genreMap, nil } -// ItemsNeedingEmbedding returns content IDs of media items that either have no -// embedding or whose embedding was generated with a different model. -// Books bypass the status='matched' gate because their scanner/plugin-derived -// metadata is authoritative as soon as the scan/enrichment completes. -func (r *Repo) ItemsNeedingEmbedding(ctx context.Context, currentModel string, limit int) ([]string, error) { - query := fmt.Sprintf(` +// buildItemsNeedingEmbeddingSQL returns the cheap missing/model-stale candidate +// query. It deliberately avoids the per-row item_people LATERAL joins that the +// text-staleness CTE (ListEmbeddingTextCandidates) needs, so the active-backfill +// pass stays cheap: a single LEFT JOIN on the embeddings table is enough to find +// rows that have no embedding or were embedded under a different model. Ordering +// by content_id makes the $2 cursor stable across calls. +// Args: $1 = current model, $2 = afterID cursor (empty for the first page), $3 = limit. +func buildItemsNeedingEmbeddingSQL() string { + return fmt.Sprintf(` SELECT mi.content_id FROM media_items mi LEFT JOIN media_item_embeddings e ON e.media_item_id = mi.content_id WHERE %s AND (e.media_item_id IS NULL OR e.model != $1) - LIMIT $2 + AND ($2 = '' OR mi.content_id > $2) + ORDER BY mi.content_id + LIMIT $3 `, embeddingEligibilityWhereClause()) - rows, err := r.pool.Query(ctx, query, currentModel, limit) +} + +// ItemsNeedingEmbedding returns content IDs of media items that either have no +// embedding or whose embedding was generated with a different model, ordered by +// content_id and paged via afterID (pass "" for the first page). This is the +// "cheap" pass of the embedding backfill: it intentionally does NOT detect +// text-staleness (which would require the expensive item_people LATERAL joins); +// re-embedding text-changed items is handled separately by +// ListEmbeddingTextCandidates. +// +// Books bypass the status='matched' gate because their scanner/plugin-derived +// metadata is authoritative as soon as the scan/enrichment completes. +func (r *Repo) ItemsNeedingEmbedding(ctx context.Context, currentModel, afterID string, limit int) ([]string, error) { + rows, err := r.pool.Query(ctx, buildItemsNeedingEmbeddingSQL(), currentModel, afterID, limit) if err != nil { return nil, fmt.Errorf("items needing embedding: %w", err) } diff --git a/internal/recommendations/repo_test.go b/internal/recommendations/repo_test.go index a67d84dd..12e849dc 100644 --- a/internal/recommendations/repo_test.go +++ b/internal/recommendations/repo_test.go @@ -1,8 +1,11 @@ package recommendations import ( + "context" "strings" "testing" + + "github.com/jackc/pgx/v5/pgxpool" ) func TestTasteSeedCandidateQueryOrdersByReliableColdStartSignals(t *testing.T) { @@ -146,6 +149,129 @@ func TestTasteProfileRefreshSubjectsQueryIncludesEbookReaderProgress(t *testing. ) } +func TestBuildItemsNeedingEmbeddingSQLIsCheap(t *testing.T) { + query := buildItemsNeedingEmbeddingSQL() + + // The whole point of the cheap pass is that it never pays the per-row + // item_people LATERAL joins that the text-staleness CTE needs. If either of + // these ever reappears here, Pass 1 has stopped being cheap. + if strings.Contains(query, "item_people") { + t.Fatalf("cheap embedding query must not reference item_people: %s", query) + } + if strings.Contains(query, "LATERAL") { + t.Fatalf("cheap embedding query must not use LATERAL joins: %s", query) + } + + normalized := strings.Join(strings.Fields(query), " ") + assertQueryTermsInOrder(t, normalized, + "LEFT JOIN media_item_embeddings e ON e.media_item_id = mi.content_id", + "e.media_item_id IS NULL OR e.model != $1", + "$2 = '' OR mi.content_id > $2", + "ORDER BY mi.content_id", + "LIMIT $3", + ) +} + +func TestItemsNeedingEmbeddingMissingAndModelStaleOnly(t *testing.T) { + pool := newEngineTestPool(t) + ctx := context.Background() + + const ( + prefix = "t7cheap-" + currentModel = "model-current" + oldModel = "model-old" + ) + cleanupRecoMediaItems(t, pool, prefix) + + missingID := prefix + "1-missing" + modelStaleID := prefix + "2-modelstale" + textStaleID := prefix + "3-textstale" + + // (i) matched movie, no embedding row at all -> cheap candidate. + seedRecoMediaItem(t, pool, missingID, "movie", "matched") + // (ii) embedding exists but under an OLD model -> cheap candidate. + seedRecoMediaItem(t, pool, modelStaleID, "movie", "matched") + seedRecoEmbedding(t, pool, modelStaleID, oldModel, modelStaleID) + // (iii) embedding under the CURRENT model. The cheap query only looks at + // model identity, so this row must NOT appear even though its cast changed + // (text-staleness is Pass 2's job, detected by the expensive CTE). + seedRecoMediaItem(t, pool, textStaleID, "movie", "matched") + seedRecoEmbedding(t, pool, textStaleID, currentModel, "stale canonical text") + + repo := NewRepo(pool) + + got, err := repo.ItemsNeedingEmbedding(ctx, currentModel, "", 100) + if err != nil { + t.Fatalf("ItemsNeedingEmbedding: %v", err) + } + gotSet := make(map[string]bool, len(got)) + for _, id := range got { + gotSet[id] = true + } + if !gotSet[missingID] { + t.Errorf("missing-embedding item %q should be a cheap candidate; got %v", missingID, got) + } + if !gotSet[modelStaleID] { + t.Errorf("model-stale item %q should be a cheap candidate; got %v", modelStaleID, got) + } + if gotSet[textStaleID] { + t.Errorf("text-stale-only item %q must NOT be a cheap candidate; got %v", textStaleID, got) + } + + // Results are ordered by content_id, so paging past the first id must skip + // it and still surface the rest. This is what lets EmbedAll advance its + // cursor over a failed/skipped item without re-fetching it forever. + if len(got) < 2 { + t.Fatalf("expected at least 2 cheap candidates to exercise the cursor, got %v", got) + } + first := got[0] + after, err := repo.ItemsNeedingEmbedding(ctx, currentModel, first, 100) + if err != nil { + t.Fatalf("ItemsNeedingEmbedding(afterID): %v", err) + } + for _, id := range after { + if id == first { + t.Fatalf("cursor afterID=%q still returned %q", first, first) + } + if id <= first { + t.Fatalf("cursor returned id %q <= afterID %q; ordering broken", id, first) + } + } +} + +// seedRecoMediaItem inserts a minimal embed-eligible media item whose title is +// its content_id. Delegates to seedRecoMediaItemTitled so both the cheap-query +// and EmbedAll tests seed rows that survive ItemRepository.GetByIDs. +func seedRecoMediaItem(t *testing.T, pool *pgxpool.Pool, contentID, mediaType, status string) { + t.Helper() + seedRecoMediaItemTitled(t, pool, contentID, mediaType, status, contentID) +} + +// seedRecoEmbedding inserts a zero-vector embedding row with an explicit model +// and canonical_text. The vector width matches the migrated column (3072). +func seedRecoEmbedding(t *testing.T, pool *pgxpool.Pool, contentID, model, canonicalText string) { + t.Helper() + if _, err := pool.Exec(context.Background(), ` + INSERT INTO media_item_embeddings (media_item_id, embedding, model, canonical_text) + VALUES ( + $1, + (SELECT array_agg(0.0::real) FROM generate_series(1, 3072))::vector, + $2, + $3 + ) + `, contentID, model, canonicalText); err != nil { + t.Fatalf("seed embedding %s: %v", contentID, err) + } +} + +func cleanupRecoMediaItems(t *testing.T, pool *pgxpool.Pool, prefix string) { + t.Helper() + t.Cleanup(func() { + // media_item_embeddings and item_people rows cascade on media_items delete. + _, _ = pool.Exec(context.Background(), `DELETE FROM media_items WHERE content_id LIKE $1`, prefix+"%") + }) +} + func assertQueryTermsInOrder(t *testing.T, query string, terms ...string) { t.Helper() diff --git a/internal/recommendations/similar.go b/internal/recommendations/similar.go index 25f9de60..2c9294dc 100644 --- a/internal/recommendations/similar.go +++ b/internal/recommendations/similar.go @@ -24,6 +24,16 @@ func embeddingTextNeedsRefresh(storedModel, storedCanonicalText, generatedCanoni return storedModel != currentModel || storedCanonicalText != generatedCanonicalText } +const ( + // embeddingBackfillBatchSize is how many items EmbedAll embeds per API call. + embeddingBackfillBatchSize = 10 + + // embeddingTextStaleQuotaPerRun bounds how many text-stale items the + // expensive Pass 2 re-embeds in a single EmbedAll run. It caps the cost of + // the one full-table text-staleness CTE scan per run. + embeddingTextStaleQuotaPerRun = 200 +) + // SimilarItems returns items most similar to the given item. Blends embedding // similarity (70%) with co-watch Jaccard score (30%), applies a validation // pipeline, MMR re-ranking, and assigns connection reasons. @@ -247,17 +257,40 @@ func (e *Engine) EmbedItem(ctx context.Context, itemID string) error { return e.repo.UpsertEmbedding(ctx, itemID, vectors[0], e.cfg.EmbeddingModel, text) } -// EmbedAll embeds all items that are missing embeddings or have stale canonical text. +// EmbedAll embeds items that are missing embeddings or have stale canonical +// text, in two passes: +// +// - Pass 1 (cheap): drain every missing or model-stale item via the +// ItemsNeedingEmbedding cursor. This query is a single LEFT JOIN with no +// item_people LATERAL joins, so active backfill (lots of brand-new items) +// stays cheap. The cursor advances past each page, so an item that fails to +// embed or store is simply retried on the next EmbedAll run rather than +// stalling the page forever. +// - Pass 2 (expensive): ONLY once Pass 1 has fully drained, make a single +// ListEmbeddingTextCandidates scan (LIMIT embeddingTextStaleQuotaPerRun) to +// find items whose embedding model is current but whose canonical text has +// drifted (e.g. a cast change). Detecting this requires recomputing each +// row's text via the full item_people LATERAL CTE, so it is run at most once +// per EmbedAll call. Re-embedding refreshes canonical_text, so handled items +// drop out of the candidate set on the next run — natural forward progress +// without a Pass 2 cursor. +// +// Coverage-first tradeoff: in steady state Pass 1 drains every run, so text-stale +// items are re-embedded promptly. Only under pathological continuous heavy +// ingest (Pass 1 never drains within a run) does Pass 2 get skipped — we +// deliberately prioritize getting NEW items covered over re-embedding +// text-changed ones. A periodic "force Pass 2 every Nth run" is a possible +// follow-up, intentionally out of scope here. func (e *Engine) EmbedAll(ctx context.Context) (int, error) { model := e.cfg.EmbeddingModel total := 0 - batchSize := 10 - afterID := "" if err := e.ensureEmbeddingLockConfig(ctx); err != nil { return total, err } + // Pass 1: drain the cheap missing/model-stale backlog via the cursor. + afterID := "" for { select { case <-ctx.Done(): @@ -265,46 +298,122 @@ func (e *Engine) EmbedAll(ctx context.Context) (int, error) { default: } - candidates, err := e.repo.ListEmbeddingTextCandidates(ctx, afterID, model, batchSize) + ids, err := e.repo.ItemsNeedingEmbedding(ctx, model, afterID, embeddingBackfillBatchSize) if err != nil { return total, err } - if len(candidates) == 0 { + if len(ids) == 0 { break } - afterID = candidates[len(candidates)-1].MediaItemID - - ids := make([]string, 0, len(candidates)) - stored := make(map[string]EmbeddingTextCandidate, len(candidates)) - for _, candidate := range candidates { - ids = append(ids, candidate.MediaItemID) - stored[candidate.MediaItemID] = candidate - } + afterID = ids[len(ids)-1] items, err := e.itemRepo.GetByIDs(ctx, ids) if err != nil { return total, fmt.Errorf("get items for embedding: %w", err) } - // Hydrate cast/crew for richer embedding text. e.hydrateItemPeople(ctx, items) texts := make([]string, len(items)) - staleItems := make([]*models.MediaItem, 0, len(items)) - staleTexts := make([]string, 0, len(items)) for i, item := range items { texts[i] = embeddings.BuildEmbeddingText(item) - candidate := stored[item.ContentID] - if embeddingTextNeedsRefresh(candidate.Model, candidate.CanonicalText, texts[i], model) { - staleItems = append(staleItems, item) - staleTexts = append(staleTexts, texts[i]) - } } - if len(staleItems) == 0 { + + // embedBatch only returns quota/billing or context errors; both mean the + // cheap backlog did not fully drain this run, so we stop here and skip the + // expensive Pass 2. Any per-item store/embed failure is swallowed inside + // embedBatch and retried next run (the cursor passes over it). + embedded, err := e.embedBatch(ctx, items, texts, model) + total += embedded + if err != nil { + return total, err + } + } + + // Reaching here means Pass 1 fully drained (every error path above returns). + + select { + case <-ctx.Done(): + return total, ctx.Err() + default: + } + + // Pass 2: one expensive text-staleness scan, bounded by the per-run quota. + candidates, err := e.repo.ListEmbeddingTextCandidates(ctx, "", model, embeddingTextStaleQuotaPerRun) + if err != nil { + return total, err + } + if len(candidates) == 0 { + return total, nil + } + + ids := make([]string, 0, len(candidates)) + stored := make(map[string]EmbeddingTextCandidate, len(candidates)) + for _, candidate := range candidates { + ids = append(ids, candidate.MediaItemID) + stored[candidate.MediaItemID] = candidate + } + + items, err := e.itemRepo.GetByIDs(ctx, ids) + if err != nil { + return total, fmt.Errorf("get items for embedding: %w", err) + } + e.hydrateItemPeople(ctx, items) + + // Pass 1 already cleared missing/model-stale rows, so the remaining + // candidates are text-stale. Re-confirm in Go (the SQL detection is an + // approximation of BuildEmbeddingText) before paying for an embed. + staleItems := make([]*models.MediaItem, 0, len(items)) + staleTexts := make([]string, 0, len(items)) + for _, item := range items { + text := embeddings.BuildEmbeddingText(item) + candidate := stored[item.ContentID] + if embeddingTextNeedsRefresh(candidate.Model, candidate.CanonicalText, text, model) { + staleItems = append(staleItems, item) + staleTexts = append(staleTexts, text) + } + } + + embedded, err := e.embedBatch(ctx, staleItems, staleTexts, model) + total += embedded + if err != nil { + return total, err + } + + return total, nil +} + +// embedBatch embeds the given items (parallel items/texts slices, same length), +// chunking internally by embeddingBackfillBatchSize so callers can pass any +// size. It returns the number of items successfully stored. +// +// Behavior preserved from the original EmbedAll inner loop: +// - quota/billing errors return immediately (retrying won't help) and are the +// only batch-embed errors propagated to the caller; +// - any other batch embed failure falls back to embedding one item at a time +// so a single oversized item can't block the rest; +// - ensureEmbeddingLock runs before each upsert; +// - a store (upsert) failure for one item is logged and skipped, not fatal. +func (e *Engine) embedBatch(ctx context.Context, items []*models.MediaItem, texts []string, model string) (int, error) { + total := 0 + for start := 0; start < len(items); start += embeddingBackfillBatchSize { + select { + case <-ctx.Done(): + return total, ctx.Err() + default: + } + + end := start + embeddingBackfillBatchSize + if end > len(items) { + end = len(items) + } + chunkItems := items[start:end] + chunkTexts := texts[start:end] + if len(chunkItems) == 0 { continue } - vectors, err := e.embClient.Embed(ctx, staleTexts) + vectors, err := e.embClient.Embed(ctx, chunkTexts) if err != nil { // Quota/billing errors won't resolve by retrying — stop immediately. if isQuotaError(err) { @@ -313,12 +422,12 @@ func (e *Engine) EmbedAll(ctx context.Context) (int, error) { // Batch failed — fall back to embedding one at a time so a single // oversized item doesn't block the entire job. - slog.Warn("batch embed failed, falling back to single-item mode", "error", err, "batch_size", len(staleItems)) - for i, item := range staleItems { + slog.Warn("batch embed failed, falling back to single-item mode", "error", err, "batch_size", len(chunkItems)) + for i, item := range chunkItems { if ctx.Err() != nil { return total, ctx.Err() } - single := staleTexts[i] + single := chunkTexts[i] vecs, embedErr := e.embClient.Embed(ctx, []string{single}) if embedErr != nil { if isQuotaError(embedErr) { @@ -344,14 +453,14 @@ func (e *Engine) EmbedAll(ctx context.Context) (int, error) { continue } - for i, item := range staleItems { + for i, item := range chunkItems { if i >= len(vectors) { break } if err := e.ensureEmbeddingLock(ctx, vectors[i]); err != nil { return total, fmt.Errorf("embed item %s: %w", item.ContentID, err) } - if err := e.repo.UpsertEmbedding(ctx, item.ContentID, vectors[i], model, staleTexts[i]); err != nil { + if err := e.repo.UpsertEmbedding(ctx, item.ContentID, vectors[i], model, chunkTexts[i]); err != nil { slog.Warn("skipping item, store failed", "item_id", item.ContentID, "error", err) continue }