feat(brain): hybrid BM25 + pgvector retrieval (opt-in)
Wires nomic-embed-text (iguana ollama) + pgvector on the shared
postgres18 into brain_query / brain_answer via Reciprocal Rank Fusion.
Pure BM25 stays the default; setting BRAIN_PG_DSN and BRAIN_EMBED_URL
together opts in. Setting one without the other is misconfiguration →
exit 1.
New packages:
- internal/embed
Client.Embed(ctx, text) → []float32 via POST {URL}/api/embed.
Defaults to nomic-embed-text:latest (768 dim). nil-on-empty-URL so
callers gate on a single nil check.
- internal/vectorstore
PGStore wraps a pgxpool against postgres18. Init creates
brain_embeddings(path PK, vector(768), updated_at) + HNSW cosine
index idempotently. Upsert / Delete / Search / KnownPaths.
Sync(brainDir, store, embedder) diffs brain/wiki/ against the store
and upserts new files / deletes removed ones; StartSync runs it on
a ticker (default 300s). Integration tests gated by BRAIN_PG_TEST_DSN.
- scripts/brain-embeddings-init.sql
One-time DBA setup: brain DB, brain_app role, vector extension,
GRANTs. Idempotent.
Search layer:
- search.QueryOptions gains Vector + Embedder fields.
- QueryContext is the cancellable variant; Query stays for callers.
- When both are set, BM25 (top-N) and pgvector (top-4N) candidates
merge via Reciprocal Rank Fusion (k=60, Cormack et al. 2009 — no
tuning knob, robust to scale differences between rankers).
- Vector-only hits are hydrated from disk so callers see uniform
Result records (path, title, excerpt, wing, hall, score).
- Wing/hall filters still apply to vector candidates via path-prefix.
- On embedder/vector errors the search falls back to BM25 — embedding
outage degrades quality but doesn't take the brain offline.
MCP wiring:
- mcp.Server.WithHybridRetrieval(v, e) opt-in setter, same shape as
WithReranker.
- brainQuery and brainAnswer pass the wired vector/embedder through
to search.QueryContext.
REST:
- POST /backfill-embeddings drives Sync synchronously. Returns
{added, deleted, errors[]}. 503 when feature is unconfigured.
cmd/server/main.go:
- BRAIN_PG_DSN + BRAIN_EMBED_URL together enable hybrid; one alone
→ exit 1.
- vectorAdapter bridges *PGStore (returns []Hit) to
search.VectorSearcher (which takes []VectorHit) without either
package importing the other.
- BRAIN_EMBED_SYNC_INTERVAL (default 300s) controls the background
Sync ticker.
Backend pivot from Qdrant to pgvector recorded in DECISIONS.md
2026-05-18 (supersedes 2026-04-08): postgres18 already runs in
databases/ ns, Qdrant was never deployed, one engine beats two.
Dependency: github.com/jackc/pgx/v5 — modern, native pgvector via
parametric vector literals.
Tests:
- embed.Client: empty-URL nil, request shape, dimension, upstream
error propagation, empty-text rejection.
- vectorstore.PGStore: dimension validation (unit); upsert/search/
KnownPaths (integration, BRAIN_PG_TEST_DSN-gated).
- vectorstore.Sync: adds new files, skips known, deletes
disappeared, skips _index.md, no-op when nil, collects embedder
errors.
- search.Query: hybrid promotes vector-only hits via RRF; falls
back to BM25 on embedder error.
Closes hyperguild#8.
This commit is contained in:
155
ingestion/internal/vectorstore/pg.go
Normal file
155
ingestion/internal/vectorstore/pg.go
Normal file
@@ -0,0 +1,155 @@
|
||||
// Package vectorstore stores brain note embeddings in pgvector on the
|
||||
// shared postgres18 instance. One row per markdown path, cosine-distance
|
||||
// indexed via HNSW for sub-millisecond top-k retrieval.
|
||||
package vectorstore
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
// Hit is a single result from a cosine-distance search.
|
||||
type Hit struct {
|
||||
Path string
|
||||
Distance float64 // 0 = identical, 2 = opposite
|
||||
}
|
||||
|
||||
// PGStore is a pgvector-backed embeddings store. Construct with New and
|
||||
// call Init once to create the table + HNSW index. Use Close to release
|
||||
// the underlying pool.
|
||||
type PGStore struct {
|
||||
pool *pgxpool.Pool
|
||||
}
|
||||
|
||||
// New opens a connection pool against dsn (a libpq-style URL). Caller
|
||||
// owns the resulting *PGStore and must invoke Close.
|
||||
func New(ctx context.Context, dsn string) (*PGStore, error) {
|
||||
pool, err := pgxpool.New(ctx, dsn)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("pgxpool: %w", err)
|
||||
}
|
||||
if err := pool.Ping(ctx); err != nil {
|
||||
pool.Close()
|
||||
return nil, fmt.Errorf("ping: %w", err)
|
||||
}
|
||||
return &PGStore{pool: pool}, nil
|
||||
}
|
||||
|
||||
// Close releases the underlying connection pool.
|
||||
func (s *PGStore) Close() {
|
||||
if s.pool != nil {
|
||||
s.pool.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// Init creates the brain_embeddings table and its HNSW index if they
|
||||
// don't already exist. Safe to call on every startup. Assumes the
|
||||
// `vector` extension is already installed (one-time DBA setup; see
|
||||
// scripts/brain-embeddings-init.sql).
|
||||
func (s *PGStore) Init(ctx context.Context) error {
|
||||
const ddl = `
|
||||
CREATE TABLE IF NOT EXISTS brain_embeddings (
|
||||
path TEXT PRIMARY KEY,
|
||||
embedding vector(768) NOT NULL,
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS brain_embeddings_embedding_idx
|
||||
ON brain_embeddings USING hnsw (embedding vector_cosine_ops);
|
||||
`
|
||||
_, err := s.pool.Exec(ctx, ddl)
|
||||
return err
|
||||
}
|
||||
|
||||
// Upsert inserts or replaces the embedding for path. Embedding must be
|
||||
// 768-dim (nomic-embed-text). Caller is responsible for normalising
|
||||
// paths to forward-slash form.
|
||||
func (s *PGStore) Upsert(ctx context.Context, path string, embedding []float32) error {
|
||||
if len(embedding) != 768 {
|
||||
return fmt.Errorf("expected 768-dim embedding, got %d", len(embedding))
|
||||
}
|
||||
_, err := s.pool.Exec(ctx, `
|
||||
INSERT INTO brain_embeddings (path, embedding, updated_at)
|
||||
VALUES ($1, $2, now())
|
||||
ON CONFLICT (path) DO UPDATE
|
||||
SET embedding = EXCLUDED.embedding, updated_at = now()
|
||||
`, path, vectorLiteral(embedding))
|
||||
return err
|
||||
}
|
||||
|
||||
// Delete removes the row at path. No-op when the row doesn't exist.
|
||||
func (s *PGStore) Delete(ctx context.Context, path string) error {
|
||||
_, err := s.pool.Exec(ctx, `DELETE FROM brain_embeddings WHERE path = $1`, path)
|
||||
return err
|
||||
}
|
||||
|
||||
// Search returns the top-limit nearest paths by cosine distance.
|
||||
func (s *PGStore) Search(ctx context.Context, query []float32, limit int) ([]Hit, error) {
|
||||
if len(query) != 768 {
|
||||
return nil, fmt.Errorf("expected 768-dim query, got %d", len(query))
|
||||
}
|
||||
if limit <= 0 {
|
||||
limit = 10
|
||||
}
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT path, embedding <=> $1 AS distance
|
||||
FROM brain_embeddings
|
||||
ORDER BY embedding <=> $1
|
||||
LIMIT $2
|
||||
`, vectorLiteral(query), limit)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var hits []Hit
|
||||
for rows.Next() {
|
||||
var h Hit
|
||||
if err := rows.Scan(&h.Path, &h.Distance); err != nil {
|
||||
return nil, fmt.Errorf("scan: %w", err)
|
||||
}
|
||||
hits = append(hits, h)
|
||||
}
|
||||
if err := rows.Err(); err != nil && !errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil, err
|
||||
}
|
||||
return hits, nil
|
||||
}
|
||||
|
||||
// KnownPaths returns the path set already present in the store. Used by
|
||||
// the watcher to diff against the wiki/ tree and decide what to upsert.
|
||||
func (s *PGStore) KnownPaths(ctx context.Context) (map[string]struct{}, error) {
|
||||
rows, err := s.pool.Query(ctx, `SELECT path FROM brain_embeddings`)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query paths: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
out := make(map[string]struct{})
|
||||
for rows.Next() {
|
||||
var p string
|
||||
if err := rows.Scan(&p); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[p] = struct{}{}
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// vectorLiteral renders a Go float32 slice as the literal representation
|
||||
// pgvector accepts as a parametric input: `[v1,v2,...,vN]`.
|
||||
func vectorLiteral(v []float32) string {
|
||||
var b strings.Builder
|
||||
b.WriteByte('[')
|
||||
for i, x := range v {
|
||||
if i > 0 {
|
||||
b.WriteByte(',')
|
||||
}
|
||||
fmt.Fprintf(&b, "%g", x)
|
||||
}
|
||||
b.WriteByte(']')
|
||||
return b.String()
|
||||
}
|
||||
Reference in New Issue
Block a user