diff --git a/Makefile b/Makefile index 40eaf25..98a7d2f 100644 --- a/Makefile +++ b/Makefile @@ -4,6 +4,13 @@ export CGO_ENABLED?=0 IMAGE?=quay.io/mudler/localrecall:latest +# Embedding model the test suites exercise. docker-compose starts LocalAI with +# this model, but it is downloaded/loaded lazily on first request - so the test +# harness must wait for it to actually answer before running specs (see +# wait-localai), otherwise the first embedding call races the download and the +# whole suite fails with "model not found". +EMBEDDING_MODEL?=granite-embedding-107m-multilingual + print-version: @echo "Version: ${VERSION}" @@ -55,6 +62,23 @@ wait-localai: docker compose logs localai | tail -20; \ exit 1; \ fi + @echo "Warming embedding model '$(EMBEDDING_MODEL)' (downloaded lazily on first use)..." + @timeout=600; \ + while [ $$timeout -gt 0 ]; do \ + if curl -fsS http://localhost:8081/v1/embeddings \ + -H 'Content-Type: application/json' \ + -d '{"model":"$(EMBEDDING_MODEL)","input":"warmup"}' >/dev/null 2>&1; then \ + echo "Embedding model '$(EMBEDDING_MODEL)' is ready"; \ + break; \ + fi; \ + sleep 3; \ + timeout=$$((timeout - 3)); \ + done; \ + if [ $$timeout -le 0 ]; then \ + echo "Error: embedding model '$(EMBEDDING_MODEL)' did not load in time"; \ + docker compose logs localai | tail -30; \ + exit 1; \ + fi # Start all test services (LocalAI and PostgreSQL from docker-compose) start-test-services: diff --git a/rag/engine/postgres.go b/rag/engine/postgres.go index 61d799d..43d6901 100644 --- a/rag/engine/postgres.go +++ b/rag/engine/postgres.go @@ -698,6 +698,74 @@ func (p *PostgresDB) GetBySource(source string) ([]types.Result, error) { return results, nil } +// hybridCandidateMultiplier controls how many candidates each retrieval arm +// (vector and BM25) pulls before fusion: max(limit*multiplier, floor). A larger +// pool trades latency for recall; 10x with a floor of 100 keeps recall high for +// typical small result limits while staying index-bound. +const ( + hybridCandidateMultiplier = 10 + hybridCandidateFloor = 100 + // rrfK is the Reciprocal Rank Fusion smoothing constant. 60 is the value used + // across the IR literature and by pgvector/Timescale's reference hybrid-search + // examples; it damps the influence of the very top ranks so the two arms blend + // smoothly. + rrfK = 60 +) + +// buildHybridSearchQuery returns the SQL for hybrid (BM25 + vector) search over +// the documents table. Bind parameters: $1 query text, $2 BM25 weight, +// $3 query embedding (::vector), $4 vector weight, $5 result limit. +// +// It follows the canonical Reciprocal Rank Fusion pattern recommended by both +// pgvector and Timescale (pg_textsearch + pgvectorscale). Each arm retrieves its +// top-N candidates via a *bare* operator - "ORDER BY embedding <=> $vec" and +// "ORDER BY full_text <@> to_bm25query(...)" - which pgvector's HNSW/DiskANN and +// the BM25 index serve directly, and assigns a rank. The arms are combined with a +// FULL OUTER JOIN and scored by weighted RRF (sum of weight/(rrfK+rank)); the +// final id list is joined back to the table by primary key to fetch payloads. +// +// RRF fuses by *rank*, not raw score, which avoids mixing BM25's unbounded scores +// with cosine similarity's [0,1] range. The previous query sorted on a wrapped +// scalar similarity expression in a single stage, which blinded the planner into +// a full sequential scan over every row and exceeded the statement timeout on +// multi-million-row collections (LocalAI issue #10186). $2/$4 weight each arm +// (equal by default), so an arm can be biased without breaking the index path. +func buildHybridSearchQuery(tableName string) string { + candidatePool := fmt.Sprintf("GREATEST($5 * %d, %d)", hybridCandidateMultiplier, hybridCandidateFloor) + return fmt.Sprintf(` + WITH bm25_results AS ( + SELECT id, ROW_NUMBER() OVER (ORDER BY full_text <@> to_bm25query($1, 'idx_%[1]s_bm25')) AS rank + FROM %[1]s + ORDER BY full_text <@> to_bm25query($1, 'idx_%[1]s_bm25') + LIMIT %[2]s + ), + vector_results AS ( + SELECT id, ROW_NUMBER() OVER (ORDER BY embedding <=> $3::vector) AS rank + FROM %[1]s + WHERE embedding IS NOT NULL + ORDER BY embedding <=> $3::vector + LIMIT %[2]s + ), + fused AS ( + SELECT + COALESCE(b.id, v.id) AS id, + COALESCE($2 / (%[3]d + b.rank), 0) + COALESCE($4 / (%[3]d + v.rank), 0) AS similarity + FROM bm25_results b + FULL OUTER JOIN vector_results v ON b.id = v.id + ) + SELECT + d.id::text, + COALESCE(d.title, '') as title, + d.content, + d.metadata, + f.similarity + FROM fused f + JOIN %[1]s d ON d.id = f.id + ORDER BY f.similarity DESC + LIMIT $5 + `, tableName, candidatePool, rrfK) +} + func (p *PostgresDB) Search(s string, similarEntries int) ([]types.Result, error) { ctx := context.Background() @@ -708,23 +776,8 @@ func (p *PostgresDB) Search(s string, similarEntries int) ([]types.Result, error } queryEmbeddingStr := formatVector(queryEmbedding) - // Build hybrid search query - // Combine BM25 score and vector similarity - query := fmt.Sprintf(` - SELECT - id::text, - COALESCE(title, '') as title, - content, - metadata, - ( - COALESCE(-(full_text <@> to_bm25query($1, 'idx_%s_bm25')), 0) * $2 + - COALESCE((1 - (embedding <=> $3::vector)), 0) * $4 - ) as similarity - FROM %s - WHERE embedding IS NOT NULL - ORDER BY similarity DESC - LIMIT $5 - `, p.tableName, p.tableName) + // Build hybrid search query (BM25 + vector similarity) + query := buildHybridSearchQuery(p.tableName) rows, err := p.pool.Query(ctx, query, s, p.bm25Weight, queryEmbeddingStr, p.vectorWeight, similarEntries) if err != nil { diff --git a/rag/engine/postgres_planning_test.go b/rag/engine/postgres_planning_test.go new file mode 100644 index 0000000..a01c2cf --- /dev/null +++ b/rag/engine/postgres_planning_test.go @@ -0,0 +1,127 @@ +package engine + +import ( + "context" + "fmt" + "os" + "strings" + + "github.com/jackc/pgx/v5/pgxpool" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// Regression test for LocalAI issue #10186: the hybrid search query wrapped the +// vector distance operator in a scalar similarity expression and sorted on the +// alias (ORDER BY similarity DESC). That blinds the planner: pgvector's +// HNSW/DiskANN index can only serve a bare "ORDER BY embedding <=> $vec" path, +// so the wrapped form degrades to a full sequential scan over every row and, on +// multi-million-row tables, blows past the statement timeout. +// +// This is a query-planning test, so it only needs a PostgreSQL with pg_textsearch +// + pgvector/vectorscale (the docker-compose stack). It builds the schema via the +// real setupDatabase() and asserts, through EXPLAIN, that the vector ordering is +// served by the index rather than a full scan. No embedding model is required. +var _ = Describe("hybrid search query planning (LocalAI issue #10186)", func() { + var ( + ctx context.Context + pool *pgxpool.Pool + tableName string + queryVec string + ) + + BeforeEach(func() { + ctx = context.Background() + + databaseURL := os.Getenv("POSTGRES_TEST_URL") + if databaseURL == "" { + databaseURL = "postgresql://localrecall:localrecall@localhost:5432/localrecall?sslmode=disable" + } + + var err error + pool, err = pgxpool.New(ctx, databaseURL) + Expect(err).ToNot(HaveOccurred()) + Expect(pool.Ping(ctx)).To(Succeed(), + "PostgreSQL with pg_textsearch + pgvector must be reachable for the planning test") + + const dims = 8 + collectionName := "plan10186" + p := &PostgresDB{ + pool: pool, + collectionName: collectionName, + tableName: sanitizeTableName(collectionName), + embeddingDims: dims, + bm25Weight: 0.5, + vectorWeight: 0.5, + } + tableName = p.tableName + + // Start from a clean slate so setupDatabase()'s CREATE TABLE IF NOT EXISTS + // builds the table at the dimensions this test expects. + _, err = pool.Exec(ctx, "DROP TABLE IF EXISTS "+tableName) + Expect(err).ToNot(HaveOccurred()) + Expect(p.setupDatabase()).To(Succeed()) + + // Seed enough rows that an index path is the cheap plan. Random vectors are + // inserted directly: a planning test needs row shape, not real embeddings. + _, err = pool.Exec(ctx, fmt.Sprintf(` + INSERT INTO %s (title, content, embedding) + SELECT 'title '||g, 'content number '||g, + ('['||random()||','||random()||','||random()||','||random()||','|| + random()||','||random()||','||random()||','||random()||']')::vector + FROM generate_series(1, 2000) g + `, tableName)) + Expect(err).ToNot(HaveOccurred()) + _, err = pool.Exec(ctx, "ANALYZE "+tableName) + Expect(err).ToNot(HaveOccurred()) + + queryVec = "[0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8]" + }) + + AfterEach(func() { + if pool != nil { + _, _ = pool.Exec(ctx, "DROP TABLE IF EXISTS "+tableName) + pool.Close() + } + }) + + It("serves the vector ordering from the index instead of a full sequential scan", func() { + query := buildHybridSearchQuery(tableName) + + // Disable plain sequential scans so the planner is forced onto an index + // path *if the query is index-compatible*. The buggy wrapped-scalar ORDER + // BY cannot be served by the vector index and still falls back to a full + // (disabled, high-cost) scan, which this test catches. + tx, err := pool.Begin(ctx) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = tx.Rollback(ctx) }() + + _, err = tx.Exec(ctx, "SET LOCAL enable_seqscan = off") + Expect(err).ToNot(HaveOccurred()) + + rows, err := tx.Query(ctx, "EXPLAIN "+query, "content", 0.5, queryVec, 0.5, 5) + Expect(err).ToNot(HaveOccurred()) + + var plan strings.Builder + for rows.Next() { + var line string + Expect(rows.Scan(&line)).To(Succeed()) + plan.WriteString(line) + plan.WriteByte('\n') + } + Expect(rows.Err()).ToNot(HaveOccurred()) + planText := plan.String() + AddReportEntry("EXPLAIN plan", planText) + + // pgvector/vectorscale emit an index "Order By:" condition on the distance + // operator only when the index actually serves the nearest-neighbour + // ordering. The buggy query produces a "Sort Key:" on the wrapped scalar + // instead, and never this line. + Expect(planText).To(ContainSubstring("Order By: (embedding <=>"), + "the vector nearest-neighbour ordering must be served by the index") + + // And the documents table must never be sequentially scanned for a search. + Expect(planText).ToNot(ContainSubstring("Seq Scan on "+tableName), + "the documents table must not be sequentially scanned") + }) +})