hybrid semantic search: embeddings + KNN + RRF fusion (#1)

cmd/solr-mem-backfill/main.go

@@ -0,0 +1,88 @@
+// Command solr-mem-backfill embeds existing memories so semantic search can
+// find them. It re-embeds every memory (idempotent) — title + content via the
+// configured embedder — and atomically sets the `embedding` field.
+//
+// Env: SOLR_URL (default http://localhost:8983/solr/memories), EMBED_URL,
+// EMBED_MODEL, EMBED_DIM (same as the server).
+package main
+
+import (
+	"context"
+	"flag"
+	"log"
+	"os"
+	"strings"
+	"time"
+
+	"github.com/arreyder/solr-mem/internal/embed"
+	"github.com/arreyder/solr-mem/internal/solr"
+)
+
+func main() {
+	batch := flag.Int("batch", 100, "docs per page / update batch")
+	flag.Parse()
+
+	solrURL := os.Getenv("SOLR_URL")
+	if solrURL == "" {
+		solrURL = "http://localhost:8983/solr/memories"
+	}
+	client := solr.NewClient(solrURL)
+
+	emb := embed.FromEnv()
+	if !emb.Enabled() {
+		log.Fatal("EMBED_URL not set — nothing to backfill (embeddings disabled)")
+	}
+
+	ctx := context.Background()
+	start := 0
+	total, embedded, failed := 0, 0, 0
+
+	for {
+		resp, err := client.Query(ctx, solr.QueryParams{
+			Query:     "*:*",
+			Rows:      *batch,
+			Start:     start,
+			Sort:      "id asc", // stable paging
+			Fields:    []string{"id", "title", "content"},
+			Highlight: false,
+		})
+		if err != nil {
+			log.Fatalf("query at start=%d: %v", start, err)
+		}
+		if len(resp.Docs) == 0 {
+			break
+		}
+
+		var updates []map[string]any
+		for _, d := range resp.Docs {
+			total++
+			id, _ := d["id"].(string)
+			title, _ := d["title"].(string)
+			content, _ := d["content"].(string)
+			text := strings.TrimSpace(title + "\n\n" + content)
+			if id == "" || text == "" {
+				continue
+			}
+			vec, err := emb.EmbedDocument(ctx, text)
+			if err != nil || len(vec) == 0 {
+				log.Printf("embed failed id=%s: %v", id, err)
+				failed++
+				continue
+			}
+			updates = append(updates, map[string]any{
+				"id":        id,
+				"embedding": map[string]any{"set": vec},
+			})
+			embedded++
+		}
+
+		if err := client.BulkUpdate(ctx, updates); err != nil {
+			log.Fatalf("bulk update at start=%d: %v", start, err)
+		}
+		log.Printf("progress: %d seen, %d embedded, %d failed", total, embedded, failed)
+		start += *batch
+		time.Sleep(50 * time.Millisecond) // be gentle on the embed service
+	}
+
+	log.Printf("DONE: %d memories, %d embedded, %d failed", total, embedded, failed)
+}

cmd/solr-mem-server/bulk_store_tool.go

@@ -78,6 +78,7 @@ func bulkStoreMemoriesTool(ctx context.Context, args map[string]any) (any, error
 			SessionID:  getString(m, "session_id"),
 			RelatedIDs: getStringSlice(m, "related_ids"),
 			Format:     format,
+			Embedding:  embedMemoryText(ctx, scrubbedTitle, scrubbedContent),
 		})
 	}
 

cmd/solr-mem-server/embedding.go

@@ -0,0 +1,47 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"strings"
+
+	"github.com/arreyder/solr-mem/internal/solr"
+)
+
+// embedMemoryText embeds a memory's semantic text (title + content). Returns
+// nil when embeddings are disabled or on error, so store/update degrade to
+// lexical-only instead of failing the write.
+func embedMemoryText(ctx context.Context, title, content string) []float32 {
+	if !embedder.Enabled() {
+		return nil
+	}
+	text := strings.TrimSpace(title + "\n\n" + content)
+	if text == "" {
+		return nil
+	}
+	vec, err := embedder.EmbedDocument(ctx, text)
+	if err != nil {
+		log.Printf("embedding failed (proceeding without vector): %v", err)
+		return nil
+	}
+	return vec
+}
+
+// currentTitleContent fetches a memory's stored title and content. Used when
+// re-embedding on update where only one of the two was supplied, so the vector
+// still reflects both fields.
+func currentTitleContent(ctx context.Context, id string) (title, content string) {
+	resp, err := solrClient.Query(ctx, solr.QueryParams{
+		Query:     fmt.Sprintf("id:%q", id),
+		Rows:      1,
+		Fields:    []string{"title", "content"},
+		Highlight: false,
+	})
+	if err != nil || resp == nil || len(resp.Docs) == 0 {
+		return "", ""
+	}
+	title, _ = resp.Docs[0]["title"].(string)
+	content, _ = resp.Docs[0]["content"].(string)
+	return title, content
+}

