diff --git a/CHANGELOG.md b/CHANGELOG.md index a0a6bca..4824a14 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,13 @@ The project publishes 0.x prerelease versions; a stable release line is not yet ### Added +- Cosine HNSW indexes on `embeddings_text` (768), `embeddings_visual` (512), + and `embeddings_face` (512) via migration 0025 (`#173`). Text search walks + `ORDER BY embedding <=> query LIMIT n` (planner-usable) and falls back to + exact per-file `DISTINCT ON` when a bounded scan underfills after + deduplication. Visual cosine-order queries can use the visual index. Face + clustering remains in-process; the face index is DDL only. Recall is not + claimed here; the live harness is `#175`. - Advertise the model-free lexical file-search route in the MCP `mem_search` schema and verify `tools/list` plus route/filter forwarding through `tools/call`. diff --git a/SPEC.md b/SPEC.md index 277caa8..8755146 100644 --- a/SPEC.md +++ b/SPEC.md @@ -537,7 +537,10 @@ embeddings_face ( - `files` FTS + trigram — 显式指定 `route=lexical` 的文件名无模型词法召回; `auto` 只融合 text/visual,不自动回退到 lexical,worker 不可用时仍报错。 仅搜索 `files.name`,路径只用于筛选,不检索文件正文或路径片段。 -- `embeddings_* (embedding)` — pgvector HNSW +- `embeddings_* (embedding)` — pgvector HNSW (`vector_cosine_ops`, migration 0025) + on `embeddings_text` (768), `embeddings_visual` (512), `embeddings_face` (512). + Text search uses cosine-order candidates plus exact per-file fallback. + Face clustering is still in-process. `index_generation_vectors` is not indexed. - `file_entities (entity_id)` — 反查"和某人有关的所有文件" ### 6.3 文件夹一致性规则(重要) diff --git a/docs/MIGRATION_SEQUENCE.md b/docs/MIGRATION_SEQUENCE.md index 7ed41df..df3c8d1 100644 --- a/docs/MIGRATION_SEQUENCE.md +++ b/docs/MIGRATION_SEQUENCE.md @@ -4,9 +4,9 @@ These draft changes are cumulative, not independently deployable: | Order | Draft / original PR | Migration | Required predecessor | | --- | --- | --- | --- | -| 1 | #194 / #183 | 0024 file lexical lane | released/main schema 23 | -| 2 | #197 / #180 | 0025 HNSW indexes | #194, schema 24 | -| 3 | #195 / #185 | 0026 data-plane hygiene | #197, schema 25 | +| 1 | #194 / #183 | 0024 file lexical lane | released/main schema 23 (merged) | +| 2 | #173 HNSW completion (supersedes #197 HOLD) | 0025 HNSW indexes + text continuation | #194, schema 24 | +| 3 | #195 / #185 | 0026 data-plane hygiene | schema 25 | The PR base chain is `main` → `codex/fix-pr-183` → `codex/fix-pr-180` → `codex/fix-pr-185`. Successor branches must include their predecessor schema and source. Local @@ -19,13 +19,14 @@ tags `v0.1.0` / `v0.1.1` contain only migrations 0001–0023. This does not prov that a private deployment never applied a draft. Consequently migration numbers and SQL identities are retained, not renumbered on an assumption. -HNSW DDL and shipping text-query planner acceptance are separate evidence: -creating indexes does not prove that a text query uses them. An index-only -successor must describe that partial scope and leave the broader #173 acceptance -open. #195 still requires the predecessor schema 25 regardless of query strategy. -This document does not approve a product decision or query-strategy change, -waive a review gate, or authorize deployment. #176's model-free file-lane RFC -also requires a maintainer decision before this draft is made review-ready. +Migration 0025 creates the three cosine HNSW indexes. The shipping text route +no longer uses `DISTINCT ON (f.id) ORDER BY f.id` as its primary plan: it walks +cosine-ordered candidates and falls back to that exact query only when a bounded +scan underfills. Visual cosine-order already matched HNSW. Face DDL is not a +face-query speedup. #195 still requires predecessor schema 25. Recall and live +latency remain `#175`, not this migration. + +This document does not waive a review gate or authorize deployment. Goose startup remains strict: no `WithAllowMissing` or equivalent option is enabled. A database that already applied 26 while omitting 24/25 will correctly diff --git a/docs/VALIDATION_HNSW.md b/docs/VALIDATION_HNSW.md new file mode 100644 index 0000000..d59b1c8 --- /dev/null +++ b/docs/VALIDATION_HNSW.md @@ -0,0 +1,60 @@ +# HNSW index + text continuation for #173 + +Migration `0025_ann_hnsw_indexes.sql` adds cosine HNSW indexes to text (768), +visual (512), and face (512) embeddings. Main already shipped lexical +migration 0024; this branch's head is 25. + +## What this change proves + +- Populated 24 → 25 → 24 → 25 preserves text/visual/face vectors and rebuilds + three `VALID` `vector_cosine_ops` HNSW indexes (`TestHNSWMigrationPostgres`). +- Post-index INSERT succeeds; an UPDATE to the wrong dimension is rejected by + the `vector(N)` column type (failure mode: PostgreSQL dimension error, not a + silent pad/truncate). +- `EXPLAIN (ANALYZE)` of the shipping text cosine-order query and the visual + cosine-order query names `idx_embeddings_text_embedding_hnsw` and + `idx_embeddings_visual_embedding_hnsw` on a 2,000-row corpus. Planner + settings are not forced. +- `TestTextANNFileSemanticsPostgres` keeps best-chunk-per-file top-k when one + file owns 101 nearest chunks, and still enforces owner, literal path, + allow-list, MIME, and time filters. Invalid allow-lists fail closed. + +## Text continuation / fallback + +A bounded `ORDER BY distance LIMIT n` scan can underfill after per-file +deduplication (`ef_search=40` returning 40 chunks of one file). The shipping +path: + +1. Run a CTE `ORDER BY embedding <=> $1 LIMIT remaining` on `embeddings_text` + (HNSW-compatible; omit `ANY(exclude)` when the exclude list is empty). +2. Join those candidates to `files` and apply owner/path/MIME/time filters. +3. Keep the first sighting of each file (that chunk is the file's best). +4. Repeat, excluding selected files, until k files are collected. +5. If a round returns no new files, fill the remainder with the original + exact `DISTINCT ON (f.id) ORDER BY f.id, distance` query. + +Step 1 is the planner-usable shape. Step 4 preserves the previous result +contract on pathological corpora. Iterative-scan GUC is not enabled. + +## What this change does not prove + +- Live embedding quality, production latency, index build time, or numerical + recall. Those belong to [#175](https://github.com/bytefolk/mem/issues/175) + (shipping search-path producer) and the closed producer attempt + [#184](https://github.com/bytefolk/mem/pull/184). Fixture scores are not + substituted. +- Face query speedup. `assignCluster` still averages centroids in Go. +- `index_generation_vectors` ANN. The column is undimensioned. + +## Local gates + +```bash +MEM_TEST_DB="$MEM_TEST_DB" ./scripts/verify.sh integration +``` + +`run_hnsw_migration` creates a fresh `_test` database and runs +`TestHNSWMigrationPostgres`, which records EXPLAIN ANALYZE. `scripts/verify_hnsw_indexes.sh` +is a manual `psql` helper; CI does not call it because libpq rejects some pgx URIs. + +Face evidence is valid DDL and populated-table migration, not a measured +face-query speedup. diff --git a/scripts/verify.sh b/scripts/verify.sh index b8a45ad..0c74b8f 100755 --- a/scripts/verify.sh +++ b/scripts/verify.sh @@ -4,7 +4,7 @@ set -euo pipefail REPO_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)" MODE="${1:-unit}" -EXPECTED_MIGRATION_HEAD=24 +EXPECTED_MIGRATION_HEAD=25 MIGRATION_ROLLBACK_TARGET=11 MODEL_TEXT_CANONICAL_BASE=15 WORKSPACE_AI_PROFILE_BASE=16 @@ -349,6 +349,7 @@ run_postgres_tests() { TestManagedAISettlementOutboxPostgres TestReleasedFileStageRetryPostgres TestDurableContextPostgres + TestTextANNFileSemanticsPostgres ) integration_log="$(mktemp "${TMPDIR:-/tmp}/mem-integration.XXXXXX")" @@ -359,7 +360,7 @@ run_postgres_tests() { MEM_TEST_DB="$MEM_TEST_DB" go test \ ${race_flag:+"$race_flag"} \ -v -count=1 -p 1 -timeout 20m \ - -run '^(TestMemoryPostgres|TestHandoffPostgres|TestWorkspaceTransferPostgres|TestWorkspaceTransferMergeConservativePostgres|TestHandoffCrossAgentHTTPIntegration|TestRelocateHTTPPostgres|TestMemoryPathLifecycleIntegration|TestWorkspacePathLockingIntegration|TestFilePathLockingIntegration|TestAnnotationDecisionIntegration|TestIndexerEnrichmentIntegration|TestRecomputePerson|TestManagedEmbeddingEntitlementPostgres|TestManagedSearchReplayPostgres|TestManagedEmbeddingHTTPAuthorizationPostgres|TestAIProfilePostgres|TestIndexGenerationPostgres|TestManagedAISettlementOutboxPostgres|TestReleasedFileStageRetryPostgres|TestDurableContextPostgres)$' \ + -run '^(TestMemoryPostgres|TestHandoffPostgres|TestWorkspaceTransferPostgres|TestWorkspaceTransferMergeConservativePostgres|TestHandoffCrossAgentHTTPIntegration|TestRelocateHTTPPostgres|TestMemoryPathLifecycleIntegration|TestWorkspacePathLockingIntegration|TestFilePathLockingIntegration|TestAnnotationDecisionIntegration|TestIndexerEnrichmentIntegration|TestRecomputePerson|TestManagedEmbeddingEntitlementPostgres|TestManagedSearchReplayPostgres|TestManagedEmbeddingHTTPAuthorizationPostgres|TestAIProfilePostgres|TestIndexGenerationPostgres|TestManagedAISettlementOutboxPostgres|TestReleasedFileStageRetryPostgres|TestDurableContextPostgres|TestTextANNFileSemanticsPostgres)$' \ ./internal/memory \ ./internal/handoff \ ./internal/workspacetransfer \ @@ -397,9 +398,20 @@ run_integration() { validate_test_database with_fresh_test_database migration_sequence run_migration_sequence with_fresh_test_database migration run_migration_round_trip + with_fresh_test_database hnsw_migration run_hnsw_migration with_fresh_test_database integration run_postgres_integration } +run_hnsw_migration() { + log "Populated HNSW migration, rollback, ingest, dimension rejection and EXPLAIN" + ( + cd "${REPO_ROOT}/server" + MEM_HNSW_TEST_DB="$MEM_TEST_DB" go test -v -count=1 \ + -run '^TestHNSWMigrationPostgres$' ./internal/db + ) + log "Planner EXPLAIN is recorded by TestHNSWMigrationPostgres (psql URI script is manual)" +} + run_integration_race() { validate_test_database with_fresh_test_database integration_race run_postgres_integration_race diff --git a/scripts/verify_hnsw_indexes.sh b/scripts/verify_hnsw_indexes.sh new file mode 100755 index 0000000..0eb0cac --- /dev/null +++ b/scripts/verify_hnsw_indexes.sh @@ -0,0 +1,58 @@ +#!/usr/bin/env bash +# Read-only planner verification for shipping text and visual query shapes. +# Requires a populated disposable database; does not force planner settings. +set -euo pipefail +trap 'echo "ERROR: HNSW verification aborted on an execution error; assertions are incomplete" >&2' ERR +DB_URL="${1:?Usage: $0 }" +CORPUS_USER="${2:?Supply the user UUID that owns the populated corpus}" +[[ "$CORPUS_USER" =~ ^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$ ]] || { + echo 'ERROR: corpus user must be a UUID' >&2; exit 1; +} +sql() { psql -X -A -t -v ON_ERROR_STOP=1 "$DB_URL" -c "$1"; } +db_name="$(sql 'SELECT current_database()')" +[[ "$db_name" == *_test ]] || { echo 'ERROR: database must end in _test' >&2; exit 1; } +pass=0 +fail=0 +for kind in text visual face; do + index="idx_embeddings_${kind}_embedding_hnsw" + valid="$(sql "SELECT count(*) FROM pg_index i + JOIN pg_class c ON c.oid = i.indexrelid JOIN pg_am a ON a.oid = c.relam + WHERE i.indrelid = 'embeddings_${kind}'::regclass + AND c.relname = '${index}' AND a.amname = 'hnsw' AND i.indisvalid + AND pg_get_indexdef(i.indexrelid) LIKE '%vector_cosine_ops%'")" + rows="$(sql "SELECT count(*) FROM embeddings_${kind} e JOIN files f ON f.id=e.file_id + WHERE f.user_id='${CORPUS_USER}'::uuid AND e.embedding IS NOT NULL")" + if [[ "$valid" == 1 && "$rows" -gt 0 ]]; then + echo "PASS: ${index} is valid; corpus contains ${rows} non-null vectors" + pass=$((pass + 1)) + else + echo "FAIL: ${index}: valid=${valid}, corpus vectors=${rows}" + fail=$((fail + 1)) + fi +done +assert_plan() { + local route="$1" query="$2" plan + plan="$(sql "EXPLAIN (ANALYZE, BUFFERS) ${query}")" + echo "$plan" + if grep -q "idx_embeddings_${route}_embedding_hnsw" <<<"$plan"; then + echo "PASS: shipping ${route} query uses HNSW" + pass=$((pass + 1)) + else + echo "FAIL: shipping ${route} query does not use HNSW" + fail=$((fail + 1)) + fi +} +# Match queryTextDistanceOrder: ANN CTE then file join. DISTINCT ON is fallback. +assert_plan text "WITH nearest AS ( + SELECT e.id, e.file_id, e.embedding <=> array_fill(0.1::real, ARRAY[768])::vector AS dist + FROM embeddings_text e + ORDER BY e.embedding <=> array_fill(0.1::real, ARRAY[768])::vector ASC + LIMIT 10 +) +SELECT e.id, f.id FROM nearest e JOIN files f ON f.id=e.file_id + WHERE f.user_id='${CORPUS_USER}'::uuid ORDER BY e.dist ASC" +assert_plan visual "SELECT e.file_id + FROM embeddings_visual e + ORDER BY e.embedding <=> array_fill(0.1::real, ARRAY[512])::vector ASC LIMIT 10" +echo "Results: ${pass} passed, ${fail} failed" +[[ "$fail" -eq 0 ]] diff --git a/server/internal/db/hnsw_migration_test.go b/server/internal/db/hnsw_migration_test.go new file mode 100644 index 0000000..4bdc846 --- /dev/null +++ b/server/internal/db/hnsw_migration_test.go @@ -0,0 +1,172 @@ +package db + +import ( + "context" + "database/sql" + "fmt" + "os" + "strings" + "testing" + "time" + + "github.com/jackc/pgx/v5" + "github.com/pressly/goose/v3" +) + +const hnswCorpusUser = "00000000-0000-0000-0000-000000000173" + +// Populated 24→25→24→25, ingest after index build, wrong-dimension rejection, +// and EXPLAIN of the shipping cosine-order text/visual shapes. +func TestHNSWMigrationPostgres(t *testing.T) { + dsn := os.Getenv("MEM_HNSW_TEST_DB") + if dsn == "" { + t.Skip("MEM_HNSW_TEST_DB not set; requires a fresh owned test database") + } + cfg, err := pgx.ParseConfig(dsn) + if err != nil { + t.Fatal(err) + } + if !strings.HasSuffix(cfg.Database, "_test") { + t.Fatalf("refusing non-test database %q", cfg.Database) + } + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute) + defer cancel() + db, err := sql.Open("pgx", dsn) + if err != nil { + t.Fatal(err) + } + defer db.Close() + var history sql.NullString + if err := db.QueryRowContext(ctx, "SELECT to_regclass('goose_db_version')::text").Scan(&history); err != nil { + t.Fatal(err) + } + if history.Valid { + t.Fatal("refusing existing migration history; provide a new owned test database") + } + goose.SetBaseFS(migrationsFS) + if err := goose.SetDialect("postgres"); err != nil { + t.Fatal(err) + } + if err := goose.UpToContext(ctx, db, "migrations", 24); err != nil { + t.Fatal(err) + } + exec := func(query string) { + t.Helper() + if _, err := db.ExecContext(ctx, query); err != nil { + t.Fatal(err) + } + } + exec(`INSERT INTO users(id,email,password_hash) + VALUES ('` + hnswCorpusUser + `','hnsw@example.test','test'); + INSERT INTO files(id,user_id,name,path,size,sha256,mime,storage_key) + SELECT md5(i::text)::uuid,'` + hnswCorpusUser + `', + 'fixture-' || i, '/hnsw', 0, 'sha-' || i, 'text/plain', 'fixture-' || i + FROM generate_series(1,2000) i`) + for _, kind := range []string{"text", "visual", "face"} { + dim, extraCols, extraValues := 512, "", "" + if kind == "text" { + dim, extraCols, extraValues = 768, ",chunk_index,chunk_text", ",0,'fixture'" + } + exec(fmt.Sprintf(`INSERT INTO embeddings_%s(file_id,embedding%s) + SELECT id,array_fill(0.1::real,ARRAY[%d])::vector%s FROM files`, kind, extraCols, dim, extraValues)) + } + assertState := func(wantIndexes, wantRows int) { + t.Helper() + for _, kind := range []string{"text", "visual", "face"} { + var indexes, rows int + if err := db.QueryRowContext(ctx, `SELECT count(*) FROM pg_index i + JOIN pg_class c ON c.oid=i.indexrelid JOIN pg_am a ON a.oid=c.relam + WHERE i.indrelid=($1::text)::regclass AND c.relname=$2 AND a.amname='hnsw' + AND i.indisvalid AND pg_get_indexdef(i.indexrelid) LIKE '%vector_cosine_ops%'`, + "embeddings_"+kind, "idx_embeddings_"+kind+"_embedding_hnsw").Scan(&indexes); err != nil { + t.Fatal(err) + } + if err := db.QueryRowContext(ctx, "SELECT count(*) FROM embeddings_"+kind+" WHERE embedding IS NOT NULL").Scan(&rows); err != nil { + t.Fatal(err) + } + if indexes != wantIndexes || rows != wantRows { + t.Fatalf("%s: valid cosine indexes=%d want=%d, preserved vectors=%d want=%d", kind, indexes, wantIndexes, rows, wantRows) + } + } + } + assertState(0, 2000) + if err := goose.UpToContext(ctx, db, "migrations", 25); err != nil { + t.Fatal(err) + } + assertState(1, 2000) + if err := goose.DownToContext(ctx, db, "migrations", 24); err != nil { + t.Fatal(err) + } + assertState(0, 2000) + if err := goose.UpToContext(ctx, db, "migrations", 25); err != nil { + t.Fatal(err) + } + assertState(1, 2000) + exec(`INSERT INTO files(id,user_id,name,path,size,sha256,mime,storage_key) + VALUES (md5('2001')::uuid,'` + hnswCorpusUser + `', + 'after-index','/hnsw',0,'sha-2001','text/plain','after-index')`) + for _, kind := range []string{"text", "visual", "face"} { + dim, extraCols, extraValues := 512, "", "" + if kind == "text" { + dim, extraCols, extraValues = 768, ",chunk_index,chunk_text", ",0,'after-index'" + } + exec(fmt.Sprintf(`INSERT INTO embeddings_%s(file_id,embedding%s) + VALUES (md5('2001')::uuid,array_fill(0.2::real,ARRAY[%d])::vector%s)`, kind, extraCols, dim, extraValues)) + _, err := db.ExecContext(ctx, fmt.Sprintf(`UPDATE embeddings_%s + SET embedding=array_fill(0.1::real,ARRAY[%d])::vector WHERE file_id=md5('2001')::uuid`, kind, dim-1)) + if err == nil || !strings.Contains(strings.ToLower(err.Error()), "dimension") { + t.Fatalf("%s: wrong dimensionality must fail, got %v", kind, err) + } + } + assertState(1, 2001) + if _, err := db.ExecContext(ctx, `ANALYZE embeddings_text; ANALYZE embeddings_visual; ANALYZE files`); err != nil { + t.Fatal(err) + } + assertIndexScan(t, ctx, db, "idx_embeddings_text_embedding_hnsw", ` + WITH nearest AS ( + SELECT e.id, e.file_id, e.embedding <=> array_fill(0.1::real, ARRAY[768])::vector AS dist + FROM embeddings_text e + ORDER BY e.embedding <=> array_fill(0.1::real, ARRAY[768])::vector ASC + LIMIT 10 + ) + SELECT e.id, f.id + FROM nearest e + JOIN files f ON f.id = e.file_id + WHERE f.user_id = '`+hnswCorpusUser+`'::uuid + ORDER BY e.dist ASC`) + assertIndexScan(t, ctx, db, "idx_embeddings_visual_embedding_hnsw", ` + SELECT e.file_id + FROM embeddings_visual e + ORDER BY e.embedding <=> array_fill(0.1::real, ARRAY[512])::vector ASC + LIMIT 10`) + if err := (&DB{url: dsn}).Migrate(ctx); err != nil { + t.Fatal(err) + } + t.Log("PASS: populated 24->25->24->25, valid cosine HNSW, ingest, dimension rejection, text/visual EXPLAIN uses HNSW") +} + +func assertIndexScan(t *testing.T, ctx context.Context, db *sql.DB, index, query string) { + t.Helper() + rows, err := db.QueryContext(ctx, "EXPLAIN (ANALYZE, BUFFERS) "+query) + if err != nil { + t.Fatalf("EXPLAIN %s: %v", index, err) + } + defer rows.Close() + var plan strings.Builder + for rows.Next() { + var line string + if err := rows.Scan(&line); err != nil { + t.Fatal(err) + } + plan.WriteString(line) + plan.WriteByte('\n') + } + if err := rows.Err(); err != nil { + t.Fatal(err) + } + text := plan.String() + t.Logf("EXPLAIN %s:\n%s", index, text) + if !strings.Contains(text, index) { + t.Fatalf("shipping plan did not use %s:\n%s", index, text) + } +} diff --git a/server/internal/db/migrations/0001_init.sql b/server/internal/db/migrations/0001_init.sql index a5f9bf2..e6a6328 100644 --- a/server/internal/db/migrations/0001_init.sql +++ b/server/internal/db/migrations/0001_init.sql @@ -106,8 +106,8 @@ CREATE TABLE IF NOT EXISTS embeddings_text ( embedding vector(768) ); CREATE INDEX IF NOT EXISTS idx_embeddings_text_file ON embeddings_text (file_id); --- HNSW index will be added by worker once we settle on a model dimension. Kept off here --- because pgvector requires the table to have data of consistent dim before building. +-- Cosine HNSW for embeddings_text/visual/face is created in migration 0025 +-- once vector(768)/vector(512) dimensions are fixed. -- +goose StatementEnd -- +goose StatementBegin diff --git a/server/internal/db/migrations/0019_versioned_index_generations.sql b/server/internal/db/migrations/0019_versioned_index_generations.sql index bd173c8..31a06d9 100644 --- a/server/internal/db/migrations/0019_versioned_index_generations.sql +++ b/server/internal/db/migrations/0019_versioned_index_generations.sql @@ -264,8 +264,10 @@ CREATE INDEX idx_index_generation_targets_file_hash -- +goose StatementBegin -- `vector` intentionally has no table-wide dimension. Every row is validated -- against its immutable generation.output_dimension by the canonical service. --- Future ANN indexes must be route/dimension-specific expression or partition --- indexes; silently padding or truncating vectors is never allowed. +-- ANN indexes on this undimensioned table must be route/dimension-specific +-- expression or partition indexes; silently padding or truncating vectors is +-- never allowed. Legacy embeddings_text/visual/face tables are indexed by +-- migration 0025; this table waits for the generation executor. CREATE TABLE index_generation_vectors ( generation_id uuid NOT NULL REFERENCES index_generations(id) ON DELETE CASCADE, workspace_id uuid NOT NULL REFERENCES workspaces(id) ON DELETE CASCADE, diff --git a/server/internal/db/migrations/0025_ann_hnsw_indexes.sql b/server/internal/db/migrations/0025_ann_hnsw_indexes.sql new file mode 100644 index 0000000..fdfb03d --- /dev/null +++ b/server/internal/db/migrations/0025_ann_hnsw_indexes.sql @@ -0,0 +1,47 @@ +-- +goose Up +-- Transactional startup migration: CREATE INDEX blocks writes on these tables +-- until commit. Schedule a maintenance window for a populated deployment. +-- +-- Operator class vector_cosine_ops matches the <=> operator used by search +-- and relator. pgvector defaults m=16, ef_construction=64; tune only after +-- representative recall/build/ingest measurements. +-- +-- Wrong-dimension vectors fail at INSERT/UPDATE against the fixed +-- vector(768)/vector(512) columns, before index maintenance. NULL embeddings +-- are allowed by the table DDL and are not present in a cosine HNSW index. +-- This file does not use CREATE INDEX CONCURRENTLY: a failed concurrent build +-- leaves an INVALID index that IF NOT EXISTS will skip. +-- +-- index_generation_vectors is intentionally not indexed (undimensioned; no +-- generation executor yet). Face clustering is still in-process; the face +-- index is DDL for a future SQL kNN, not a measured face-query speedup. +-- +-- Refs: https://github.com/bytefolk/mem/issues/173 + +-- +goose StatementBegin +CREATE INDEX IF NOT EXISTS idx_embeddings_text_embedding_hnsw + ON embeddings_text USING hnsw (embedding vector_cosine_ops); +-- +goose StatementEnd + +-- +goose StatementBegin +CREATE INDEX IF NOT EXISTS idx_embeddings_visual_embedding_hnsw + ON embeddings_visual USING hnsw (embedding vector_cosine_ops); +-- +goose StatementEnd + +-- +goose StatementBegin +CREATE INDEX IF NOT EXISTS idx_embeddings_face_embedding_hnsw + ON embeddings_face USING hnsw (embedding vector_cosine_ops); +-- +goose StatementEnd + +-- +goose Down +-- +goose StatementBegin +DROP INDEX IF EXISTS idx_embeddings_face_embedding_hnsw; +-- +goose StatementEnd + +-- +goose StatementBegin +DROP INDEX IF EXISTS idx_embeddings_visual_embedding_hnsw; +-- +goose StatementEnd + +-- +goose StatementBegin +DROP INDEX IF EXISTS idx_embeddings_text_embedding_hnsw; +-- +goose StatementEnd diff --git a/server/internal/face/face.go b/server/internal/face/face.go index 146ba99..af157e6 100644 --- a/server/internal/face/face.go +++ b/server/internal/face/face.go @@ -10,9 +10,10 @@ // to that entity. Otherwise we create a new entity (unnamed). // 4. Insert embeddings_face (file_id, entity_id, bbox, embedding). // -// This is intentionally O(n) per insert — fine for a personal drive up to -// thousands of faces. For larger corpora swap in pgvector HNSW + offline -// re-clustering. +// Clustering is intentionally O(n) per insert (in-process centroid distance). +// Migration 0025 adds a cosine HNSW index on embeddings_face; assignCluster +// does not query it yet. A future SQL kNN plus offline re-clustering can use +// that index without changing the 512-d insightface space. package face import ( diff --git a/server/internal/relator/relator.go b/server/internal/relator/relator.go index 30d009b..e61f868 100644 --- a/server/internal/relator/relator.go +++ b/server/internal/relator/relator.go @@ -178,8 +178,9 @@ func (s *Service) fileMeta(ctx context.Context, id uuid.UUID) (userID uuid.UUID, // recomputeText finds the top-K text-embedding nearest neighbors for srcID // (within the same user) and rewrites file_relations rows of type same_topic. // -// Strategy: take the first chunk of src as the seed; ANN against ALL chunks -// of OTHER files, DISTINCT ON dst file (best chunk wins). +// Strategy: take the first chunk of src as the seed; walk cosine-ordered +// chunks of OTHER files (HNSW-compatible) and fall back to DISTINCT ON when +// a bounded scan underfills after per-file deduplication. func (s *Service) recomputeText(ctx context.Context, srcID, userID uuid.UUID, topK int) error { tx, err := s.pool.Begin(ctx) if err != nil { @@ -194,47 +195,22 @@ func (s *Service) recomputeText(ctx context.Context, srcID, userID uuid.UUID, to return fmt.Errorf("clear: %w", err) } - rows, err := tx.Query(ctx, ` - WITH seed AS ( - SELECT embedding FROM embeddings_text - WHERE file_id = $1 AND chunk_index = 0 - LIMIT 1 - ) - SELECT DISTINCT ON (e.file_id) - e.file_id, - (1 - (e.embedding <=> (SELECT embedding FROM seed)))::real AS score - FROM embeddings_text e - JOIN files f ON f.id = e.file_id - WHERE f.user_id = $2 - AND e.file_id != $1 - AND (SELECT embedding FROM seed) IS NOT NULL - ORDER BY e.file_id, e.embedding <=> (SELECT embedding FROM seed) ASC - LIMIT $3 - `, srcID, userID, topK) + neighbors, err := textNeighbors(ctx, tx, srcID, userID, topK) if err != nil { return fmt.Errorf("knn: %w", err) } - defer rows.Close() batch := &pgx.Batch{} count := 0 - for rows.Next() { - var dstID uuid.UUID - var score float32 - if err := rows.Scan(&dstID, &score); err != nil { - return fmt.Errorf("scan: %w", err) - } + for _, n := range neighbors { batch.Queue(` INSERT INTO file_relations (src_id, dst_id, type, score, computed_at) VALUES ($1, $2, $3, $4, now()) ON CONFLICT (src_id, dst_id, type) DO UPDATE SET score = EXCLUDED.score, computed_at = EXCLUDED.computed_at - `, srcID, dstID, TypeSameTopic, score) + `, srcID, n.fileID, TypeSameTopic, n.score) count++ } - if err := rows.Err(); err != nil { - return err - } if count > 0 { br := tx.SendBatch(ctx, batch) for i := 0; i < count; i++ { @@ -250,6 +226,158 @@ func (s *Service) recomputeText(ctx context.Context, srcID, userID uuid.UUID, to return tx.Commit(ctx) } +type textNeighbor struct { + fileID uuid.UUID + score float32 +} + +// textNeighbors preserves best-chunk-per-file top-K. It walks cosine-order +// candidates (HNSW-compatible) and falls back to exact DISTINCT ON when a +// bounded scan underfills after per-file deduplication. +func textNeighbors(ctx context.Context, tx pgx.Tx, srcID, userID uuid.UUID, topK int) ([]textNeighbor, error) { + if topK <= 0 { + return nil, nil + } + selected := make([]uuid.UUID, 0, topK) + seen := make(map[uuid.UUID]struct{}, topK) + out := make([]textNeighbor, 0, topK) + for round := 0; round <= topK && len(out) < topK; round++ { + remaining := topK - len(out) + batch, err := queryTextNeighborsDistanceOrder(ctx, tx, srcID, userID, selected, remaining) + if err != nil { + return nil, err + } + added := 0 + for _, n := range batch { + if _, ok := seen[n.fileID]; ok { + continue + } + seen[n.fileID] = struct{}{} + selected = append(selected, n.fileID) + out = append(out, n) + added++ + if len(out) >= topK { + break + } + } + if added == 0 { + rest, err := queryTextNeighborsExact(ctx, tx, srcID, userID, selected, remaining) + if err != nil { + return nil, err + } + out = append(out, rest...) + break + } + } + if len(out) > topK { + out = out[:topK] + } + return out, nil +} + +func queryTextNeighborsDistanceOrder( + ctx context.Context, + tx pgx.Tx, + srcID, userID uuid.UUID, + selected []uuid.UUID, + limit int, +) ([]textNeighbor, error) { + var seed string + err := tx.QueryRow(ctx, ` + SELECT embedding::text FROM embeddings_text + WHERE file_id = $1 AND chunk_index = 0 + LIMIT 1 + `, srcID).Scan(&seed) + if err == pgx.ErrNoRows { + return nil, nil + } + if err != nil { + return nil, err + } + excludeSQL := "TRUE" + args := []any{seed, srcID, userID, limit} + if len(selected) > 0 { + args = append(args, selected) + excludeSQL = fmt.Sprintf("NOT (e.file_id = ANY($%d::uuid[]))", len(args)) + } + rows, err := tx.Query(ctx, fmt.Sprintf(` + WITH nearest AS ( + SELECT e.file_id, e.embedding <=> $1::vector AS dist + FROM embeddings_text e + WHERE e.file_id != $2 + AND %s + ORDER BY e.embedding <=> $1::vector ASC + LIMIT $4 + ) + SELECT n.file_id, (1 - n.dist)::real AS score + FROM nearest n + JOIN files f ON f.id = n.file_id + WHERE f.user_id = $3 + ORDER BY n.dist ASC + `, excludeSQL), args...) + if err != nil { + return nil, err + } + defer rows.Close() + return scanTextNeighbors(rows) +} + +func queryTextNeighborsExact( + ctx context.Context, + tx pgx.Tx, + srcID, userID uuid.UUID, + selected []uuid.UUID, + limit int, +) ([]textNeighbor, error) { + excludeSQL := "TRUE" + args := []any{srcID, userID, limit} + if len(selected) > 0 { + args = append(args, selected) + excludeSQL = fmt.Sprintf("NOT (e.file_id = ANY($%d::uuid[]))", len(args)) + } + rows, err := tx.Query(ctx, fmt.Sprintf(` + WITH seed AS ( + SELECT embedding FROM embeddings_text + WHERE file_id = $1 AND chunk_index = 0 + LIMIT 1 + ) + SELECT file_id, score FROM ( + SELECT DISTINCT ON (e.file_id) + e.file_id, + (1 - (e.embedding <=> (SELECT embedding FROM seed)))::real AS score + FROM embeddings_text e + JOIN files f ON f.id = e.file_id + WHERE f.user_id = $2 + AND e.file_id != $1 + AND (SELECT embedding FROM seed) IS NOT NULL + AND %s + ORDER BY e.file_id, e.embedding <=> (SELECT embedding FROM seed) ASC + ) hits + ORDER BY score DESC + LIMIT $3 + `, excludeSQL), args...) + if err != nil { + return nil, err + } + defer rows.Close() + return scanTextNeighbors(rows) +} + +func scanTextNeighbors(rows pgx.Rows) ([]textNeighbor, error) { + out := make([]textNeighbor, 0, 8) + for rows.Next() { + var n textNeighbor + if err := rows.Scan(&n.fileID, &n.score); err != nil { + return nil, fmt.Errorf("scan: %w", err) + } + out = append(out, n) + } + if err := rows.Err(); err != nil { + return nil, err + } + return out, nil +} + func (s *Service) recomputeVisual(ctx context.Context, srcID, userID uuid.UUID, topK int) error { tx, err := s.pool.Begin(ctx) if err != nil { diff --git a/server/internal/search/hnsw_semantics_test.go b/server/internal/search/hnsw_semantics_test.go new file mode 100644 index 0000000..2f69ccc --- /dev/null +++ b/server/internal/search/hnsw_semantics_test.go @@ -0,0 +1,116 @@ +package search + +import ( + "context" + "os" + "strings" + "testing" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" + + memdb "github.com/PeterGuy326/mem/server/internal/db" +) + +// A chunk candidate budget must not become a file result budget or bypass +// owner/path/MIME/time scope. One file holding 101 nearest chunks must still +// yield k distinct eligible files, using HNSW continuation plus exact fallback. +func TestTextANNFileSemanticsPostgres(t *testing.T) { + dsn := os.Getenv("MEM_TEST_DB") + if dsn == "" { + t.Skip("MEM_TEST_DB not set; skipping text ANN PostgreSQL regression") + } + cfg, err := pgxpool.ParseConfig(dsn) + if err != nil { + t.Fatal(err) + } + if !strings.HasSuffix(cfg.ConnConfig.Database, "_test") { + t.Fatalf("refusing non-test database %q", cfg.ConnConfig.Database) + } + ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second) + defer cancel() + database, err := memdb.Open(ctx, dsn) + if err != nil { + t.Fatal(err) + } + defer database.Close() + if err := database.Migrate(ctx); err != nil { + t.Fatal(err) + } + owner, other := uuid.New(), uuid.New() + for _, id := range []uuid.UUID{owner, other} { + if _, err := database.Pool.Exec(ctx, "INSERT INTO users(id,email,password_hash) VALUES($1,$2,'test')", id, id.String()+"@example.test"); err != nil { + t.Fatal(err) + } + } + defer func() { + cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cleanupCancel() + if _, err := database.Pool.Exec(cleanupCtx, "DELETE FROM users WHERE id=ANY($1::uuid[])", []uuid.UUID{owner, other}); err != nil { + t.Errorf("cleanup: %v", err) + } + }() + now := time.Now().UTC().Truncate(time.Second) + since, until := now.Add(-time.Hour), now.Add(time.Hour) + vec := make([]float32, textEmbeddingSchemaDim) + vec[0] = 1 + addFile := func(user uuid.UUID, path, mime string, at time.Time, chunks int, far bool) uuid.UUID { + t.Helper() + id := uuid.New() + _, err := database.Pool.Exec(ctx, `INSERT INTO files(id,user_id,name,path,size,sha256,mime,storage_key,created_at,timeline_at) + VALUES($1,$2,'fixture',$3,0,$4,$5,$6,$7,$7)`, id, user, path, id.String(), mime, id.String(), at) + if err != nil { + t.Fatal(err) + } + for chunk := 0; chunk < chunks; chunk++ { + v := make([]float32, textEmbeddingSchemaDim) + if far { + v[1] = 1 + } else { + v[0], v[1] = 1, float32(chunk)*0.001 + } + _, err := database.Pool.Exec(ctx, `INSERT INTO embeddings_text(file_id,chunk_index,chunk_text,embedding) + VALUES($1,$2,'source chunk',$3::vector)`, id, chunk, vectorLiteral(v)) + if err != nil { + t.Fatal(err) + } + } + return id + } + best := addFile(owner, "/Work_%/Docs", "text/plain", now, 101, false) + eligible := map[uuid.UUID]bool{best: true} + for i := 0; i < 19; i++ { + eligible[addFile(owner, "/Work_%/Docs", "application/pdf", now, 1, true)] = true + } + addFile(other, "/Work_%/Docs", "text/plain", now, 1, false) + addFile(owner, "/Work_AB/Docs", "text/plain", now, 1, false) + addFile(owner, "/Work_%/Private", "text/plain", now, 1, false) + addFile(owner, "/Work_%/Docs", "image/png", now, 1, false) + addFile(owner, "/Work_%/Docs", "text/plain", since.Add(-time.Second), 1, false) + addFile(owner, "/Work_%/Docs", "text/plain", until.Add(time.Second), 1, false) + service := New(database.Pool, nil) + q := Query{UserID: owner, Limit: 10, PathPrefix: "/Work_%", AllowedPaths: []string{"/Work_%/Docs"}, Type: "doc", Since: &since, Until: &until, SnippetChars: 200} + hits, err := service.runTextANN(ctx, q, vec) + if err != nil { + t.Fatal(err) + } + if len(hits) != 10 { + t.Fatalf("got %d files, want 10 despite 101 nearest chunks belonging to one file", len(hits)) + } + seen := map[uuid.UUID]bool{} + for _, hit := range hits { + if !eligible[hit.FileID] || seen[hit.FileID] { + t.Fatalf("out-of-scope or duplicate file: %+v", hit) + } + seen[hit.FileID] = true + } + if hits[0].FileID != best || hits[0].ChunkIndex != 0 || hits[0].Score != 1 { + t.Fatalf("best chunk was not preserved: %+v", hits[0]) + } + q.AllowedPaths = []string{""} + hits, err = service.runTextANN(ctx, q, vec) + if err != nil || len(hits) != 0 { + t.Fatalf("invalid allow-list must fail closed: hits=%v err=%v", hits, err) + } +} diff --git a/server/internal/search/search.go b/server/internal/search/search.go index 5f0152a..899a635 100644 --- a/server/internal/search/search.go +++ b/server/internal/search/search.go @@ -640,23 +640,134 @@ func (s *Service) mergeAutoResults(q Query, tr, vr autoResult) ([]Hit, error) { return out, nil } -// runTextANN issues the text-route SQL and scans results. +// runTextANN returns the k files whose best chunk is nearest the query. +// +// The previous DISTINCT ON (f.id) ORDER BY f.id, distance shape cannot use a +// cosine HNSW index. The shipping path now walks globally ordered chunks +// (HNSW-compatible ORDER BY distance LIMIT n), keeps the first sighting of +// each file (that chunk is the file's best), and excludes selected files on +// the next round. If a bounded approximate scan underfills — one file owning +// many near chunks, or post-filters emptying the HNSW candidate list — the +// original exact DISTINCT ON query fills the remaining slots. func (s *Service) runTextANN(ctx context.Context, q Query, vec []float32) ([]Hit, error) { + const maxTextANNLimit = 100 + if q.Limit <= 0 { + q.Limit = 10 + } + if q.Limit > maxTextANNLimit { + q.Limit = maxTextANNLimit + } args := []any{vectorLiteral(vec), q.UserID} where := []string{"f.user_id = $2"} args, where = appendPathFilters(args, where, q.PathPrefix, q.AllowedPaths) args, where = appendMIMEFilter(args, where, q.Type) - if q.Since != nil { - args = append(args, *q.Since) - where = append(where, fmt.Sprintf("COALESCE(f.timeline_at, f.created_at) >= $%d", len(args))) + args, where = appendTimeFilters(args, where, q.Since, q.Until) + + // Constant caps: CodeQL still treats a sanitized q.Limit as user-controlled. + selected := make([]uuid.UUID, 0, maxTextANNLimit) + seen := make(map[uuid.UUID]struct{}, maxTextANNLimit) + out := make([]Hit, 0, maxTextANNLimit) + for round := 0; round <= q.Limit && len(out) < q.Limit; round++ { + remaining := q.Limit - len(out) + batch, err := s.queryTextDistanceOrder(ctx, q, args, where, selected, remaining) + if err != nil { + return nil, err + } + added := 0 + for _, h := range batch { + if _, ok := seen[h.FileID]; ok { + continue + } + seen[h.FileID] = struct{}{} + selected = append(selected, h.FileID) + out = append(out, h) + added++ + if len(out) >= q.Limit { + break + } + } + if added == 0 { + rest, err := s.queryTextExactRemaining(ctx, q, args, where, selected, remaining) + if err != nil { + return nil, err + } + out = append(out, rest...) + break + } } - if q.Until != nil { - args = append(args, *q.Until) - where = append(where, fmt.Sprintf("COALESCE(f.timeline_at, f.created_at) <= $%d", len(args))) + sortHitsByScoreDesc(out) + if len(out) > q.Limit { + out = out[:q.Limit] } - args = append(args, q.Limit) - limitIdx := len(args) + return out, nil +} + +func cloneArgs(args []any) []any { + out := make([]any, len(args)) + copy(out, args) + return out +} +func excludeFileIDs(selected []uuid.UUID) []uuid.UUID { + if selected == nil { + return []uuid.UUID{} + } + return selected +} + +func (s *Service) queryTextDistanceOrder( + ctx context.Context, + q Query, + args []any, + where []string, + selected []uuid.UUID, + limit int, +) ([]Hit, error) { + queryArgs := cloneArgs(args) + excludeSQL := "TRUE" + if len(selected) > 0 { + queryArgs = append(queryArgs, selected) + excludeSQL = fmt.Sprintf("NOT (e.file_id = ANY($%d::uuid[]))", len(queryArgs)) + } + queryArgs = append(queryArgs, limit) + limitIdx := len(queryArgs) + // ANN first so the planner can use HNSW; file filters apply after. + sql := fmt.Sprintf(` + WITH nearest AS ( + SELECT e.id, e.file_id, e.chunk_index, e.chunk_text, + e.embedding <=> $1::vector AS dist + FROM embeddings_text e + WHERE %s + ORDER BY e.embedding <=> $1::vector ASC + LIMIT $%d + ) + SELECT e.id::text, f.id, f.name, f.path, f.mime, f.sha256, + e.chunk_index, (1 - e.dist) AS score, e.chunk_text, f.summary, + f.timeline_at, f.created_at + FROM nearest e + JOIN files f ON f.id = e.file_id + WHERE %s + ORDER BY e.dist ASC + `, excludeSQL, limitIdx, strings.Join(where, " AND ")) + return s.scanHits(ctx, sql, queryArgs, RouteText, q.SnippetChars) +} + +func (s *Service) queryTextExactRemaining( + ctx context.Context, + q Query, + args []any, + where []string, + selected []uuid.UUID, + limit int, +) ([]Hit, error) { + queryArgs := cloneArgs(args) + excludeSQL := "TRUE" + if len(selected) > 0 { + queryArgs = append(queryArgs, selected) + excludeSQL = fmt.Sprintf("NOT (e.file_id = ANY($%d::uuid[]))", len(queryArgs)) + } + queryArgs = append(queryArgs, limit) + limitIdx := len(queryArgs) sql := fmt.Sprintf(` SELECT evidence_id, file_id, name, path, mime, content_sha256, chunk_index, score, snippet, summary, timeline_at, created_at @@ -676,14 +787,14 @@ func (s *Service) runTextANN(ctx context.Context, q Query, vec []float32) ([]Hit f.created_at AS created_at FROM embeddings_text e JOIN files f ON f.id = e.file_id - WHERE %s - ORDER BY f.id, e.embedding <=> $1::vector ASC + WHERE %s + AND %s + ORDER BY f.id, e.embedding <=> $1::vector ASC ) hits ORDER BY score DESC LIMIT $%d - `, strings.Join(where, " AND "), limitIdx) - - return s.scanHits(ctx, sql, args, RouteText, q.SnippetChars) + `, strings.Join(where, " AND "), excludeSQL, limitIdx) + return s.scanHits(ctx, sql, queryArgs, RouteText, q.SnippetChars) } // runVisualANN issues the visual-route SQL.