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:
@@ -3,6 +3,7 @@ package search
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
@@ -13,6 +14,26 @@ import (
|
||||
"github.com/mathiasbq/hyperguild/ingestion/internal/brain"
|
||||
)
|
||||
|
||||
// VectorSearcher returns the top-limit nearest paths by cosine
|
||||
// distance. The vectorstore package implements this against pgvector.
|
||||
type VectorSearcher interface {
|
||||
Search(ctx context.Context, query []float32, limit int) ([]VectorHit, error)
|
||||
}
|
||||
|
||||
// VectorHit is a single path + distance pair from a vector search.
|
||||
// Re-declared here (rather than imported) to keep search package
|
||||
// free of vectorstore/embed deps and to make stubbing trivial in tests.
|
||||
type VectorHit struct {
|
||||
Path string
|
||||
Distance float64
|
||||
}
|
||||
|
||||
// Embedder turns a query string into a dense vector. The embed package
|
||||
// implements this against Ollama's /api/embed.
|
||||
type Embedder interface {
|
||||
Embed(ctx context.Context, text string) ([]float32, error)
|
||||
}
|
||||
|
||||
// Result is a single search hit from the brain wiki.
|
||||
type Result struct {
|
||||
Path string `json:"path"`
|
||||
@@ -29,16 +50,30 @@ type Result struct {
|
||||
// When Hall is additionally set, the walk is restricted to
|
||||
// brain/wiki/<wing>/<hall>/. Without either, the legacy walk over
|
||||
// brain/knowledge/ and brain/wiki/ is used.
|
||||
//
|
||||
// When both Vector and Embedder are non-nil, results are computed
|
||||
// hybridly: BM25 and vector candidate lists are merged via Reciprocal
|
||||
// Rank Fusion. With either nil the function falls back to BM25 only,
|
||||
// keeping behaviour unchanged for callers that have not opted in.
|
||||
type QueryOptions struct {
|
||||
Query string
|
||||
Limit int
|
||||
Wing string
|
||||
Hall string
|
||||
Query string
|
||||
Limit int
|
||||
Wing string
|
||||
Hall string
|
||||
Vector VectorSearcher
|
||||
Embedder Embedder
|
||||
}
|
||||
|
||||
// Query searches the brain. Returns up to opts.Limit results sorted by
|
||||
// score descending. Empty query returns nil.
|
||||
func Query(brainDir string, opts QueryOptions) ([]Result, error) {
|
||||
return QueryContext(context.Background(), brainDir, opts)
|
||||
}
|
||||
|
||||
// QueryContext is the cancellable variant of Query. Hybrid retrieval
|
||||
// requires a context because both the embedder and the vector store are
|
||||
// network calls.
|
||||
func QueryContext(ctx context.Context, brainDir string, opts QueryOptions) ([]Result, error) {
|
||||
if opts.Limit <= 0 {
|
||||
opts.Limit = 5
|
||||
}
|
||||
@@ -102,12 +137,108 @@ func Query(brainDir string, opts QueryOptions) ([]Result, error) {
|
||||
sort.Slice(results, func(i, j int) bool {
|
||||
return results[i].Score > results[j].Score
|
||||
})
|
||||
|
||||
// Hybrid scoring kicks in only when both the embedder and the
|
||||
// vector store are wired and BM25 actually returned candidates.
|
||||
if opts.Vector != nil && opts.Embedder != nil && len(results) > 0 {
|
||||
merged, err := hybridMerge(ctx, brainDir, opts, results)
|
||||
if err != nil {
|
||||
slog.Warn("search: hybrid merge failed, falling back to BM25", "err", err)
|
||||
} else {
|
||||
results = merged
|
||||
}
|
||||
}
|
||||
|
||||
if len(results) > opts.Limit {
|
||||
results = results[:opts.Limit]
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// rrfK is the constant in the Reciprocal Rank Fusion formula. 60 is
|
||||
// standard (Cormack et al. 2009) and parameter-free in practice.
|
||||
const rrfK = 60.0
|
||||
|
||||
// hybridMerge embeds the query, runs a vector search, and merges its
|
||||
// candidates with the BM25 list via Reciprocal Rank Fusion. Results
|
||||
// that came only from the vector side are hydrated by reading the
|
||||
// note's frontmatter for title/wing/hall and excerpting the body.
|
||||
//
|
||||
// rrf(d) = sum_r 1 / (k + rank_r(d)) over rankers r ∈ {BM25, vector}.
|
||||
func hybridMerge(ctx context.Context, brainDir string, opts QueryOptions, bm25 []Result) ([]Result, error) {
|
||||
q, err := opts.Embedder.Embed(ctx, opts.Query)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("embed query: %w", err)
|
||||
}
|
||||
vectorLimit := opts.Limit * 4
|
||||
if vectorLimit < 20 {
|
||||
vectorLimit = 20
|
||||
}
|
||||
hits, err := opts.Vector.Search(ctx, q, vectorLimit)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("vector search: %w", err)
|
||||
}
|
||||
|
||||
rrf := make(map[string]float64)
|
||||
byPath := make(map[string]Result)
|
||||
for rank, r := range bm25 {
|
||||
rrf[r.Path] += 1.0 / (rrfK + float64(rank+1))
|
||||
byPath[r.Path] = r
|
||||
}
|
||||
for rank, h := range hits {
|
||||
if opts.Wing != "" && !pathInScope(h.Path, opts.Wing, opts.Hall) {
|
||||
continue
|
||||
}
|
||||
rrf[h.Path] += 1.0 / (rrfK + float64(rank+1))
|
||||
if _, seen := byPath[h.Path]; !seen {
|
||||
r, err := hydrate(brainDir, h.Path)
|
||||
if err != nil {
|
||||
slog.Warn("search: hydrate failed for vector hit", "path", h.Path, "err", err)
|
||||
continue
|
||||
}
|
||||
byPath[h.Path] = r
|
||||
}
|
||||
}
|
||||
|
||||
merged := make([]Result, 0, len(byPath))
|
||||
for p, r := range byPath {
|
||||
r.Score = int(rrf[p] * 1e6) // scale to int for stable JSON; relative order is what matters
|
||||
merged = append(merged, r)
|
||||
}
|
||||
sort.Slice(merged, func(i, j int) bool {
|
||||
return merged[i].Score > merged[j].Score
|
||||
})
|
||||
return merged, nil
|
||||
}
|
||||
|
||||
// pathInScope reports whether a wiki path satisfies the wing/hall filter.
|
||||
func pathInScope(relPath, wing, hall string) bool {
|
||||
prefix := "wiki/" + brain.Sanitise(wing) + "/"
|
||||
if hall != "" {
|
||||
prefix += hall + "/"
|
||||
}
|
||||
return strings.HasPrefix(relPath, prefix)
|
||||
}
|
||||
|
||||
// hydrate reads a single note from disk and returns a Result with title,
|
||||
// excerpt, wing, and hall populated. Used for paths that surface only
|
||||
// via vector search.
|
||||
func hydrate(brainDir, relPath string) (Result, error) {
|
||||
full := filepath.Join(brainDir, filepath.FromSlash(relPath))
|
||||
content, err := os.ReadFile(full)
|
||||
if err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
wing, hall := extractWingHall(string(content), relPath)
|
||||
return Result{
|
||||
Path: relPath,
|
||||
Title: extractTitle(string(content), filepath.Base(relPath)),
|
||||
Excerpt: excerpt(string(content), 300),
|
||||
Wing: wing,
|
||||
Hall: hall,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// resolveRoots returns the directories to walk for the given wing/hall
|
||||
// filters. Validates hall against the closed vocabulary when set.
|
||||
func resolveRoots(brainDir, wing, hall string) ([]string, error) {
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
package search_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
@@ -12,6 +13,69 @@ import (
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type stubEmbedder struct{ vec []float32 }
|
||||
|
||||
func (s stubEmbedder) Embed(_ context.Context, _ string) ([]float32, error) { return s.vec, nil }
|
||||
|
||||
type stubVector struct{ hits []search.VectorHit }
|
||||
|
||||
func (s stubVector) Search(_ context.Context, _ []float32, _ int) ([]search.VectorHit, error) {
|
||||
return s.hits, nil
|
||||
}
|
||||
|
||||
func TestSearch_HybridRRFPromotesVectorOnlyHit(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
for _, p := range []struct{ rel, body string }{
|
||||
// BM25-keyword note (matches "lejpa" once)
|
||||
{"wiki/jepa-fx/facts/foo.md", "---\ntitle: Foo\n---\nlejpa keyword\n"},
|
||||
// Semantically related note that does NOT contain the keyword.
|
||||
{"wiki/jepa-fx/facts/semantic.md", "---\ntitle: Semantic\n---\nNo keyword in body.\n"},
|
||||
} {
|
||||
full := filepath.Join(dir, p.rel)
|
||||
require.NoError(t, os.MkdirAll(filepath.Dir(full), 0o755))
|
||||
require.NoError(t, os.WriteFile(full, []byte(p.body), 0o644))
|
||||
}
|
||||
|
||||
embedder := stubEmbedder{vec: []float32{0.1}}
|
||||
vector := stubVector{hits: []search.VectorHit{
|
||||
{Path: "wiki/jepa-fx/facts/semantic.md", Distance: 0.05}, // best vector match
|
||||
{Path: "wiki/jepa-fx/facts/foo.md", Distance: 0.10},
|
||||
}}
|
||||
|
||||
got, err := search.Query(dir, search.QueryOptions{
|
||||
Query: "lejpa",
|
||||
Limit: 5,
|
||||
Vector: vector,
|
||||
Embedder: embedder,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, got, 2, "vector-only hit should be hydrated into results")
|
||||
paths := []string{got[0].Path, got[1].Path}
|
||||
assert.Contains(t, paths, "wiki/jepa-fx/facts/foo.md")
|
||||
assert.Contains(t, paths, "wiki/jepa-fx/facts/semantic.md")
|
||||
}
|
||||
|
||||
func TestSearch_HybridFallsBackOnEmbedderError(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
require.NoError(t, os.MkdirAll(filepath.Join(dir, "wiki"), 0o755))
|
||||
require.NoError(t, os.WriteFile(filepath.Join(dir, "wiki", "x.md"), []byte("keyword foo"), 0o644))
|
||||
|
||||
embedder := errorEmbedder{}
|
||||
vector := stubVector{}
|
||||
got, err := search.Query(dir, search.QueryOptions{
|
||||
Query: "keyword", Limit: 5, Vector: vector, Embedder: embedder,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, got, 1, "BM25 result should still come back when embedder fails")
|
||||
assert.Equal(t, "wiki/x.md", got[0].Path)
|
||||
}
|
||||
|
||||
type errorEmbedder struct{}
|
||||
|
||||
func (errorEmbedder) Embed(_ context.Context, _ string) ([]float32, error) {
|
||||
return nil, assert.AnError
|
||||
}
|
||||
|
||||
func TestSearch_ReturnsMatchingPages(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
require.NoError(t, os.MkdirAll(filepath.Join(dir, "knowledge"), 0o755))
|
||||
|
||||
Reference in New Issue
Block a user