cmd/solr-mem-server/fuse.go

@@ -0,0 +1,84 @@
+package main
+
+import (
+	"sort"
+
+	"github.com/arreyder/solr-mem/internal/solr"
+)
+
+// rrfK is the reciprocal-rank-fusion constant. 60 is the widely-used default
+// from the original RRF paper; it damps the influence of any single ranker's
+// top positions so the two lists combine smoothly.
+const rrfK = 60
+
+func docID(d map[string]any) string {
+	s, _ := d["id"].(string)
+	return s
+}
+
+// fuseResponses combines lexical and semantic (KNN) result lists with
+// reciprocal rank fusion: score(d) = Σ 1/(rrfK + rank) across the lists it
+// appears in. Returns a response with docs ordered by fused score (desc),
+// de-duplicated by id and capped to limit, with highlighting merged from both.
+// Ties break by lexical order first, then semantic — deterministic regardless
+// of map iteration.
+func fuseResponses(lexical, semantic *solr.QueryResponse, limit int) *solr.QueryResponse {
+	scores := map[string]float64{}
+	docByID := map[string]map[string]any{}
+	var order []string // deterministic seed order: lexical first, then semantic-only
+	seen := map[string]bool{}
+
+	accumulate := func(resp *solr.QueryResponse) {
+		if resp == nil {
+			return
+		}
+		for rank, d := range resp.Docs {
+			id := docID(d)
+			if id == "" {
+				continue
+			}
+			scores[id] += 1.0 / float64(rrfK+rank+1) // rank is 0-based
+			if !seen[id] {
+				seen[id] = true
+				order = append(order, id)
+				docByID[id] = d
+			}
+		}
+	}
+	accumulate(lexical)
+	accumulate(semantic)
+
+	sort.SliceStable(order, func(i, j int) bool {
+		return scores[order[i]] > scores[order[j]]
+	})
+	if limit > 0 && len(order) > limit {
+		order = order[:limit]
+	}
+
+	docs := make([]map[string]any, 0, len(order))
+	for _, id := range order {
+		docs = append(docs, docByID[id])
+	}
+
+	hl := map[string]map[string][]string{}
+	for _, resp := range []*solr.QueryResponse{lexical, semantic} {
+		if resp == nil {
+			continue
+		}
+		for k, v := range resp.Highlighting {
+			hl[k] = v
+		}
+	}
+
+	out := &solr.QueryResponse{Docs: docs, Highlighting: hl}
+	if lexical != nil {
+		out.Facets = lexical.Facets
+		out.NumFound = lexical.NumFound
+	}
+	// Don't report fewer than we're actually returning — semantic-only hits
+	// aren't counted in the lexical NumFound.
+	if len(docs) > out.NumFound {
+		out.NumFound = len(docs)
+	}
+	return out
+}

cmd/solr-mem-server/fuse_test.go

@@ -0,0 +1,61 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/arreyder/solr-mem/internal/solr"
+)
+
+func resp(ids ...string) *solr.QueryResponse {
+	docs := make([]map[string]any, len(ids))
+	for i, id := range ids {
+		docs[i] = map[string]any{"id": id}
+	}
+	return &solr.QueryResponse{NumFound: len(ids), Docs: docs}
+}
+
+func order(r *solr.QueryResponse) []string {
+	out := make([]string, len(r.Docs))
+	for i, d := range r.Docs {
+		out[i] = docID(d)
+	}
+	return out
+}
+
+func TestFuseResponses_RRF(t *testing.T) {
+	// b ranks high in both lists -> should win. d is semantic-only -> still included.
+	lexical := resp("a", "b", "c")
+	semantic := resp("b", "d", "a")
+
+	fused := fuseResponses(lexical, semantic, 10)
+	got := order(fused)
+
+	// b: 1/61 + 1/61 (rank0 both) = highest.
+	if got[0] != "b" {
+		t.Fatalf("expected 'b' first, got %v", got)
+	}
+	// All four unique ids present, deduped.
+	if len(got) != 4 {
+		t.Fatalf("expected 4 unique docs, got %v", got)
+	}
+	// d (semantic-only) is included.
+	if !contains(got, "d") {
+		t.Errorf("semantic-only 'd' missing: %v", got)
+	}
+}
+
+func TestFuseResponses_LimitAndNilSemantic(t *testing.T) {
+	fused := fuseResponses(resp("a", "b", "c"), nil, 2)
+	if got := order(fused); len(got) != 2 || got[0] != "a" {
+		t.Fatalf("limit/nil-semantic: got %v", got)
+	}
+}
+
+func contains(s []string, v string) bool {
+	for _, x := range s {
+		if x == v {
+			return true
+		}
+	}
+	return false
+}

