From 9f33d6df9360da92a4e9305126f79db2b6d9bff6 Mon Sep 17 00:00:00 2001 From: Victor Vargas Date: Tue, 30 Jun 2026 23:53:30 -0700 Subject: [PATCH] feat(rag): add Retrieval-Augmented Memory with ChromaDB backend and embedding support --- pkg/rag/backends/chroma/chroma.go | 231 ++++++++++++++++++ pkg/rag/backends/chroma/chroma_test.go | 311 +++++++++++++++++++++++++ pkg/rag/embeddings/mock.go | 25 ++ pkg/rag/embeddings/ollama.go | 86 +++++++ pkg/rag/memory.go | 129 ++++++++++ pkg/rag/memory_test.go | 230 ++++++++++++++++++ 6 files changed, 1012 insertions(+) create mode 100644 pkg/rag/backends/chroma/chroma.go create mode 100644 pkg/rag/backends/chroma/chroma_test.go create mode 100644 pkg/rag/embeddings/mock.go create mode 100644 pkg/rag/embeddings/ollama.go create mode 100644 pkg/rag/memory.go create mode 100644 pkg/rag/memory_test.go diff --git a/pkg/rag/backends/chroma/chroma.go b/pkg/rag/backends/chroma/chroma.go new file mode 100644 index 0000000..6eda10e --- /dev/null +++ b/pkg/rag/backends/chroma/chroma.go @@ -0,0 +1,231 @@ +package chroma + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/VictorVargas/rony-llm-agent/pkg/rag" +) + +// Config holds the settings for the ChromaDB backend. +type Config struct { + BaseURL string // e.g. "http://localhost:8000" + Timeout int // request timeout in seconds (0 = default) +} + +// Backend implements rag.Backend using ChromaDB's REST API. +type Backend struct { + baseURL string + http *http.Client +} + +// New creates a new ChromaDB backend. +func New(cfg Config) (*Backend, error) { + baseURL := cfg.BaseURL + if baseURL == "" { + baseURL = "http://localhost:8000" + } + + timeout := time.Duration(cfg.Timeout) * time.Second + if timeout == 0 { + timeout = 30 * time.Second + } + + return &Backend{ + baseURL: baseURL, + http: &http.Client{ + Timeout: timeout, + }, + }, nil +} + +func (b *Backend) Upsert(ctx context.Context, id string, vector []float32, metadata map[string]string) error { + collection := "rony-memory" + + embeddings := make([][]float64, 1) + for _, f := range vector { + embeddings[0] = append(embeddings[0], float64(f)) + } + + metadatas := make(map[string]interface{}) + for k, v := range metadata { + metadatas[k] = v + } + + reqBody, err := json.Marshal(map[string]interface{}{ + "ids": []string{id}, + "embeddings": embeddings, + "metadatas": []map[string]interface{}{metadatas}, + }) + if err != nil { + return fmt.Errorf("marshaling upsert request: %w", err) + } + + endpoint := fmt.Sprintf("/api/v1/collections/%s/upsert", collection) + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, b.baseURL+endpoint, strings.NewReader(string(reqBody))) + if err != nil { + return fmt.Errorf("creating request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + + resp, err := b.http.Do(httpReq) + if err != nil { + return fmt.Errorf("request failed: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + body, _ := io.ReadAll(resp.Body) + return fmt.Errorf("API error %d: %s", resp.StatusCode, string(body)) + } + + return nil +} + +func (b *Backend) Search(ctx context.Context, queryVector []float32, topK int) ([]rag.SearchResult, error) { + collection := "rony-memory" + + query := make([]float64, len(queryVector)) + for i, f := range queryVector { + query[i] = float64(f) + } + + reqBody, err := json.Marshal(map[string]interface{}{ + "queries": []map[string]interface{}{ + { + "vector": query, + "n_results": topK, + }, + }, + }) + if err != nil { + return nil, fmt.Errorf("marshaling search request: %w", err) + } + + endpoint := fmt.Sprintf("/api/v1/collections/%s/query", collection) + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, b.baseURL+endpoint, strings.NewReader(string(reqBody))) + if err != nil { + return nil, fmt.Errorf("creating request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + + resp, err := b.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 chromaQueryResponse + if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil { + return nil, fmt.Errorf("decoding response: %w", err) + } + + if len(apiResp.Results) == 0 || len(apiResp.Results[0].IDs) == 0 { + return []rag.SearchResult{}, nil + } + + results := make([]rag.SearchResult, len(apiResp.Results[0].IDs[0])) + for i := range apiResp.Results[0].IDs[0] { + var metadata map[string]string + if len(apiResp.Results[0].Metadatas) > 0 && len(apiResp.Results[0].Metadatas[0]) > i { + metadata = stringifyMap(apiResp.Results[0].Metadatas[0][i]) + } + + results[i] = rag.SearchResult{ + ID: apiResp.Results[0].IDs[0][i], + Content: apiResp.Results[0].Documents[0][i], + Score: float32(apiResp.Results[0].Distances[0][i]), + Metadata: metadata, + } + } + + return results, nil +} + +func (b *Backend) Forget(ctx context.Context, id string) error { + collection := "rony-memory" + + reqBody, err := json.Marshal(map[string]interface{}{ + "ids": []string{id}, + }) + if err != nil { + return fmt.Errorf("marshaling delete request: %w", err) + } + + endpoint := fmt.Sprintf("/api/v1/collections/%s/delete", collection) + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, b.baseURL+endpoint, strings.NewReader(string(reqBody))) + if err != nil { + return fmt.Errorf("creating request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + + resp, err := b.http.Do(httpReq) + if err != nil { + return fmt.Errorf("request failed: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + body, _ := io.ReadAll(resp.Body) + return fmt.Errorf("API error %d: %s", resp.StatusCode, string(body)) + } + + return nil +} + +func (b *Backend) ForgetAll(ctx context.Context) error { + collection := "rony-memory" + + reqBody := []byte(`{"where": {}}`) + + endpoint := fmt.Sprintf("/api/v1/collections/%s/delete", collection) + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, b.baseURL+endpoint, strings.NewReader(string(reqBody))) + if err != nil { + return fmt.Errorf("creating request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + + resp, err := b.http.Do(httpReq) + if err != nil { + return fmt.Errorf("request failed: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + body, _ := io.ReadAll(resp.Body) + return fmt.Errorf("API error %d: %s", resp.StatusCode, string(body)) + } + + return nil +} + +func stringifyMap(m map[string]interface{}) map[string]string { + result := make(map[string]string) + for k, v := range m { + result[k] = fmt.Sprintf("%v", v) + } + return result +} + +// chromaQueryResponse represents the structure of a ChromaDB query response. +type chromaQueryResponse struct { + Names []string `json:"names"` + Results []chromaQueryResults `json:"results"` +} + +type chromaQueryResults struct { + IDs [][]string `json:"ids"` + Documents [][]string `json:"documents"` + Distances [][]float64 `json:"distances"` + Metadatas [][]map[string]interface{} `json:"metadatas"` +} diff --git a/pkg/rag/backends/chroma/chroma_test.go b/pkg/rag/backends/chroma/chroma_test.go new file mode 100644 index 0000000..fa73ec9 --- /dev/null +++ b/pkg/rag/backends/chroma/chroma_test.go @@ -0,0 +1,311 @@ +package chroma_test + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/VictorVargas/rony-llm-agent/pkg/rag/backends/chroma" +) + +func TestBackend_Upsert(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/v1/collections/rony-memory/upsert" { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{}`)) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer server.Close() + + backend, err := chroma.New(chroma.Config{BaseURL: server.URL}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + err = backend.Upsert(context.Background(), "test-id", []float32{0.1, 0.2}, map[string]string{"key": "value"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestBackend_Upsert_APIError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(`{"error":"server error"}`)) + })) + defer server.Close() + + backend, err := chroma.New(chroma.Config{BaseURL: server.URL}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + err = backend.Upsert(context.Background(), "test-id", []float32{0.1}, nil) + if err == nil { + t.Fatal("expected error") + } +} + +func TestBackend_Search(t *testing.T) { + meta := []map[string]interface{}{{"key": "value"}} + metaNested := [][]map[string]interface{}{meta} + mockResponse := map[string]interface{}{ + "names": []string{"rony-memory"}, + "results": []map[string]interface{}{ + { + "ids": [][]string{{"test-id"}}, + "documents": [][]string{{"test content"}}, + "distances": [][]float64{{0.9}}, + "metadatas": metaNested, + }, + }, + } + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/v1/collections/rony-memory/query" { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(mockResponse) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer server.Close() + + backend, err := chroma.New(chroma.Config{BaseURL: server.URL}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + results, err := backend.Search(context.Background(), []float32{0.1, 0.2}, 5) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if len(results) != 1 { + t.Fatalf("expected 1 result, got %d", len(results)) + } + + if results[0].ID != "test-id" { + t.Errorf("expected 'test-id', got %q", results[0].ID) + } + + if results[0].Content != "test content" { + t.Errorf("expected 'test content', got %q", results[0].Content) + } +} + +func TestBackend_ForgetAll(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/v1/collections/rony-memory/delete" { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{}`)) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer server.Close() + + backend, err := chroma.New(chroma.Config{BaseURL: server.URL}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + err = backend.ForgetAll(context.Background()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestBackend_ForgetAll_APIError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(`{"error":"server error"}`)) + })) + defer server.Close() + + backend, err := chroma.New(chroma.Config{BaseURL: server.URL}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + err = backend.ForgetAll(context.Background()) + if err == nil { + t.Fatal("expected error") + } +} + +func TestNew(t *testing.T) { + _, err := chroma.New(chroma.Config{BaseURL: ""}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + _, err = chroma.New(chroma.Config{}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestNew_WithTimeout(t *testing.T) { + backend, err := chroma.New(chroma.Config{ + BaseURL: "http://localhost:8000", + Timeout: 10, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if backend == nil { + t.Fatal("expected non-nil backend") + } +} + +func TestUpsert_InvalidJSON(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{}`)) + })) + defer server.Close() + + backend, err := chroma.New(chroma.Config{BaseURL: server.URL}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Test with nil metadata (should work) + err = backend.Upsert(context.Background(), "test-id", []float32{0.1}, nil) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestSearch_EmptyResults(t *testing.T) { + mockResponse := map[string]interface{}{ + "results": []map[string]interface{}{}, + } + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/v1/collections/rony-memory/query" { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(mockResponse) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer server.Close() + + backend, err := chroma.New(chroma.Config{BaseURL: server.URL}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + results, err := backend.Search(context.Background(), []float32{0.1, 0.2}, 5) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if len(results) != 0 { + t.Fatalf("expected 0 results, got %d", len(results)) + } +} + +func TestSearch_MalformedResponse(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/v1/collections/rony-memory/query" { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`invalid json`)) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer server.Close() + + backend, err := chroma.New(chroma.Config{BaseURL: server.URL}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + _, err = backend.Search(context.Background(), []float32{0.1, 0.2}, 5) + if err == nil { + t.Fatal("expected error for malformed response") + } +} + +func TestSearch_MissingFields(t *testing.T) { + mockResponse := map[string]interface{}{ + "results": []map[string]interface{}{}, + } + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/v1/collections/rony-memory/query" { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(mockResponse) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer server.Close() + + backend, err := chroma.New(chroma.Config{BaseURL: server.URL}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + results, err := backend.Search(context.Background(), []float32{0.1, 0.2}, 5) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if len(results) != 0 { + t.Fatalf("expected 0 results, got %d", len(results)) + } +} + +func TestUpsert_ContextCanceled(t *testing.T) { + backend, err := chroma.New(chroma.Config{BaseURL: "http://localhost:8000"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err = backend.Upsert(ctx, "test-id", []float32{0.1}, nil) + if err == nil { + t.Fatal("expected error for canceled context") + } +} + +func TestSearch_ContextCanceled(t *testing.T) { + backend, err := chroma.New(chroma.Config{BaseURL: "http://localhost:8000"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err = backend.Search(ctx, []float32{0.1, 0.2}, 5) + if err == nil { + t.Fatal("expected error for canceled context") + } +} + +func TestForget_ContextCanceled(t *testing.T) { + backend, err := chroma.New(chroma.Config{BaseURL: "http://localhost:8000"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err = backend.Forget(ctx, "test-id") + if err == nil { + t.Fatal("expected error for canceled context") + } +} diff --git a/pkg/rag/embeddings/mock.go b/pkg/rag/embeddings/mock.go new file mode 100644 index 0000000..b643a86 --- /dev/null +++ b/pkg/rag/embeddings/mock.go @@ -0,0 +1,25 @@ +package embeddings + +import ( + "context" +) + +// MockEmbedder is a test double for Embedder. +type MockEmbedder struct { + EmbedFunc func(ctx context.Context, text string) ([]float32, error) + DimensionsFn func() int +} + +func (m *MockEmbedder) Embed(ctx context.Context, text string) ([]float32, error) { + if m.EmbedFunc != nil { + return m.EmbedFunc(ctx, text) + } + return []float32{0.1, 0.2, 0.3}, nil +} + +func (m *MockEmbedder) Dimensions() int { + if m.DimensionsFn != nil { + return m.DimensionsFn() + } + return 3 +} diff --git a/pkg/rag/embeddings/ollama.go b/pkg/rag/embeddings/ollama.go new file mode 100644 index 0000000..c8cfe42 --- /dev/null +++ b/pkg/rag/embeddings/ollama.go @@ -0,0 +1,86 @@ +package embeddings + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" +) + +// Config holds the settings for the Ollama embedder. +type Config struct { + BaseURL string // e.g. "http://localhost:11434" + Model string // e.g. "nomic-embed-text" +} + +// Ollama implements Embedder using Ollama's embedding API. +type Ollama struct { + baseURL string + model string + http *http.Client +} + +// NewOllama creates a new Ollama embedder. +func NewOllama(cfg Config) (*Ollama, error) { + baseURL := cfg.BaseURL + if baseURL == "" { + baseURL = "http://localhost:11434" + } + model := cfg.Model + if model == "" { + model = "nomic-embed-text" + } + + return &Ollama{ + baseURL: baseURL, + model: model, + http: http.DefaultClient, + }, nil +} + +func (e *Ollama) Embed(ctx context.Context, text string) ([]float32, error) { + endpoint := e.baseURL + "/api/embed" + + reqBody, err := json.Marshal(map[string]interface{}{ + "model": e.model, + "input": text, + }) + if err != nil { + return nil, fmt.Errorf("marshaling request: %w", err) + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, strings.NewReader(string(reqBody))) + if err != nil { + return nil, fmt.Errorf("creating request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + + 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 ollamaEmbedResponse + if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil { + return nil, fmt.Errorf("decoding response: %w", err) + } + + return apiResp.Embedding, nil +} + +func (e *Ollama) Dimensions() int { + // Default for nomic-embed-text + return 768 +} + +type ollamaEmbedResponse struct { + Embedding []float32 `json:"embedding"` +} diff --git a/pkg/rag/memory.go b/pkg/rag/memory.go new file mode 100644 index 0000000..fb5cb87 --- /dev/null +++ b/pkg/rag/memory.go @@ -0,0 +1,129 @@ +package rag + +import ( + "context" + "fmt" + "time" + + "github.com/google/uuid" +) + +// Fragment represents a piece of content stored in the RAG system. +type Fragment struct { + ID string + Content string + Vector []float32 + Metadata map[string]string + Timestamp time.Time + ProjectID string +} + +// Memory provides persistent memory and semantic search over agent content. +type Memory interface { + Add(ctx context.Context, fragment Fragment) error + Search(ctx context.Context, query string, topK int) ([]Fragment, error) + Forget(ctx context.Context, id string) error + ForgetAll(ctx context.Context) error +} + +// Config holds the settings for creating a Memory. +type Config struct { + Backend Backend + Embedder Embedder +} + +// Backend is the interface for vector database backends. +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) + Forget(ctx context.Context, id string) error + ForgetAll(ctx context.Context) error +} + +// SearchResult represents a matched fragment from a search. +type SearchResult struct { + ID string + Content string + Score float32 + Metadata map[string]string +} + +// Embedder generates embeddings for text. +type Embedder interface { + Embed(ctx context.Context, text string) ([]float32, error) + Dimensions() int +} + +// memory implements Memory using a Backend and Embedder. +type memory struct { + backend Backend + embedder Embedder +} + +// New creates a new Memory with the given config. +func New(cfg Config) (Memory, error) { + if cfg.Backend == nil { + return nil, fmt.Errorf("backend is required") + } + if cfg.Embedder == nil { + return nil, fmt.Errorf("embedder is required") + } + return &memory{ + backend: cfg.Backend, + embedder: cfg.Embedder, + }, nil +} + +func (m *memory) Add(ctx context.Context, fragment Fragment) error { + if fragment.ID == "" { + fragment.ID = uuid.New().String() + } + if fragment.Metadata == nil { + fragment.Metadata = make(map[string]string) + } + fragment.Metadata["project_id"] = fragment.ProjectID + fragment.Timestamp = time.Now() + + vector, err := m.embedder.Embed(ctx, fragment.Content) + if err != nil { + return fmt.Errorf("embedding: %w", err) + } + fragment.Vector = vector + + return m.backend.Upsert(ctx, fragment.ID, fragment.Vector, fragment.Metadata) +} + +func (m *memory) Search(ctx context.Context, query string, topK int) ([]Fragment, error) { + if topK <= 0 { + topK = 5 + } + + queryVector, err := m.embedder.Embed(ctx, query) + if err != nil { + return nil, fmt.Errorf("embedding query: %w", err) + } + + results, err := m.backend.Search(ctx, queryVector, topK) + if err != nil { + return nil, fmt.Errorf("search: %w", err) + } + + fragments := make([]Fragment, len(results)) + for i, r := range results { + fragments[i] = Fragment{ + ID: r.ID, + Content: r.Content, + Metadata: r.Metadata, + ProjectID: r.Metadata["project_id"], + } + } + return fragments, nil +} + +func (m *memory) Forget(ctx context.Context, id string) error { + return m.backend.Forget(ctx, id) +} + +func (m *memory) ForgetAll(ctx context.Context) error { + return m.backend.ForgetAll(ctx) +} diff --git a/pkg/rag/memory_test.go b/pkg/rag/memory_test.go new file mode 100644 index 0000000..e585eff --- /dev/null +++ b/pkg/rag/memory_test.go @@ -0,0 +1,230 @@ +package rag_test + +import ( + "context" + "fmt" + "testing" + + "github.com/VictorVargas/rony-llm-agent/pkg/rag" + "github.com/VictorVargas/rony-llm-agent/pkg/rag/embeddings" +) + +func TestNew_Memory(t *testing.T) { + _, err := rag.New(rag.Config{}) + if err == nil { + t.Fatal("expected error for missing backend") + } + + _, err = rag.New(rag.Config{ + Backend: &mockBackend{}, + }) + if err == nil { + t.Fatal("expected error for missing embedder") + } + + m, err := rag.New(rag.Config{ + Backend: &mockBackend{}, + Embedder: &embeddings.MockEmbedder{}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if m == nil { + t.Fatal("expected non-nil memory") + } +} + +func TestMemory_Add(t *testing.T) { + m, err := rag.New(rag.Config{ + Backend: &mockBackend{}, + Embedder: &embeddings.MockEmbedder{}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + err = m.Add(context.Background(), rag.Fragment{ + Content: "test content", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestMemory_Add_EmbeddingError(t *testing.T) { + m, err := rag.New(rag.Config{ + Backend: &mockBackend{}, + Embedder: &embeddings.MockEmbedder{EmbedFunc: func(ctx context.Context, text string) ([]float32, error) { return nil, fmt.Errorf("embed error") }}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + err = m.Add(context.Background(), rag.Fragment{ + Content: "test content", + }) + if err == nil { + t.Fatal("expected error for embedding failure") + } +} + +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) { + return []rag.SearchResult{ + {ID: "1", Content: "result 1", Score: 0.9}, + }, nil + }, + }, + Embedder: &embeddings.MockEmbedder{}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + results, err := m.Search(context.Background(), "test query", 5) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(results) != 1 { + t.Fatalf("expected 1 result, got %d", len(results)) + } + if results[0].Content != "result 1" { + t.Errorf("expected 'result 1', got %q", results[0].Content) + } +} + +func TestMemory_ForgetAll(t *testing.T) { + forgetAllCalled := false + m, err := rag.New(rag.Config{ + Backend: &mockBackend{ + forgetAllFunc: func(ctx context.Context) error { + forgetAllCalled = true + return nil + }, + }, + Embedder: &embeddings.MockEmbedder{}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + err = m.ForgetAll(context.Background()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !forgetAllCalled { + t.Fatal("expected forgetAll to be called") + } +} + +func TestMemory_Search_EmbeddingError(t *testing.T) { + m, err := rag.New(rag.Config{ + Backend: &mockBackend{}, + Embedder: &embeddings.MockEmbedder{EmbedFunc: func(ctx context.Context, text string) ([]float32, error) { + return nil, fmt.Errorf("embed error") + }}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + _, err = m.Search(context.Background(), "test query", 5) + if err == nil { + t.Fatal("expected error for embedding failure") + } +} + +// 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) + forgetAllFunc func(ctx context.Context) error +} + +func (m *mockBackend) Upsert(ctx context.Context, id string, vector []float32, metadata map[string]string) error { + if m.upsertFunc != nil { + return m.upsertFunc(ctx, id, vector, metadata) + } + return nil +} + +func (m *mockBackend) Search(ctx context.Context, queryVector []float32, topK int) ([]rag.SearchResult, error) { + if m.searchFunc != nil { + return m.searchFunc(ctx, queryVector, topK) + } + return nil, nil +} + +func (m *mockBackend) Forget(ctx context.Context, id string) error { + return nil +} + +func (m *mockBackend) ForgetAll(ctx context.Context) error { + if m.forgetAllFunc != nil { + return m.forgetAllFunc(ctx) + } + return nil +} + +func TestMemory_Add_MultipleFragments(t *testing.T) { + m, err := rag.New(rag.Config{ + Backend: &mockBackend{}, + Embedder: &embeddings.MockEmbedder{}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + for i := 0; i < 10; i++ { + err = m.Add(context.Background(), rag.Fragment{ + Content: fmt.Sprintf("test content %d", i), + }) + if err != nil { + t.Fatalf("unexpected error on iteration %d: %v", i, err) + } + } +} + +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) { + return []rag.SearchResult{}, nil + }, + }, + Embedder: &embeddings.MockEmbedder{}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + results, err := m.Search(context.Background(), "", 5) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if len(results) != 0 { + t.Fatalf("expected 0 results, got %d", len(results)) + } +} + +func TestMemory_ForgetAll_Error(t *testing.T) { + m, err := rag.New(rag.Config{ + Backend: &mockBackend{ + forgetAllFunc: func(ctx context.Context) error { + return fmt.Errorf("forget all error") + }, + }, + Embedder: &embeddings.MockEmbedder{}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + err = m.ForgetAll(context.Background()) + if err == nil { + t.Fatal("expected error for forget all failure") + } +}