feat(rag): add SQLite+FTS5 backend, fix content/usage plumbing bugs
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 <noreply@anthropic.com>
This commit is contained in:
parent
2cde90d8f7
commit
0652023037
14 changed files with 1092 additions and 39 deletions
16
go.mod
16
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
|
||||
)
|
||||
|
|
|
|||
49
go.sum
49
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=
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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{
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -60,6 +60,7 @@ func (b *Backend) Upsert(ctx context.Context, id string, vector []float32, metad
|
|||
reqBody, err := json.Marshal(map[string]interface{}{
|
||||
"ids": []string{id},
|
||||
"embeddings": embeddings,
|
||||
"documents": []string{content},
|
||||
"metadatas": []map[string]interface{}{metadatas},
|
||||
})
|
||||
if err != nil {
|
||||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
296
pkg/rag/backends/sqlitevec/sqlitevec.go
Normal file
296
pkg/rag/backends/sqlitevec/sqlitevec.go
Normal file
|
|
@ -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)))
|
||||
}
|
||||
240
pkg/rag/backends/sqlitevec/sqlitevec_test.go
Normal file
240
pkg/rag/backends/sqlitevec/sqlitevec_test.go
Normal file
|
|
@ -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))
|
||||
}
|
||||
}
|
||||
123
pkg/rag/embeddings/openai_compatible.go
Normal file
123
pkg/rag/embeddings/openai_compatible.go
Normal file
|
|
@ -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 <key>" 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"`
|
||||
}
|
||||
134
pkg/rag/embeddings/openai_compatible_test.go
Normal file
134
pkg/rag/embeddings/openai_compatible_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
},
|
||||
},
|
||||
|
|
|
|||
Loading…
Reference in a new issue