diff --git a/cmd/solr-mem-backfill/main.go b/cmd/solr-mem-backfill/main.go new file mode 100644 index 0000000..e376a2d --- /dev/null +++ b/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) +} diff --git a/cmd/solr-mem-server/bulk_store_tool.go b/cmd/solr-mem-server/bulk_store_tool.go index 6f64a75..bfadadd 100644 --- a/cmd/solr-mem-server/bulk_store_tool.go +++ b/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), }) } diff --git a/cmd/solr-mem-server/embedding.go b/cmd/solr-mem-server/embedding.go new file mode 100644 index 0000000..43f6ed8 --- /dev/null +++ b/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 +} diff --git a/cmd/solr-mem-server/fuse.go b/cmd/solr-mem-server/fuse.go new file mode 100644 index 0000000..04c9793 --- /dev/null +++ b/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 +} diff --git a/cmd/solr-mem-server/fuse_test.go b/cmd/solr-mem-server/fuse_test.go new file mode 100644 index 0000000..e6cf9ca --- /dev/null +++ b/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 +} diff --git a/cmd/solr-mem-server/main.go b/cmd/solr-mem-server/main.go index 3d01776..5f2b2c7 100644 --- a/cmd/solr-mem-server/main.go +++ b/cmd/solr-mem-server/main.go @@ -8,6 +8,7 @@ import ( "net/http" "os" + "github.com/arreyder/solr-mem/internal/embed" "github.com/arreyder/solr-mem/internal/solr" "github.com/modelcontextprotocol/go-sdk/mcp" ) @@ -15,6 +16,10 @@ import ( 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() diff --git a/cmd/solr-mem-server/search_tool.go b/cmd/solr-mem-server/search_tool.go index 8aa70aa..31ef0b2 100644 --- a/cmd/solr-mem-server/search_tool.go +++ b/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)) diff --git a/cmd/solr-mem-server/store_tool.go b/cmd/solr-mem-server/store_tool.go index c17459c..2054ea3 100644 --- a/cmd/solr-mem-server/store_tool.go +++ b/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) diff --git a/cmd/solr-mem-server/tools.go b/cmd/solr-mem-server/tools.go index 0a47e8b..c1857e6 100644 --- a/cmd/solr-mem-server/tools.go +++ b/cmd/solr-mem-server/tools.go @@ -56,6 +56,8 @@ func ToolSchemas() []ToolDefinition { **When to use**: Find relevant memories by content, tags, type, agent, or time range. Uses edismax with field boosting (content^3, title^2, tags^1.5) and recency boost. +**Semantic (hybrid) search**: when embeddings are enabled, results blend lexical (BM25) ranking with vector similarity (KNN) via reciprocal rank fusion — so conceptually-related memories surface even with no shared words. On by default; pass semantic=false for pure lexical (exact/debugging). Ignored when start>0 (pagination is lexical-only). + **Match modes**: By default a query requires most terms to match (mm 75%), so a long OR-style list of synonyms can return nothing. Pass match="any" for OR-style recall (any term matches) — best when throwing several candidate terms at the store; match="all" requires every term. **Lean payloads**: Pass fields (e.g. ["id","title","importance","memory_type"]) to project only those fields and avoid returning full content bodies. Use start for pagination (offset). @@ -66,6 +68,7 @@ func ToolSchemas() []ToolDefinition { **Optional**: match, fields, start, agent_id, memory_type, tags, source, importance_min, from, to, limit, highlight, facet, session_id, lifetime, session_cap, track`, InputSchema: NewObjectSchema(map[string]any{ "query": prop("string", "Full-text search query (required)"), + "semantic": prop("boolean", "Hybrid semantic+lexical search (default true when embeddings enabled). Set false for pure lexical/exact matching."), "match": prop("string", "Match mode: 'most' (default, mm 75%), 'any' (OR — any term), 'all' (every term). Use 'any' for synonym/OR-style recall."), "track": prop("boolean", "Record a retrieval for surfaced memories (default true). Set false for maintenance/bulk scans so they don't inflate usage stats."), "fields": arrayPropSchema(prop("string", "Field name"), "Project only these fields (Solr fl) for lean payloads; id is always included"), diff --git a/cmd/solr-mem-server/update_tool.go b/cmd/solr-mem-server/update_tool.go index 26d25af..929fc34 100644 --- a/cmd/solr-mem-server/update_tool.go +++ b/cmd/solr-mem-server/update_tool.go @@ -75,6 +75,26 @@ func updateMemoryTool(ctx context.Context, args map[string]any) (any, error) { return nil, fmt.Errorf("no fields to update") } + // Re-embed only when the semantic text (content/title) changed. Tag / + // importance / related_ids updates (e.g. the sleep-pass) skip this. If only + // one of the two changed, fetch the other so the vector reflects both. + newContent, contentChanged := fields["content"].(string) + newTitle, titleChanged := fields["title"].(string) + if embedder.Enabled() && (contentChanged || titleChanged) { + if !contentChanged || !titleChanged { + curTitle, curContent := currentTitleContent(ctx, id) + if !titleChanged { + newTitle = curTitle + } + if !contentChanged { + newContent = curContent + } + } + if vec := embedMemoryText(ctx, newTitle, newContent); vec != nil { + fields["embedding"] = vec + } + } + fields["updated_at"] = time.Now().UTC().Format(time.RFC3339) if err := solrClient.Update(ctx, id, fields); err != nil { diff --git a/internal/embed/embed.go b/internal/embed/embed.go new file mode 100644 index 0000000..eb17875 --- /dev/null +++ b/internal/embed/embed.go @@ -0,0 +1,182 @@ +// Package embed provides text embeddings for semantic memory search. +// +// The default backend is an Ollama-compatible HTTP endpoint. When no endpoint +// is configured the package returns a disabled embedder, so the rest of the +// system degrades gracefully to lexical-only search instead of failing. +package embed + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + "strconv" + "strings" + "time" +) + +// Embedder turns text into a dense vector. Implementations must be safe for +// concurrent use. +// +// Query and document embedding are distinct because some models (e.g. +// nomic-embed-text) are trained for asymmetric retrieval and expect different +// task prefixes on the stored text vs. the search query. Always embed stored +// memories with EmbedDocument and search queries with EmbedQuery so the two +// land in the same space. +type Embedder interface { + // Enabled reports whether embeddings are available. When false, callers + // should fall back to lexical-only behavior. + Enabled() bool + // Dim is the vector dimension (must match the Solr DenseVectorField). + Dim() int + // EmbedDocument embeds text to be stored/indexed. + EmbedDocument(ctx context.Context, text string) ([]float32, error) + // EmbedQuery embeds a search query. + EmbedQuery(ctx context.Context, text string) ([]float32, error) +} + +// FromEnv builds an Embedder from EMBED_URL / EMBED_MODEL / EMBED_DIM. +// If EMBED_URL is empty, returns a disabled embedder (semantic search off). +func FromEnv() Embedder { + url := os.Getenv("EMBED_URL") + if url == "" { + return Disabled{} + } + model := os.Getenv("EMBED_MODEL") + if model == "" { + model = "nomic-embed-text" + } + dim := 768 + if v := os.Getenv("EMBED_DIM"); v != "" { + if n, err := strconv.Atoi(v); err == nil && n > 0 { + dim = n + } + } + o := NewOllama(url, model, dim) + if v := os.Getenv("EMBED_MAX_CHARS"); v != "" { + if n, err := strconv.Atoi(v); err == nil && n > 0 { + o.maxChars = n + } + } + // Task prefixes: default to nomic's for nomic models, off otherwise. + // EMBED_*_PREFIX env overrides either way (set to " " to force-clear). + if strings.Contains(strings.ToLower(model), "nomic") { + o.docPrefix = "search_document: " + o.queryPrefix = "search_query: " + } + if v, ok := os.LookupEnv("EMBED_DOC_PREFIX"); ok { + o.docPrefix = v + } + if v, ok := os.LookupEnv("EMBED_QUERY_PREFIX"); ok { + o.queryPrefix = v + } + return o +} + +// Disabled is a no-op embedder used when no backend is configured. +type Disabled struct{} + +func (Disabled) Enabled() bool { return false } +func (Disabled) Dim() int { return 0 } +func (Disabled) EmbedDocument(context.Context, string) ([]float32, error) { + return nil, nil +} +func (Disabled) EmbedQuery(context.Context, string) ([]float32, error) { + return nil, nil +} + +// Ollama embeds via an Ollama-compatible /api/embeddings endpoint. +type Ollama struct { + baseURL string + model string + dim int + maxChars int // truncate input to this many runes (model context guard) + docPrefix string // task prefix for stored documents + queryPrefix string // task prefix for search queries + client *http.Client +} + +// defaultMaxChars keeps embed input under typical small-model context windows +// (nomic-embed-text ~2048 tokens). Conservative at ~4 chars/token. +const defaultMaxChars = 6000 + +// NewOllama builds an Ollama embedder. baseURL is the host root, e.g. +// "http://pax99.local:11434". +func NewOllama(baseURL, model string, dim int) *Ollama { + return &Ollama{ + baseURL: baseURL, + model: model, + dim: dim, + maxChars: defaultMaxChars, + client: &http.Client{Timeout: 30 * time.Second}, + } +} + +func (o *Ollama) Enabled() bool { return true } +func (o *Ollama) Dim() int { return o.dim } + +func (o *Ollama) EmbedDocument(ctx context.Context, text string) ([]float32, error) { + return o.embed(ctx, o.docPrefix+text) +} + +func (o *Ollama) EmbedQuery(ctx context.Context, text string) ([]float32, error) { + return o.embed(ctx, o.queryPrefix+text) +} + +func (o *Ollama) embed(ctx context.Context, text string) ([]float32, error) { + text = truncateRunes(text, o.maxChars) + body, err := json.Marshal(map[string]any{"model": o.model, "prompt": text}) + if err != nil { + return nil, err + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, + o.baseURL+"/api/embeddings", bytes.NewReader(body)) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + + resp, err := o.client.Do(req) + if err != nil { + return nil, fmt.Errorf("embed request: %w", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("embed returned %d: %s", resp.StatusCode, b) + } + return parseEmbedding(resp.Body, o.dim) +} + +// truncateRunes caps s to at most max runes (UTF-8 safe). max<=0 means no cap. +func truncateRunes(s string, max int) string { + if max <= 0 { + return s + } + r := []rune(s) + if len(r) <= max { + return s + } + return string(r[:max]) +} + +// parseEmbedding decodes an Ollama embeddings response and validates the +// dimension. Split out for testing. +func parseEmbedding(r io.Reader, wantDim int) ([]float32, error) { + var out struct { + Embedding []float32 `json:"embedding"` + } + if err := json.NewDecoder(r).Decode(&out); err != nil { + return nil, fmt.Errorf("decode embedding: %w", err) + } + if len(out.Embedding) == 0 { + return nil, fmt.Errorf("embedding response had no vector") + } + if wantDim > 0 && len(out.Embedding) != wantDim { + return nil, fmt.Errorf("embedding dim %d != expected %d", len(out.Embedding), wantDim) + } + return out.Embedding, nil +} diff --git a/internal/embed/embed_test.go b/internal/embed/embed_test.go new file mode 100644 index 0000000..679154d --- /dev/null +++ b/internal/embed/embed_test.go @@ -0,0 +1,83 @@ +package embed + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestDisabled(t *testing.T) { + var e Embedder = Disabled{} + if e.Enabled() { + t.Fatal("Disabled must report not enabled") + } + if v, err := e.EmbedDocument(context.Background(), "x"); err != nil || v != nil { + t.Fatalf("Disabled.EmbedDocument = %v, %v; want nil, nil", v, err) + } + if v, err := e.EmbedQuery(context.Background(), "x"); err != nil || v != nil { + t.Fatalf("Disabled.EmbedQuery = %v, %v; want nil, nil", v, err) + } +} + +func TestOllamaEmbedAndPrefixes(t *testing.T) { + var gotPrompt string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/embeddings" { + t.Errorf("path = %q", r.URL.Path) + } + var req map[string]any + b, _ := io.ReadAll(r.Body) + json.Unmarshal(b, &req) + gotPrompt, _ = req["prompt"].(string) + io.WriteString(w, `{"embedding":[0.1,0.2,0.3]}`) + })) + defer srv.Close() + + o := NewOllama(srv.URL, "nomic-embed-text", 3) + o.docPrefix = "search_document: " + o.queryPrefix = "search_query: " + + v, err := o.EmbedDocument(context.Background(), "hello") + if err != nil || len(v) != 3 || v[0] != 0.1 { + t.Fatalf("EmbedDocument = %v, %v", v, err) + } + if gotPrompt != "search_document: hello" { + t.Errorf("doc prompt = %q, want prefixed", gotPrompt) + } + + if _, err := o.EmbedQuery(context.Background(), "hello"); err != nil { + t.Fatalf("EmbedQuery err: %v", err) + } + if gotPrompt != "search_query: hello" { + t.Errorf("query prompt = %q, want prefixed", gotPrompt) + } +} + +func TestTruncateRunes(t *testing.T) { + if got := truncateRunes("hello", 3); got != "hel" { + t.Errorf("got %q", got) + } + if got := truncateRunes("hi", 10); got != "hi" { + t.Errorf("no-trunc got %q", got) + } + if got := truncateRunes("hello", 0); got != "hello" { + t.Errorf("max<=0 must be no-op, got %q", got) + } + // multi-byte safe + if got := truncateRunes("héllo", 2); got != "hé" { + t.Errorf("utf8 got %q", got) + } +} + +func TestParseEmbeddingDimMismatch(t *testing.T) { + if _, err := parseEmbedding(strings.NewReader(`{"embedding":[1,2]}`), 3); err == nil { + t.Fatal("expected dim-mismatch error") + } + if _, err := parseEmbedding(strings.NewReader(`{"embedding":[]}`), 0); err == nil { + t.Fatal("expected empty-vector error") + } +} diff --git a/internal/solr/client.go b/internal/solr/client.go index 893a282..a98aab5 100644 --- a/internal/solr/client.go +++ b/internal/solr/client.go @@ -9,6 +9,7 @@ import ( "net/http" "net/url" "path" + "strconv" "strings" "time" ) @@ -311,6 +312,61 @@ func (c *Client) MoreLikeThis(ctx context.Context, id string, rows int, filterQu return ParseQueryResponse(resp.Body) } +// KNNQuery runs an approximate-nearest-neighbor search over a dense vector +// field. The query is POSTed (not GET) because the vector text can exceed URL +// length limits. Returns docs in similarity order. filterQueries are applied as +// pre-filters; fields projects the returned docs. +func (c *Client) KNNQuery(ctx context.Context, field string, vec []float32, topK int, filterQueries, fields []string) (*QueryResponse, error) { + form := url.Values{} + form.Set("q", fmt.Sprintf("{!knn f=%s topK=%d}%s", field, topK, formatVector(vec))) + // Force the lucene parser: the /select handler defaults to edismax, which + // does NOT honor the leading {!knn} parser switch and would tokenize the + // vector literal as text (blowing maxClauseCount). Also drop the handler's + // facet/highlight/recency-boost defaults — irrelevant to a KNN sub-query. + form.Set("defType", "lucene") + form.Set("facet", "false") + form.Set("hl", "off") + form.Set("rows", strconv.Itoa(topK)) + form.Set("wt", "json") + if len(fields) > 0 { + form.Set("fl", strings.Join(fields, ",")) + } + for _, fq := range filterQueries { + form.Add("fq", fq) + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, + c.baseURL+"/select", strings.NewReader(form.Encode())) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("solr knn: %w", err) + } + defer resp.Body.Close() + if err := c.checkResponse(resp); err != nil { + return nil, err + } + return ParseQueryResponse(resp.Body) +} + +// formatVector renders a float vector as Solr's "[a,b,c]" literal. +func formatVector(vec []float32) string { + var b strings.Builder + b.WriteByte('[') + for i, f := range vec { + if i > 0 { + b.WriteByte(',') + } + b.WriteString(strconv.FormatFloat(float64(f), 'g', -1, 32)) + } + b.WriteByte(']') + return b.String() +} + // DeleteByQuery removes all documents matching the given query. func (c *Client) DeleteByQuery(ctx context.Context, query string) error { payload := map[string]any{ diff --git a/internal/solr/types.go b/internal/solr/types.go index d22ff70..e6a003a 100644 --- a/internal/solr/types.go +++ b/internal/solr/types.go @@ -20,6 +20,9 @@ type Document struct { SessionID string `json:"session_id,omitempty"` RelatedIDs []string `json:"related_ids,omitempty"` Format string `json:"format,omitempty"` + // Embedding is the dense semantic vector. Omitted when embeddings are + // disabled so the field never reaches Solr in lexical-only mode. + Embedding []float32 `json:"embedding,omitempty"` } // QueryParams holds parameters for a Solr search query. diff --git a/solr/managed-schema.xml b/solr/managed-schema.xml index 39f9cd4..272c1d9 100644 --- a/solr/managed-schema.xml +++ b/solr/managed-schema.xml @@ -9,6 +9,13 @@ + + + + @@ -54,4 +61,12 @@ + + + +