From 06520230376151b13411ac10ce323786dff8ebce Mon Sep 17 00:00:00 2001 From: Victor Vargas Date: Mon, 6 Jul 2026 00:05:30 -0700 Subject: [PATCH] feat(rag): add SQLite+FTS5 backend, fix content/usage plumbing bugs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Backend.Upsert never received the fragment's Content, so ChromaDB (and any backend) stored the vector but silently dropped the actual text — saved memories had nothing to retrieve later. Backend.Search now also takes the raw query text, and a failed/missing embedding no longer hard-fails Add/Search: it degrades to a nil vector so a lexical-capable backend can still index/find the content (Chroma has no such fallback and now says so explicitly instead of misbehaving). Adds pkg/rag/backends/sqlitevec: a zero-dependency backend (pure-Go SQLite, no external service) that does cosine similarity when a real embedding vector is available and falls back to FTS5/BM25 full-text search otherwise. Adds pkg/rag/embeddings.OpenAICompatible, covering both a local llama.cpp server (`--embeddings` enabled) and real OpenAI (or any OpenAI-shaped /embeddings endpoint) through the same client. Also fixes token usage tracking for llama.cpp streaming: the client never requested `stream_options.include_usage` nor parsed a usage-only SSE event, and even when present, the agent loop's RunStream dropped any chunk with no Delta/ReasoningDelta — silently discarding the only chunk that carries usage. Co-Authored-By: Claude Sonnet 5 --- go.mod | 16 +- go.sum | 49 +++ pkg/agent/loop.go | 8 +- pkg/agent/loop_test.go | 36 +++ pkg/llm/providers/llamacpp/client.go | 33 +++ pkg/llm/providers/llamacpp/client_test.go | 41 +++ pkg/rag/backends/chroma/chroma.go | 13 +- pkg/rag/backends/chroma/chroma_test.go | 30 +- pkg/rag/backends/sqlitevec/sqlitevec.go | 296 +++++++++++++++++++ pkg/rag/backends/sqlitevec/sqlitevec_test.go | 240 +++++++++++++++ pkg/rag/embeddings/openai_compatible.go | 123 ++++++++ pkg/rag/embeddings/openai_compatible_test.go | 134 +++++++++ pkg/rag/memory.go | 23 +- pkg/rag/memory_test.go | 89 ++++-- 14 files changed, 1092 insertions(+), 39 deletions(-) create mode 100644 pkg/rag/backends/sqlitevec/sqlitevec.go create mode 100644 pkg/rag/backends/sqlitevec/sqlitevec_test.go create mode 100644 pkg/rag/embeddings/openai_compatible.go create mode 100644 pkg/rag/embeddings/openai_compatible_test.go diff --git a/go.mod b/go.mod index 091c508..d40084d 100644 --- a/go.mod +++ b/go.mod @@ -4,4 +4,18 @@ go 1.26 require gopkg.in/yaml.v3 v3.0.1 -require github.com/google/uuid v1.6.0 +require ( + github.com/google/uuid v1.6.0 + modernc.org/sqlite v1.53.0 +) + +require ( + github.com/dustin/go-humanize v1.0.1 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect + github.com/ncruces/go-strftime v1.0.0 // indirect + github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect + golang.org/x/sys v0.44.0 // indirect + modernc.org/libc v1.73.4 // indirect + modernc.org/mathutil v1.7.1 // indirect + modernc.org/memory v1.11.0 // indirect +) diff --git a/go.sum b/go.sum index b4c5744..8d8dee1 100644 --- a/go.sum +++ b/go.sum @@ -1,6 +1,55 @@ +github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= +github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= +github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= +github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= +github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= +golang.org/x/mod v0.36.0 h1:JJjpVx6myfUsUdAzZuOSTTmRE0PfZeNWzzvKrP7amb4= +golang.org/x/mod v0.36.0/go.mod h1:moc6ELqsWcOw5Ef3xVprK5ul/MvtVvkIXLziUOICjUQ= +golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= +golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.44.0 h1:ildZl3J4uzeKP07r2F++Op7E9B29JRUy+a27EibtBTQ= +golang.org/x/sys v0.44.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/tools v0.45.0 h1:18qN3FAooORvApf5XjCXgsuayZOEtXf6JK18I3+ONa8= +golang.org/x/tools v0.45.0/go.mod h1:LuUGqqaXcXMEFEruIVJVm5mgDD8vww/z/SR1gQ4uE/0= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +modernc.org/cc/v4 v4.28.4 h1:Hd/4Es+MBj+/7hSdZaisNyu6bv3V0Dp2MdllyfqaH+c= +modernc.org/cc/v4 v4.28.4/go.mod h1:OnovgIhbbMXMu1aISnJ0wvVD1KnW+cAUJkIrAWh+kVI= +modernc.org/ccgo/v4 v4.34.4 h1:OVnSOWQjVKOYkFxoHYB+qQmSHK5gqMqARM+K9DpR/Ws= +modernc.org/ccgo/v4 v4.34.4/go.mod h1:qdKqE8FNIYyysougB1RX9MxCzp5oJOcQXSobANJ4TuE= +modernc.org/fileutil v1.4.0 h1:j6ZzNTftVS054gi281TyLjHPp6CPHr2KCxEXjEbD6SM= +modernc.org/fileutil v1.4.0/go.mod h1:EqdKFDxiByqxLk8ozOxObDSfcVOv/54xDs/DUHdvCUU= +modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI= +modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito= +modernc.org/gc/v3 v3.1.3 h1:6QAplYyVO+KdPW3pGnqmJDUxtkec8ooEWvks/hhU3lc= +modernc.org/gc/v3 v3.1.3/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY= +modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks= +modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI= +modernc.org/libc v1.73.4 h1:+ra4Ui8ngyt8HDcO1FTDPWlkAh6yOdaO2yAoh8MddQA= +modernc.org/libc v1.73.4/go.mod h1:DXZ3eO8qMCNn2SnmTNCiC71nJ9Rcq3PsnpU6Vc4rWK8= +modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU= +modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg= +modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI= +modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw= +modernc.org/opt v0.2.0 h1:tGyef5ApycA7FSEOMraay9SaTk5zmbx7Tu+cJs4QKZg= +modernc.org/opt v0.2.0/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns= +modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w= +modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE= +modernc.org/sqlite v1.53.0 h1:20WG8N9q4ji/dEqGk4uiI0c6OPjSeLTNYGFCc3+7c1M= +modernc.org/sqlite v1.53.0/go.mod h1:xoEpOIpGrgT48H5iiyt/YXPCZPEzlfmfFwtk8Lklw8s= +modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0= +modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A= +modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y= +modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 914becd..64981c3 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -181,7 +181,13 @@ func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Messa } } - if !hasToolCalls && (chunk.Delta != "" || chunk.ReasoningDelta != "") { + // A trailing usage-only chunk (no Delta/ReasoningDelta, per + // providers that report token usage in a separate final + // event) must still be forwarded, or callers can never see + // real token counts. + hasContent := chunk.Delta != "" || chunk.ReasoningDelta != "" + hasUsage := chunk.Usage.TotalTokens > 0 + if !hasToolCalls && (hasContent || hasUsage) { responseBuilder.WriteString(chunk.Delta) if !yield(chunk, nil) { return diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index a4d5c86..29e0de7 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -403,6 +403,42 @@ func TestRun_Stream_NoToolCalls(t *testing.T) { } } +func TestRun_Stream_ForwardsTrailingUsageOnlyChunk(t *testing.T) { + mockClient := &mockLLM{ + streamFunc: func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] { + return func(yield func(llm.StreamChunk, error) bool) { + yield(llm.StreamChunk{Delta: "Hello"}, nil) + // No Delta/ReasoningDelta, as providers report usage in a + // separate trailing event; it must still be forwarded. + yield(llm.StreamChunk{Usage: llm.TokenUsage{InputTokens: 10, OutputTokens: 3, TotalTokens: 13}}, nil) + } + }, + } + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: tools.NewRegistry(), + }) + + var chunks []llm.StreamChunk + stream := loop.RunStream(context.Background(), "test") + for chunk, err := range stream { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + chunks = append(chunks, chunk) + } + + if len(chunks) != 2 { + t.Fatalf("expected 2 chunks (content + usage-only), got %d", len(chunks)) + } + usage := chunks[len(chunks)-1].Usage + if usage.InputTokens != 10 || usage.OutputTokens != 3 || usage.TotalTokens != 13 { + t.Errorf("expected the trailing usage-only chunk to be forwarded, got %+v", usage) + } +} + func TestRun_Stream_WithToolCalls(t *testing.T) { registry := tools.NewRegistry() registry.Register(tools.Tool{ diff --git a/pkg/llm/providers/llamacpp/client.go b/pkg/llm/providers/llamacpp/client.go index 425e9d8..9554ff3 100644 --- a/pkg/llm/providers/llamacpp/client.go +++ b/pkg/llm/providers/llamacpp/client.go @@ -122,10 +122,31 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq return } + var usage llm.TokenUsage + if event.Usage != nil { + usage = llm.TokenUsage{ + InputTokens: event.Usage.PromptTokens, + OutputTokens: event.Usage.CompletionTokens, + TotalTokens: event.Usage.TotalTokens, + } + } + + if len(event.Choices) == 0 { + // The usage-only event (per stream_options.include_usage) + // carries no choices, so it needs its own chunk. + if event.Usage != nil { + if !yield(llm.StreamChunk{Usage: usage}, nil) { + return + } + } + continue + } + for _, choice := range event.Choices { chunk := llm.StreamChunk{ Delta: choice.Delta.Content, ReasoningDelta: choice.Delta.ReasoningContent, + Usage: usage, } if choice.FinishReason != "" { chunk.FinishReason = choice.FinishReason @@ -179,6 +200,12 @@ func (c *Client) buildRequest(req llm.CompletionRequest, stream bool) (io.Reader Stream: stream, ChatTemplateKwargs: req.ChatTemplateKwargs, } + if stream { + // Ask for a final SSE event carrying token usage (OpenAI-style + // streaming omits it otherwise), so Rony can track real token + // counts per turn instead of always seeing zero. + openReq.StreamOptions = &llamaStreamOptions{IncludeUsage: true} + } if len(tools) > 0 { openReq.Tools = tools } @@ -245,9 +272,14 @@ type llamaChatRequest struct { TopP float32 `json:"top_p,omitempty"` Stop []string `json:"stop,omitempty"` Stream bool `json:"stream"` + StreamOptions *llamaStreamOptions `json:"stream_options,omitempty"` ChatTemplateKwargs map[string]any `json:"chat_template_kwargs,omitempty"` } +type llamaStreamOptions struct { + IncludeUsage bool `json:"include_usage"` +} + type llamaMessage struct { Role string `json:"role"` Content string `json:"content"` @@ -300,6 +332,7 @@ type llamaUsage struct { type llamaStreamEvent struct { ID string `json:"id"` Choices []llamaStreamChoice `json:"choices"` + Usage *llamaUsage `json:"usage"` } type llamaStreamChoice struct { diff --git a/pkg/llm/providers/llamacpp/client_test.go b/pkg/llm/providers/llamacpp/client_test.go index 7b98377..0de722a 100644 --- a/pkg/llm/providers/llamacpp/client_test.go +++ b/pkg/llm/providers/llamacpp/client_test.go @@ -3,8 +3,10 @@ package llamacpp import ( "context" "encoding/json" + "io" "net/http" "net/http/httptest" + "strings" "testing" "github.com/VictorVargas/rony-llm-agent/pkg/llm" @@ -230,3 +232,42 @@ func TestClient_Stream_FinishReason(t *testing.T) { t.Errorf("expected 'stop' finish reason, got %q", chunks[0].FinishReason) } } + +func TestClient_Stream_RequestsAndParsesUsage(t *testing.T) { + var gotBody string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + gotBody = string(body) + w.Header().Set("Content-Type", "text/event-stream") + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"hi\"},\"finish_reason\":\"stop\"}]}\n")) + w.Write([]byte("data: {\"choices\":[],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":2,\"total_tokens\":12}}\n")) + w.Write([]byte("data: [DONE]\n")) + })) + defer server.Close() + + client, err := New(Config{BaseURL: server.URL + "/v1"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + var chunks []llm.StreamChunk + stream := client.Stream(context.Background(), llm.CompletionRequest{}) + for chunk, err := range stream { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + chunks = append(chunks, chunk) + } + + if !strings.Contains(gotBody, `"stream_options":{"include_usage":true}`) { + t.Errorf("expected the request to ask for usage via stream_options, got body: %s", gotBody) + } + + if len(chunks) != 2 { + t.Fatalf("expected 2 chunks (content + usage-only), got %d", len(chunks)) + } + usage := chunks[len(chunks)-1].Usage + if usage.InputTokens != 10 || usage.OutputTokens != 2 || usage.TotalTokens != 12 { + t.Errorf("expected usage to be parsed from the final event, got %+v", usage) + } +} diff --git a/pkg/rag/backends/chroma/chroma.go b/pkg/rag/backends/chroma/chroma.go index 6eda10e..3900cb2 100644 --- a/pkg/rag/backends/chroma/chroma.go +++ b/pkg/rag/backends/chroma/chroma.go @@ -44,7 +44,7 @@ func New(cfg Config) (*Backend, error) { }, nil } -func (b *Backend) Upsert(ctx context.Context, id string, vector []float32, metadata map[string]string) error { +func (b *Backend) Upsert(ctx context.Context, id string, vector []float32, content string, metadata map[string]string) error { collection := "rony-memory" embeddings := make([][]float64, 1) @@ -58,9 +58,10 @@ func (b *Backend) Upsert(ctx context.Context, id string, vector []float32, metad } reqBody, err := json.Marshal(map[string]interface{}{ - "ids": []string{id}, + "ids": []string{id}, "embeddings": embeddings, - "metadatas": []map[string]interface{}{metadatas}, + "documents": []string{content}, + "metadatas": []map[string]interface{}{metadatas}, }) if err != nil { return fmt.Errorf("marshaling upsert request: %w", err) @@ -87,7 +88,11 @@ func (b *Backend) Upsert(ctx context.Context, id string, vector []float32, metad return nil } -func (b *Backend) Search(ctx context.Context, queryVector []float32, topK int) ([]rag.SearchResult, error) { +func (b *Backend) Search(ctx context.Context, _ string, queryVector []float32, topK int) ([]rag.SearchResult, error) { + if len(queryVector) == 0 { + return nil, fmt.Errorf("chroma backend requires an embedding vector; it has no lexical fallback") + } + collection := "rony-memory" query := make([]float64, len(queryVector)) diff --git a/pkg/rag/backends/chroma/chroma_test.go b/pkg/rag/backends/chroma/chroma_test.go index fa73ec9..dfe6324 100644 --- a/pkg/rag/backends/chroma/chroma_test.go +++ b/pkg/rag/backends/chroma/chroma_test.go @@ -26,7 +26,7 @@ func TestBackend_Upsert(t *testing.T) { t.Fatalf("unexpected error: %v", err) } - err = backend.Upsert(context.Background(), "test-id", []float32{0.1, 0.2}, map[string]string{"key": "value"}) + err = backend.Upsert(context.Background(), "test-id", []float32{0.1, 0.2}, "test content", map[string]string{"key": "value"}) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -44,7 +44,7 @@ func TestBackend_Upsert_APIError(t *testing.T) { t.Fatalf("unexpected error: %v", err) } - err = backend.Upsert(context.Background(), "test-id", []float32{0.1}, nil) + err = backend.Upsert(context.Background(), "test-id", []float32{0.1}, "test content", nil) if err == nil { t.Fatal("expected error") } @@ -80,7 +80,7 @@ func TestBackend_Search(t *testing.T) { t.Fatalf("unexpected error: %v", err) } - results, err := backend.Search(context.Background(), []float32{0.1, 0.2}, 5) + results, err := backend.Search(context.Background(), "query", []float32{0.1, 0.2}, 5) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -177,7 +177,7 @@ func TestUpsert_InvalidJSON(t *testing.T) { } // Test with nil metadata (should work) - err = backend.Upsert(context.Background(), "test-id", []float32{0.1}, nil) + err = backend.Upsert(context.Background(), "test-id", []float32{0.1}, "test content", nil) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -203,7 +203,7 @@ func TestSearch_EmptyResults(t *testing.T) { t.Fatalf("unexpected error: %v", err) } - results, err := backend.Search(context.Background(), []float32{0.1, 0.2}, 5) + results, err := backend.Search(context.Background(), "query", []float32{0.1, 0.2}, 5) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -229,7 +229,7 @@ func TestSearch_MalformedResponse(t *testing.T) { t.Fatalf("unexpected error: %v", err) } - _, err = backend.Search(context.Background(), []float32{0.1, 0.2}, 5) + _, err = backend.Search(context.Background(), "query", []float32{0.1, 0.2}, 5) if err == nil { t.Fatal("expected error for malformed response") } @@ -255,7 +255,7 @@ func TestSearch_MissingFields(t *testing.T) { t.Fatalf("unexpected error: %v", err) } - results, err := backend.Search(context.Background(), []float32{0.1, 0.2}, 5) + results, err := backend.Search(context.Background(), "query", []float32{0.1, 0.2}, 5) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -274,7 +274,7 @@ func TestUpsert_ContextCanceled(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() - err = backend.Upsert(ctx, "test-id", []float32{0.1}, nil) + err = backend.Upsert(ctx, "test-id", []float32{0.1}, "test content", nil) if err == nil { t.Fatal("expected error for canceled context") } @@ -289,12 +289,24 @@ func TestSearch_ContextCanceled(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() - _, err = backend.Search(ctx, []float32{0.1, 0.2}, 5) + _, err = backend.Search(ctx, "query", []float32{0.1, 0.2}, 5) if err == nil { t.Fatal("expected error for canceled context") } } +func TestSearch_NoVectorReturnsError(t *testing.T) { + backend, err := chroma.New(chroma.Config{BaseURL: "http://localhost:8000"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + _, err = backend.Search(context.Background(), "query", nil, 5) + if err == nil { + t.Fatal("expected error since chroma has no lexical fallback") + } +} + func TestForget_ContextCanceled(t *testing.T) { backend, err := chroma.New(chroma.Config{BaseURL: "http://localhost:8000"}) if err != nil { diff --git a/pkg/rag/backends/sqlitevec/sqlitevec.go b/pkg/rag/backends/sqlitevec/sqlitevec.go new file mode 100644 index 0000000..c193b56 --- /dev/null +++ b/pkg/rag/backends/sqlitevec/sqlitevec.go @@ -0,0 +1,296 @@ +// Package sqlitevec is a zero-dependency rag.Backend backed by SQLite. When +// given a query vector it scores candidates with a brute-force cosine +// similarity scan in Go (no ANN index, so it trades scale — fine up to a few +// tens of thousands of fragments, comfortably covering a single user's saved +// notes/processes — for requiring nothing beyond the pure-Go sqlite driver +// already used elsewhere in Rony). When no query vector is available (no +// embedder configured, or it failed), it falls back to SQLite's built-in +// FTS5 full-text search over the fragment's content, ranked by BM25. +package sqlitevec + +import ( + "context" + "database/sql" + "encoding/binary" + "encoding/json" + "fmt" + "math" + "sort" + "strings" + "unicode" + + "github.com/VictorVargas/rony-llm-agent/pkg/rag" + _ "modernc.org/sqlite" +) + +// Backend implements rag.Backend on top of a local SQLite file. +type Backend struct { + db *sql.DB +} + +// New opens (or creates) the SQLite-backed vector store at path. An empty +// path opens an in-memory store, useful for tests. +func New(path string) (*Backend, error) { + dsn := path + if dsn == "" { + dsn = "file::memory:?cache=shared" + } + + db, err := sql.Open("sqlite", dsn) + if err != nil { + return nil, fmt.Errorf("open sqlite: %w", err) + } + + b := &Backend{db: db} + if err := b.initSchema(); err != nil { + db.Close() + return nil, fmt.Errorf("init schema: %w", err) + } + return b, nil +} + +func (b *Backend) initSchema() error { + _, err := b.db.Exec(` + CREATE TABLE IF NOT EXISTS fragments ( + id TEXT PRIMARY KEY, + vector BLOB NOT NULL, + content TEXT NOT NULL, + metadata TEXT NOT NULL + ); + CREATE VIRTUAL TABLE IF NOT EXISTS fragments_fts USING fts5( + id UNINDEXED, + content + );`) + return err +} + +// Upsert stores or replaces a fragment's vector, text and metadata, keeping +// the FTS5 index in sync for lexical fallback search. +func (b *Backend) Upsert(ctx context.Context, id string, vector []float32, content string, metadata map[string]string) error { + metaJSON, err := json.Marshal(metadata) + if err != nil { + return fmt.Errorf("marshal metadata: %w", err) + } + + tx, err := b.db.BeginTx(ctx, nil) + if err != nil { + return fmt.Errorf("begin tx: %w", err) + } + defer tx.Rollback() + + _, err = tx.ExecContext(ctx, + `INSERT INTO fragments (id, vector, content, metadata) VALUES (?, ?, ?, ?) + ON CONFLICT(id) DO UPDATE SET vector = excluded.vector, content = excluded.content, metadata = excluded.metadata`, + id, encodeVector(vector), content, string(metaJSON), + ) + if err != nil { + return fmt.Errorf("upsert fragment: %w", err) + } + + // FTS5 tables don't support ON CONFLICT, so re-sync via delete+insert. + if _, err := tx.ExecContext(ctx, `DELETE FROM fragments_fts WHERE id = ?`, id); err != nil { + return fmt.Errorf("clear fts entry: %w", err) + } + if _, err := tx.ExecContext(ctx, `INSERT INTO fragments_fts (id, content) VALUES (?, ?)`, id, content); err != nil { + return fmt.Errorf("index fts entry: %w", err) + } + + return tx.Commit() +} + +// Search scores fragments against queryVector by cosine similarity when one +// is available; otherwise it falls back to an FTS5 lexical match on query. +func (b *Backend) Search(ctx context.Context, query string, queryVector []float32, topK int) ([]rag.SearchResult, error) { + if len(queryVector) > 0 { + return b.searchByVector(ctx, queryVector, topK) + } + return b.searchByText(ctx, query, topK) +} + +func (b *Backend) searchByVector(ctx context.Context, queryVector []float32, topK int) ([]rag.SearchResult, error) { + rows, err := b.db.QueryContext(ctx, `SELECT id, vector, content, metadata FROM fragments`) + if err != nil { + return nil, fmt.Errorf("query fragments: %w", err) + } + defer rows.Close() + + var candidates []rag.SearchResult + for rows.Next() { + var id, content, metaJSON string + var vecBlob []byte + if err := rows.Scan(&id, &vecBlob, &content, &metaJSON); err != nil { + return nil, fmt.Errorf("scan fragment: %w", err) + } + + metadata, err := decodeMetadata(metaJSON) + if err != nil { + return nil, err + } + + candidates = append(candidates, rag.SearchResult{ + ID: id, + Content: content, + Score: cosineSimilarity(queryVector, decodeVector(vecBlob)), + Metadata: metadata, + }) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("iterate fragments: %w", err) + } + + sort.Slice(candidates, func(i, j int) bool { return candidates[i].Score > candidates[j].Score }) + if topK > len(candidates) { + topK = len(candidates) + } + return candidates[:topK], nil +} + +func (b *Backend) searchByText(ctx context.Context, query string, topK int) ([]rag.SearchResult, error) { + ftsQuery := buildFTSQuery(query) + if ftsQuery == "" { + return []rag.SearchResult{}, nil + } + + rows, err := b.db.QueryContext(ctx, ` + SELECT fragments.id, fragments.content, fragments.metadata, bm25(fragments_fts) AS rank + FROM fragments_fts + JOIN fragments ON fragments.id = fragments_fts.id + WHERE fragments_fts MATCH ? + ORDER BY rank + LIMIT ?`, ftsQuery, topK) + if err != nil { + return nil, fmt.Errorf("fts query: %w", err) + } + defer rows.Close() + + var results []rag.SearchResult + for rows.Next() { + var id, content, metaJSON string + var rank float64 + if err := rows.Scan(&id, &content, &metaJSON, &rank); err != nil { + return nil, fmt.Errorf("scan fts result: %w", err) + } + + metadata, err := decodeMetadata(metaJSON) + if err != nil { + return nil, err + } + + // bm25() returns lower-is-better (often negative); negate so a + // higher Score means a better match, matching the vector path. + results = append(results, rag.SearchResult{ + ID: id, + Content: content, + Score: float32(-rank), + Metadata: metadata, + }) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("iterate fts results: %w", err) + } + return results, nil +} + +// Forget deletes a single fragment by ID. +func (b *Backend) Forget(ctx context.Context, id string) error { + tx, err := b.db.BeginTx(ctx, nil) + if err != nil { + return fmt.Errorf("begin tx: %w", err) + } + defer tx.Rollback() + + if _, err := tx.ExecContext(ctx, `DELETE FROM fragments WHERE id = ?`, id); err != nil { + return fmt.Errorf("delete fragment: %w", err) + } + if _, err := tx.ExecContext(ctx, `DELETE FROM fragments_fts WHERE id = ?`, id); err != nil { + return fmt.Errorf("delete fts entry: %w", err) + } + return tx.Commit() +} + +// ForgetAll deletes every stored fragment. +func (b *Backend) ForgetAll(ctx context.Context) error { + tx, err := b.db.BeginTx(ctx, nil) + if err != nil { + return fmt.Errorf("begin tx: %w", err) + } + defer tx.Rollback() + + if _, err := tx.ExecContext(ctx, `DELETE FROM fragments`); err != nil { + return fmt.Errorf("delete all fragments: %w", err) + } + if _, err := tx.ExecContext(ctx, `DELETE FROM fragments_fts`); err != nil { + return fmt.Errorf("delete all fts entries: %w", err) + } + return tx.Commit() +} + +// Close releases the underlying database connection. +func (b *Backend) Close() error { + return b.db.Close() +} + +func decodeMetadata(metaJSON string) (map[string]string, error) { + if metaJSON == "" { + return nil, nil + } + var metadata map[string]string + if err := json.Unmarshal([]byte(metaJSON), &metadata); err != nil { + return nil, fmt.Errorf("unmarshal metadata: %w", err) + } + return metadata, nil +} + +// buildFTSQuery turns free-form text into an FTS5 MATCH query that ORs +// together each token as a quoted phrase, so punctuation or FTS5 operator +// characters (-, *, :, "...) in the input can't produce a syntax error, and +// any subset of tokens can match (rather than requiring the exact phrase). +func buildFTSQuery(text string) string { + tokens := tokenizeForFTS(text) + if len(tokens) == 0 { + return "" + } + parts := make([]string, len(tokens)) + for i, tok := range tokens { + parts[i] = `"` + strings.ReplaceAll(tok, `"`, `""`) + `"` + } + return strings.Join(parts, " OR ") +} + +func tokenizeForFTS(text string) []string { + return strings.FieldsFunc(strings.ToLower(text), func(r rune) bool { + return !unicode.IsLetter(r) && !unicode.IsNumber(r) + }) +} + +func encodeVector(v []float32) []byte { + buf := make([]byte, 4*len(v)) + for i, f := range v { + binary.LittleEndian.PutUint32(buf[i*4:], math.Float32bits(f)) + } + return buf +} + +func decodeVector(buf []byte) []float32 { + v := make([]float32, len(buf)/4) + for i := range v { + v[i] = math.Float32frombits(binary.LittleEndian.Uint32(buf[i*4:])) + } + return v +} + +func cosineSimilarity(a, b []float32) float32 { + if len(a) == 0 || len(b) == 0 || len(a) != len(b) { + return 0 + } + var dot, na, nb float64 + for i := range a { + dot += float64(a[i]) * float64(b[i]) + na += float64(a[i]) * float64(a[i]) + nb += float64(b[i]) * float64(b[i]) + } + if na == 0 || nb == 0 { + return 0 + } + return float32(dot / (math.Sqrt(na) * math.Sqrt(nb))) +} diff --git a/pkg/rag/backends/sqlitevec/sqlitevec_test.go b/pkg/rag/backends/sqlitevec/sqlitevec_test.go new file mode 100644 index 0000000..f176ace --- /dev/null +++ b/pkg/rag/backends/sqlitevec/sqlitevec_test.go @@ -0,0 +1,240 @@ +package sqlitevec_test + +import ( + "context" + "testing" + + "github.com/VictorVargas/rony-llm-agent/pkg/rag/backends/sqlitevec" +) + +func newTestBackend(t *testing.T) *sqlitevec.Backend { + t.Helper() + b, err := sqlitevec.New("") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + t.Cleanup(func() { b.Close() }) + return b +} + +func TestBackend_UpsertAndSearchByVector(t *testing.T) { + b := newTestBackend(t) + ctx := context.Background() + + if err := b.Upsert(ctx, "a", []float32{1, 0, 0}, "cómo desplegar a producción", map[string]string{"kind": "process"}); err != nil { + t.Fatalf("upsert a: %v", err) + } + if err := b.Upsert(ctx, "b", []float32{0, 1, 0}, "receta de pan", nil); err != nil { + t.Fatalf("upsert b: %v", err) + } + + results, err := b.Search(ctx, "", []float32{1, 0, 0}, 5) + if err != nil { + t.Fatalf("search: %v", err) + } + if len(results) != 2 { + t.Fatalf("expected 2 results, got %d", len(results)) + } + if results[0].ID != "a" { + t.Fatalf("expected closest match to be 'a', got %q", results[0].ID) + } + if results[0].Content != "cómo desplegar a producción" { + t.Fatalf("expected content to round-trip, got %q", results[0].Content) + } + if results[0].Metadata["kind"] != "process" { + t.Fatalf("expected metadata to round-trip, got %v", results[0].Metadata) + } +} + +func TestBackend_UpsertReplacesExisting(t *testing.T) { + b := newTestBackend(t) + ctx := context.Background() + + if err := b.Upsert(ctx, "a", []float32{1, 0}, "first version", nil); err != nil { + t.Fatalf("upsert: %v", err) + } + if err := b.Upsert(ctx, "a", []float32{0, 1}, "second version", nil); err != nil { + t.Fatalf("upsert: %v", err) + } + + results, err := b.Search(ctx, "", []float32{0, 1}, 5) + if err != nil { + t.Fatalf("search: %v", err) + } + if len(results) != 1 { + t.Fatalf("expected the upsert to replace, not duplicate; got %d results", len(results)) + } + if results[0].Content != "second version" { + t.Fatalf("expected replaced content, got %q", results[0].Content) + } + + // The FTS5 side must also have been replaced, not duplicated. + ftsResults, err := b.Search(ctx, "second version", nil, 5) + if err != nil { + t.Fatalf("fts search: %v", err) + } + if len(ftsResults) != 1 { + t.Fatalf("expected 1 fts result after replace, got %d", len(ftsResults)) + } +} + +func TestBackend_SearchTopK(t *testing.T) { + b := newTestBackend(t) + ctx := context.Background() + + for i, id := range []string{"a", "b", "c"} { + vec := []float32{float32(i), 1, 1} + if err := b.Upsert(ctx, id, vec, id, nil); err != nil { + t.Fatalf("upsert %s: %v", id, err) + } + } + + results, err := b.Search(ctx, "", []float32{1, 1, 1}, 2) + if err != nil { + t.Fatalf("search: %v", err) + } + if len(results) != 2 { + t.Fatalf("expected topK=2 results, got %d", len(results)) + } +} + +func TestBackend_Forget(t *testing.T) { + b := newTestBackend(t) + ctx := context.Background() + + if err := b.Upsert(ctx, "a", []float32{1, 0}, "content", nil); err != nil { + t.Fatalf("upsert: %v", err) + } + if err := b.Forget(ctx, "a"); err != nil { + t.Fatalf("forget: %v", err) + } + + results, err := b.Search(ctx, "", []float32{1, 0}, 5) + if err != nil { + t.Fatalf("search: %v", err) + } + if len(results) != 0 { + t.Fatalf("expected no results after forget, got %d", len(results)) + } + + ftsResults, err := b.Search(ctx, "content", nil, 5) + if err != nil { + t.Fatalf("fts search: %v", err) + } + if len(ftsResults) != 0 { + t.Fatalf("expected no fts results after forget, got %d", len(ftsResults)) + } +} + +func TestBackend_ForgetAll(t *testing.T) { + b := newTestBackend(t) + ctx := context.Background() + + for _, id := range []string{"a", "b"} { + if err := b.Upsert(ctx, id, []float32{1, 0}, id, nil); err != nil { + t.Fatalf("upsert %s: %v", id, err) + } + } + if err := b.ForgetAll(ctx); err != nil { + t.Fatalf("forget all: %v", err) + } + + results, err := b.Search(ctx, "", []float32{1, 0}, 5) + if err != nil { + t.Fatalf("search: %v", err) + } + if len(results) != 0 { + t.Fatalf("expected no results after forget all, got %d", len(results)) + } +} + +func TestBackend_SearchEmpty(t *testing.T) { + b := newTestBackend(t) + + results, err := b.Search(context.Background(), "", []float32{1, 0}, 5) + if err != nil { + t.Fatalf("search: %v", err) + } + if len(results) != 0 { + t.Fatalf("expected no results on empty store, got %d", len(results)) + } +} + +func TestBackend_SearchByText_NoVectorFallsBackToFTS(t *testing.T) { + b := newTestBackend(t) + ctx := context.Background() + + // Upserted with an empty vector, as if the embedder had failed. + if err := b.Upsert(ctx, "a", nil, "cómo desplegar a producción con Docker", map[string]string{"kind": "process"}); err != nil { + t.Fatalf("upsert a: %v", err) + } + if err := b.Upsert(ctx, "b", nil, "receta de pan con masa madre", nil); err != nil { + t.Fatalf("upsert b: %v", err) + } + + results, err := b.Search(ctx, "desplegar producción", nil, 5) + if err != nil { + t.Fatalf("search: %v", err) + } + if len(results) != 1 { + t.Fatalf("expected 1 fts match, got %d", len(results)) + } + if results[0].ID != "a" { + t.Fatalf("expected match to be 'a', got %q", results[0].ID) + } + if results[0].Metadata["kind"] != "process" { + t.Fatalf("expected metadata to round-trip through fts path, got %v", results[0].Metadata) + } +} + +func TestBackend_SearchByText_NoMatches(t *testing.T) { + b := newTestBackend(t) + ctx := context.Background() + + if err := b.Upsert(ctx, "a", nil, "receta de pan", nil); err != nil { + t.Fatalf("upsert: %v", err) + } + + results, err := b.Search(ctx, "algo completamente distinto", nil, 5) + if err != nil { + t.Fatalf("search: %v", err) + } + if len(results) != 0 { + t.Fatalf("expected no matches, got %d", len(results)) + } +} + +func TestBackend_SearchByText_HandlesSpecialCharacters(t *testing.T) { + b := newTestBackend(t) + ctx := context.Background() + + if err := b.Upsert(ctx, "a", nil, "usa docker-compose para levantar todo", nil); err != nil { + t.Fatalf("upsert: %v", err) + } + + // FTS5 operator characters in the query must not cause a syntax error. + results, err := b.Search(ctx, `docker-compose "up" AND/OR *test*`, nil, 5) + if err != nil { + t.Fatalf("search should not error on special characters: %v", err) + } + if len(results) == 0 { + t.Fatal("expected at least one match despite special characters in the query") + } +} + +func TestBackend_SearchByText_EmptyQuery(t *testing.T) { + b := newTestBackend(t) + ctx := context.Background() + + if err := b.Upsert(ctx, "a", nil, "algo", nil); err != nil { + t.Fatalf("upsert: %v", err) + } + + results, err := b.Search(ctx, " ", nil, 5) + if err != nil { + t.Fatalf("search: %v", err) + } + if len(results) != 0 { + t.Fatalf("expected no results for an empty query, got %d", len(results)) + } +} diff --git a/pkg/rag/embeddings/openai_compatible.go b/pkg/rag/embeddings/openai_compatible.go new file mode 100644 index 0000000..77e6541 --- /dev/null +++ b/pkg/rag/embeddings/openai_compatible.go @@ -0,0 +1,123 @@ +package embeddings + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "sync/atomic" +) + +// OpenAICompatibleConfig holds the settings for any embeddings API that +// follows OpenAI's request/response shape: POST {BaseURL}/embeddings with +// {"input": ..., "model": ...}, returning {"data": [{"embedding": [...]}]}. +// This covers llama.cpp (started with --embeddings), real OpenAI, and most +// third-party providers advertised as "OpenAI-compatible". +type OpenAICompatibleConfig struct { + BaseURL string // e.g. "http://localhost:8080/v1" or "https://api.openai.com/v1" + APIKey string // sent as "Authorization: Bearer " when non-empty; local servers like llama.cpp don't need one + Model string // e.g. "text-embedding-3-small"; ignored by servers that only have one model loaded +} + +// OpenAICompatible implements Embedder against any OpenAI-shaped +// /embeddings endpoint. +type OpenAICompatible struct { + baseURL string + apiKey string + model string + http *http.Client + dims atomic.Int64 // lazily learned from the first successful response +} + +// NewOpenAICompatible creates a new OpenAI-shaped embedder. +func NewOpenAICompatible(cfg OpenAICompatibleConfig) (*OpenAICompatible, error) { + if cfg.BaseURL == "" { + return nil, fmt.Errorf("base URL is required") + } + return &OpenAICompatible{ + baseURL: cfg.BaseURL, + apiKey: cfg.APIKey, + model: cfg.Model, + http: http.DefaultClient, + }, nil +} + +// NewLlamaCpp is a convenience constructor for a local llama.cpp server +// (defaults to http://localhost:8080/v1, no API key). The server must have +// been started with the `--embeddings` flag, otherwise every call fails +// (llama.cpp returns a 501). Quality depends on the loaded model: dedicated +// embedding models (e.g. nomic-embed-text, bge-m3) work best, but a +// chat/instruct model still produces a usable semantic vector via pooling. +func NewLlamaCpp(cfg LlamaCppConfig) (*OpenAICompatible, error) { + baseURL := cfg.BaseURL + if baseURL == "" { + baseURL = "http://localhost:8080/v1" + } + return NewOpenAICompatible(OpenAICompatibleConfig{BaseURL: baseURL, Model: cfg.Model}) +} + +// LlamaCppConfig holds the settings for NewLlamaCpp. +type LlamaCppConfig struct { + BaseURL string // e.g. "http://localhost:8080/v1" + Model string // optional; llama.cpp embeds with whatever model is loaded regardless of this value +} + +func (e *OpenAICompatible) Embed(ctx context.Context, text string) ([]float32, error) { + endpoint := e.baseURL + "/embeddings" + + reqBody, err := json.Marshal(map[string]interface{}{ + "input": text, + "model": e.model, + }) + if err != nil { + return nil, fmt.Errorf("marshaling request: %w", err) + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(reqBody)) + if err != nil { + return nil, fmt.Errorf("creating request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + if e.apiKey != "" { + httpReq.Header.Set("Authorization", "Bearer "+e.apiKey) + } + + resp, err := e.http.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("request failed: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + body, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("API error %d: %s", resp.StatusCode, string(body)) + } + + var apiResp openAICompatibleEmbedResponse + if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil { + return nil, fmt.Errorf("decoding response: %w", err) + } + if len(apiResp.Data) == 0 || len(apiResp.Data[0].Embedding) == 0 { + return nil, fmt.Errorf("empty embeddings response") + } + + vector := apiResp.Data[0].Embedding + e.dims.Store(int64(len(vector))) + return vector, nil +} + +// Dimensions returns the vector size learned from the last successful Embed +// call, or 0 if none has succeeded yet (it depends on the model/provider +// behind BaseURL, so it can't be known upfront). +func (e *OpenAICompatible) Dimensions() int { + return int(e.dims.Load()) +} + +type openAICompatibleEmbedResponse struct { + Data []struct { + Embedding []float32 `json:"embedding"` + Index int `json:"index"` + } `json:"data"` +} diff --git a/pkg/rag/embeddings/openai_compatible_test.go b/pkg/rag/embeddings/openai_compatible_test.go new file mode 100644 index 0000000..fdb87d2 --- /dev/null +++ b/pkg/rag/embeddings/openai_compatible_test.go @@ -0,0 +1,134 @@ +package embeddings + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" +) + +func TestNewOpenAICompatible_RequiresBaseURL(t *testing.T) { + _, err := NewOpenAICompatible(OpenAICompatibleConfig{}) + if err == nil { + t.Fatal("expected error when BaseURL is empty") + } +} + +func TestNewLlamaCpp_Defaults(t *testing.T) { + e, err := NewLlamaCpp(LlamaCppConfig{}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if e.baseURL != "http://localhost:8080/v1" { + t.Errorf("expected default base URL, got %q", e.baseURL) + } +} + +func TestNewLlamaCpp_Custom(t *testing.T) { + e, err := NewLlamaCpp(LlamaCppConfig{BaseURL: "http://custom:9000/v1", Model: "my-model"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if e.baseURL != "http://custom:9000/v1" { + t.Errorf("expected custom base URL, got %q", e.baseURL) + } +} + +func TestOpenAICompatible_Embed_Success(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/embeddings" { + w.WriteHeader(http.StatusNotFound) + return + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"data":[{"embedding":[0.1,0.2,0.3],"index":0}]}`)) + })) + defer server.Close() + + e, err := NewLlamaCpp(LlamaCppConfig{BaseURL: server.URL + "/v1"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + vector, err := e.Embed(context.Background(), "hola") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(vector) != 3 { + t.Fatalf("expected 3 dimensions, got %d", len(vector)) + } + if e.Dimensions() != 3 { + t.Errorf("expected Dimensions() to learn 3 after a successful call, got %d", e.Dimensions()) + } +} + +func TestOpenAICompatible_Embed_NotEnabled(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotImplemented) + w.Write([]byte(`{"error":{"message":"This server does not support embeddings. Start it with ` + "`--embeddings`" + `"}}`)) + })) + defer server.Close() + + e, err := NewLlamaCpp(LlamaCppConfig{BaseURL: server.URL + "/v1"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if _, err := e.Embed(context.Background(), "hola"); err == nil { + t.Fatal("expected error when the server doesn't support embeddings") + } +} + +func TestOpenAICompatible_SendsBearerTokenWhenConfigured(t *testing.T) { + var gotAuth string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"data":[{"embedding":[0.1],"index":0}]}`)) + })) + defer server.Close() + + e, err := NewOpenAICompatible(OpenAICompatibleConfig{ + BaseURL: server.URL, + APIKey: "sk-test", + Model: "text-embedding-3-small", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if _, err := e.Embed(context.Background(), "hola"); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if gotAuth != "Bearer sk-test" { + t.Fatalf("expected Authorization header to be sent, got %q", gotAuth) + } +} + +func TestOpenAICompatible_NoAuthHeaderWhenNoAPIKey(t *testing.T) { + var gotAuth string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"data":[{"embedding":[0.1],"index":0}]}`)) + })) + defer server.Close() + + e, err := NewOpenAICompatible(OpenAICompatibleConfig{BaseURL: server.URL}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if _, err := e.Embed(context.Background(), "hola"); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if gotAuth != "" { + t.Fatalf("expected no Authorization header without an API key, got %q", gotAuth) + } +} + +func TestOpenAICompatible_Dimensions_BeforeAnyCall(t *testing.T) { + e, err := NewLlamaCpp(LlamaCppConfig{}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if dims := e.Dimensions(); dims != 0 { + t.Errorf("expected 0 dimensions before any successful call, got %d", dims) + } +} diff --git a/pkg/rag/memory.go b/pkg/rag/memory.go index fb5cb87..0e0b7d1 100644 --- a/pkg/rag/memory.go +++ b/pkg/rag/memory.go @@ -32,10 +32,14 @@ type Config struct { Embedder Embedder } -// Backend is the interface for vector database backends. +// Backend is the interface for storage backends. queryVector is nil when no +// embedder produced one (e.g. it's unavailable or failed); backends that +// can't search without a vector (e.g. a pure vector database like Chroma) +// should return an error in that case, while backends capable of lexical +// search (e.g. SQLite FTS5) can fall back to matching on query instead. type Backend interface { - Upsert(ctx context.Context, id string, vector []float32, metadata map[string]string) error - Search(ctx context.Context, queryVector []float32, topK int) ([]SearchResult, error) + Upsert(ctx context.Context, id string, vector []float32, content string, metadata map[string]string) error + Search(ctx context.Context, query string, queryVector []float32, topK int) ([]SearchResult, error) Forget(ctx context.Context, id string) error ForgetAll(ctx context.Context) error } @@ -84,13 +88,16 @@ func (m *memory) Add(ctx context.Context, fragment Fragment) error { fragment.Metadata["project_id"] = fragment.ProjectID fragment.Timestamp = time.Now() + // A failed embedding doesn't block saving: the backend still has the + // raw content and can index it for lexical search (e.g. FTS5), so the + // fragment just won't be reachable by vector similarity later. vector, err := m.embedder.Embed(ctx, fragment.Content) if err != nil { - return fmt.Errorf("embedding: %w", err) + vector = nil } fragment.Vector = vector - return m.backend.Upsert(ctx, fragment.ID, fragment.Vector, fragment.Metadata) + return m.backend.Upsert(ctx, fragment.ID, fragment.Vector, fragment.Content, fragment.Metadata) } func (m *memory) Search(ctx context.Context, query string, topK int) ([]Fragment, error) { @@ -98,12 +105,14 @@ func (m *memory) Search(ctx context.Context, query string, topK int) ([]Fragment topK = 5 } + // Same fallback as Add: if embedding the query fails, search proceeds + // with no vector so the backend can fall back to lexical matching. queryVector, err := m.embedder.Embed(ctx, query) if err != nil { - return nil, fmt.Errorf("embedding query: %w", err) + queryVector = nil } - results, err := m.backend.Search(ctx, queryVector, topK) + results, err := m.backend.Search(ctx, query, queryVector, topK) if err != nil { return nil, fmt.Errorf("search: %w", err) } diff --git a/pkg/rag/memory_test.go b/pkg/rag/memory_test.go index e585eff..dbe77f8 100644 --- a/pkg/rag/memory_test.go +++ b/pkg/rag/memory_test.go @@ -51,9 +51,41 @@ func TestMemory_Add(t *testing.T) { } } -func TestMemory_Add_EmbeddingError(t *testing.T) { +func TestMemory_Add_PassesContentToBackend(t *testing.T) { + var gotContent string m, err := rag.New(rag.Config{ - Backend: &mockBackend{}, + Backend: &mockBackend{ + upsertFunc: func(_ context.Context, _ string, _ []float32, content string, _ map[string]string) error { + gotContent = content + return nil + }, + }, + Embedder: &embeddings.MockEmbedder{}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + err = m.Add(context.Background(), rag.Fragment{Content: "remember this process"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if gotContent != "remember this process" { + t.Fatalf("expected backend to receive the fragment content, got %q", gotContent) + } +} + +func TestMemory_Add_EmbeddingErrorStillSavesWithNoVector(t *testing.T) { + var gotVector []float32 + sawCall := false + m, err := rag.New(rag.Config{ + Backend: &mockBackend{ + upsertFunc: func(_ context.Context, _ string, vector []float32, _ string, _ map[string]string) error { + sawCall = true + gotVector = vector + return nil + }, + }, Embedder: &embeddings.MockEmbedder{EmbedFunc: func(ctx context.Context, text string) ([]float32, error) { return nil, fmt.Errorf("embed error") }}, }) if err != nil { @@ -63,15 +95,21 @@ func TestMemory_Add_EmbeddingError(t *testing.T) { err = m.Add(context.Background(), rag.Fragment{ Content: "test content", }) - if err == nil { - t.Fatal("expected error for embedding failure") + if err != nil { + t.Fatalf("expected Add to succeed so the backend can still index the text lexically, got: %v", err) + } + if !sawCall { + t.Fatal("expected the backend to still be called despite the embedding failure") + } + if len(gotVector) != 0 { + t.Fatalf("expected no vector to be passed through, got %v", gotVector) } } func TestMemory_Search(t *testing.T) { m, err := rag.New(rag.Config{ Backend: &mockBackend{ - searchFunc: func(ctx context.Context, queryVector []float32, topK int) ([]rag.SearchResult, error) { + searchFunc: func(ctx context.Context, query string, queryVector []float32, topK int) ([]rag.SearchResult, error) { return []rag.SearchResult{ {ID: "1", Content: "result 1", Score: 0.9}, }, nil @@ -119,9 +157,17 @@ func TestMemory_ForgetAll(t *testing.T) { } } -func TestMemory_Search_EmbeddingError(t *testing.T) { +func TestMemory_Search_EmbeddingErrorFallsBackToLexicalSearch(t *testing.T) { + var gotQuery string + var gotVector []float32 m, err := rag.New(rag.Config{ - Backend: &mockBackend{}, + Backend: &mockBackend{ + searchFunc: func(_ context.Context, query string, queryVector []float32, _ int) ([]rag.SearchResult, error) { + gotQuery = query + gotVector = queryVector + return []rag.SearchResult{{ID: "1", Content: "matched lexically"}}, nil + }, + }, Embedder: &embeddings.MockEmbedder{EmbedFunc: func(ctx context.Context, text string) ([]float32, error) { return nil, fmt.Errorf("embed error") }}, @@ -130,29 +176,38 @@ func TestMemory_Search_EmbeddingError(t *testing.T) { t.Fatalf("unexpected error: %v", err) } - _, err = m.Search(context.Background(), "test query", 5) - if err == nil { - t.Fatal("expected error for embedding failure") + results, err := m.Search(context.Background(), "test query", 5) + if err != nil { + t.Fatalf("expected Search to fall back to the backend's lexical search, got error: %v", err) + } + if len(results) != 1 || results[0].Content != "matched lexically" { + t.Fatalf("expected the backend's fallback result to come through, got %v", results) + } + if gotQuery != "test query" { + t.Fatalf("expected the raw query text to reach the backend, got %q", gotQuery) + } + if len(gotVector) != 0 { + t.Fatalf("expected no query vector to be passed through, got %v", gotVector) } } // mockBackend implements chroma.Backend for testing. type mockBackend struct { - upsertFunc func(ctx context.Context, id string, vector []float32, metadata map[string]string) error - searchFunc func(ctx context.Context, queryVector []float32, topK int) ([]rag.SearchResult, error) + upsertFunc func(ctx context.Context, id string, vector []float32, content string, metadata map[string]string) error + searchFunc func(ctx context.Context, query string, queryVector []float32, topK int) ([]rag.SearchResult, error) forgetAllFunc func(ctx context.Context) error } -func (m *mockBackend) Upsert(ctx context.Context, id string, vector []float32, metadata map[string]string) error { +func (m *mockBackend) Upsert(ctx context.Context, id string, vector []float32, content string, metadata map[string]string) error { if m.upsertFunc != nil { - return m.upsertFunc(ctx, id, vector, metadata) + return m.upsertFunc(ctx, id, vector, content, metadata) } return nil } -func (m *mockBackend) Search(ctx context.Context, queryVector []float32, topK int) ([]rag.SearchResult, error) { +func (m *mockBackend) Search(ctx context.Context, query string, queryVector []float32, topK int) ([]rag.SearchResult, error) { if m.searchFunc != nil { - return m.searchFunc(ctx, queryVector, topK) + return m.searchFunc(ctx, query, queryVector, topK) } return nil, nil } @@ -190,7 +245,7 @@ func TestMemory_Add_MultipleFragments(t *testing.T) { func TestMemory_Search_EmptyQuery(t *testing.T) { m, err := rag.New(rag.Config{ Backend: &mockBackend{ - searchFunc: func(ctx context.Context, queryVector []float32, topK int) ([]rag.SearchResult, error) { + searchFunc: func(ctx context.Context, query string, queryVector []float32, topK int) ([]rag.SearchResult, error) { return []rag.SearchResult{}, nil }, },