cmd/solr-mem-server/main.go

@@ -8,13 +8,18 @@ import (
 	"net/http"
 	"os"
 
+	"github.com/arreyder/solr-mem/internal/embed"
 	"github.com/arreyder/solr-mem/internal/solr"
 	"github.com/modelcontextprotocol/go-sdk/mcp"
 )
 
 var solrClient *solr.Client
 var codeClient *solr.Client
 
+// embedder produces query/document embeddings for semantic search. Disabled
+// (lexical-only) unless EMBED_URL is configured.
+var embedder embed.Embedder = embed.Disabled{}
+
 // indexerControlURL is the base URL of the indexer's force-reindex control
 // endpoint. Defaults to the co-located indexer on localhost.
 var indexerControlURL = envOrDefault("INDEXER_CONTROL_URL", "http://127.0.0.1:7071")
@@ -39,6 +44,13 @@ func main() {
 	}
 	codeClient = solr.NewClient(codeURL)
 
+	embedder = embed.FromEnv()
+	if embedder.Enabled() {
+		log.Printf("Semantic search enabled: embeddings via %s (dim %d)", os.Getenv("EMBED_URL"), embedder.Dim())
+	} else {
+		log.Printf("Semantic search disabled (EMBED_URL unset); lexical-only")
+	}
+
 	// Start expiration sweeper
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()

cmd/solr-mem-server/search_tool.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"log"
 	"strings"
 
 	"github.com/arreyder/solr-mem/internal/solr"
@@ -15,9 +16,22 @@ func searchMemoriesTool(ctx context.Context, args map[string]any) (any, error) {
 		return nil, fmt.Errorf("query is required")
 	}
 
+	limit := getInt(args, "limit", 10)
+
+	// Hybrid semantic search: blend lexical ranking with vector KNN. On when an
+	// embedder is configured and the caller hasn't opted out. We over-fetch
+	// (fusionK) from each ranker so fusion has enough candidates, then trim to
+	// limit at the end. Pagination (start) forces lexical-only — KNN+RRF has no
+	// stable global offset.
+	semantic := embedder.Enabled() && getBool(args, "semantic", true) && getInt(args, "start", 0) == 0
+	fusionK := limit
+	if semantic && fusionK < 50 {
+		fusionK = 50
+	}
+
 	params := solr.QueryParams{
 		Query:     query,
-		Rows:      getInt(args, "limit", 10),
+		Rows:      fusionK,
 		Start:     getInt(args, "start", 0),
 		Highlight: getBool(args, "highlight", true),
 		Facet:     getBool(args, "facet", false),
@@ -83,6 +97,22 @@ func searchMemoriesTool(ctx context.Context, args map[string]any) (any, error) {
 		return nil, fmt.Errorf("search failed: %w", err)
 	}
 
+	// Blend in semantic (KNN) ranking via reciprocal rank fusion. Degrades to
+	// lexical-only on any embed/KNN error so search never hard-fails on the
+	// optional path.
+	if semantic {
+		if vec, eerr := embedder.EmbedQuery(ctx, query); eerr != nil {
+			log.Printf("query embedding failed (lexical-only): %v", eerr)
+		} else if len(vec) > 0 {
+			knn, kerr := solrClient.KNNQuery(ctx, "embedding", vec, fusionK, params.FilterQueries, params.Fields)
+			if kerr != nil {
+				log.Printf("knn search failed (lexical-only): %v", kerr)
+			} else {
+				resp = fuseResponses(resp, knn, fusionK)
+			}
+		}
+	}
+
 	// Cap per session so one chatty session can't dominate results.
 	// Default 3; pass 0 to disable.
 	sessionCap := getInt(args, "session_cap", 3)
@@ -93,6 +123,11 @@ func searchMemoriesTool(ctx context.Context, args map[string]any) (any, error) {
 		}, sessionCap)
 	}
 
+	// Trim to the requested limit (we over-fetched fusionK for fusion headroom).
+	if limit > 0 && len(resp.Docs) > limit {
+		resp.Docs = resp.Docs[:limit]
+	}
+
 	// Credit a retrieval for the memories actually surfaced (fire-and-forget).
 	// Pass track:false for maintenance/bulk scans so they don't inflate the signal.
 	recordRetrievalsAsync(resp.Docs, getBool(args, "track", true))

cmd/solr-mem-server/store_tool.go

@@ -50,6 +50,7 @@ func storeMemoryTool(ctx context.Context, args map[string]any) (any, error) {
 		RelatedIDs: getStringSlice(args, "related_ids"),
 		Format:     format,
 	}
+	doc.Embedding = embedMemoryText(ctx, scrubbedTitle, scrubbedContent)
 
 	if err := solrClient.Add(ctx, doc); err != nil {
 		return nil, fmt.Errorf("failed to store memory: %w", err)