diff --git a/internal/recommendations/accuracy_test.go b/internal/recommendations/accuracy_test.go index 42b2b361f1..e78d4247db 100644 --- a/internal/recommendations/accuracy_test.go +++ b/internal/recommendations/accuracy_test.go @@ -33,7 +33,7 @@ func TestApplyGenreCapCountsAllGenres(t *testing.T) { } } -func TestHNSWEfSearchUsesCandidateLimitFloor(t *testing.T) { +func TestHNSWEfSearchStaysWithinPgvectorRange(t *testing.T) { tests := []struct { name string candidateLimit int @@ -42,6 +42,10 @@ func TestHNSWEfSearchUsesCandidateLimitFloor(t *testing.T) { {name: "raises small scans", candidateLimit: 40, want: minHNSWEfSearch}, {name: "keeps exact floor", candidateLimit: minHNSWEfSearch, want: minHNSWEfSearch}, {name: "keeps larger scans", candidateLimit: 900, want: 900}, + {name: "keeps pgvector maximum", candidateLimit: maxHNSWEfSearch, want: maxHNSWEfSearch}, + {name: "caps just above maximum", candidateLimit: maxHNSWEfSearch + 1, want: maxHNSWEfSearch}, + {name: "caps genre discover pool", candidateLimit: 1200, want: maxHNSWEfSearch}, + {name: "caps largest ANN limit", candidateLimit: 2000, want: maxHNSWEfSearch}, } for _, tt := range tests { diff --git a/internal/recommendations/hnsw_ef_search_db_test.go b/internal/recommendations/hnsw_ef_search_db_test.go new file mode 100644 index 0000000000..8cf49afdcf --- /dev/null +++ b/internal/recommendations/hnsw_ef_search_db_test.go @@ -0,0 +1,80 @@ +package recommendations + +import ( + "context" + "fmt" + "os" + "slices" + "testing" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/Silo-Server/silo-server/internal/catalog" +) + +// TestFindTasteProfileCandidatesWithGenresAboveEfSearchMaximumDB covers the +// Watch Tonight discover path: a genre filter multiplies the candidate pool by +// five, so a pool of 240 asks pgvector for 1200 ANN candidates, above the +// largest hnsw.ef_search it accepts. +func TestFindTasteProfileCandidatesWithGenresAboveEfSearchMaximumDB(t *testing.T) { + newEngineTestPool(t) // skips without a migrated database + ctx := t.Context() + cfg, err := pgxpool.ParseConfig(os.Getenv("SILO_TEST_DATABASE_URL")) + if err != nil { + t.Fatal(err) + } + // pgvector checks hnsw.ef_search once its library is loaded in the backend. + // A pooled production connection has usually run a vector query already, + // so load it on every connection rather than relying on a fresh one. + cfg.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error { + _, err := conn.Exec(ctx, `SELECT '[1]'::vector`) + return err + } + pool, err := pgxpool.NewWithConfig(ctx, cfg) + if err != nil { + t.Fatal(err) + } + t.Cleanup(pool.Close) + + const prefix = "test-hnsw-ef-search-" + cleanupRecoMediaItems(t, pool, prefix) + var want []string + for i := range 3 { + id := fmt.Sprintf("%smovie-%d", prefix, i) + seedRecoMediaItem(t, pool, id, "movie", "matched") + if _, err := pool.Exec(ctx, `UPDATE media_items SET genres = '{Drama}' WHERE content_id = $1`, id); err != nil { + t.Fatal(err) + } + if _, err := pool.Exec(ctx, ` + INSERT INTO media_item_embeddings (media_item_id, embedding, model, canonical_text) + VALUES ($1, (SELECT array_agg((g % ($2 + 2) + 1)::real) FROM generate_series(1, 3072) g)::vector, 'test-model', $1) + `, id, i); err != nil { + t.Fatalf("seed embedding %s: %v", id, err) + } + want = append(want, id) + } + + embedding := make([]float32, CanonicalEmbeddingDimensions) + for i := range embedding { + embedding[i] = float32(i%2 + 1) + } + const watchTonightGenrePool = 240 + items, genres, err := NewRepo(pool).FindTasteProfileCandidates(ctx, embedding, nil, []string{"Drama"}, watchTonightGenrePool, catalog.AccessFilter{}) + if err != nil { + t.Fatalf("FindTasteProfileCandidates with genres and limit %d: %v", watchTonightGenrePool, err) + } + var got []string + for _, item := range items { + if slices.Contains(want, item.MediaItemID) { + got = append(got, item.MediaItemID) + if !slices.Equal(genres[item.MediaItemID], []string{"Drama"}) { + t.Fatalf("genres for %s = %v, want [Drama]", item.MediaItemID, genres[item.MediaItemID]) + } + } + } + slices.Sort(got) + if !slices.Equal(got, want) { + t.Fatalf("taste profile candidates = %v, want the seeded Drama movies %v", got, want) + } +} diff --git a/internal/recommendations/repo.go b/internal/recommendations/repo.go index 0375b683d5..917985f6f8 100644 --- a/internal/recommendations/repo.go +++ b/internal/recommendations/repo.go @@ -24,6 +24,12 @@ func ensureCanonicalDimensions(vec []float32) ([]float32, error) { const embeddingLockSettingKey = "recommendations.embedding_lock" const minHNSWEfSearch = 200 +// maxHNSWEfSearch is the largest hnsw.ef_search pgvector accepts. A larger +// value fails set_config, so the candidate scan never runs. With iterative +// scans the index keeps returning rows past ef_search, so the cap does not +// stop a larger LIMIT at 1000 candidates. +const maxHNSWEfSearch = 1000 + // watchedActivityCTE unifies video watch progress and ebook reader progress // into one activity stream. Episodes roll up to their parent series via // COALESCE(e.series_id, ...). The ebook completed flag is derived from the @@ -116,7 +122,7 @@ func NewRepo(pool *pgxpool.Pool) *Repo { } func hnswEfSearch(candidateLimit int) int { - return max(candidateLimit, minHNSWEfSearch) + return min(max(candidateLimit, minHNSWEfSearch), maxHNSWEfSearch) } func (r *Repo) withHNSWCandidateScan(ctx context.Context, candidateLimit int, fn func(pgx.Tx) error) error { diff --git a/scripts/ci/db-contracts.txt b/scripts/ci/db-contracts.txt index c6a94e2e66..2726418dab 100644 --- a/scripts/ci/db-contracts.txt +++ b/scripts/ci/db-contracts.txt @@ -51,6 +51,9 @@ ./internal/notifications TestEmailVerificationDispatchRetriesUncertainSendOnceWithSameMessage ./internal/notifications TestEmailVerificationDispatchStopsAfterSecondUncertainty +# Recommendation candidate scans stay within pgvector's hnsw.ef_search range. +./internal/recommendations TestFindTasteProfileCandidatesWithGenresAboveEfSearchMaximumDB + # Image cache jobs run again for artwork the target no longer holds. ./internal/metadata TestEnqueueRequeuesSucceededJobOnlyWhenAskedDB