fix(postgres): index-backed RRF hybrid search (#10186)

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:

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 {

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")
+	})
+})