feat(rag): add Retrieval-Augmented Memory with ChromaDB backend and embedding support
This commit is contained in:
parent
3f2dc5ccc0
commit
9f33d6df93
6 changed files with 1012 additions and 0 deletions
231
pkg/rag/backends/chroma/chroma.go
Normal file
231
pkg/rag/backends/chroma/chroma.go
Normal file
|
|
@ -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"`
|
||||||
|
}
|
||||||
311
pkg/rag/backends/chroma/chroma_test.go
Normal file
311
pkg/rag/backends/chroma/chroma_test.go
Normal file
|
|
@ -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")
|
||||||
|
}
|
||||||
|
}
|
||||||
25
pkg/rag/embeddings/mock.go
Normal file
25
pkg/rag/embeddings/mock.go
Normal file
|
|
@ -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
|
||||||
|
}
|
||||||
86
pkg/rag/embeddings/ollama.go
Normal file
86
pkg/rag/embeddings/ollama.go
Normal file
|
|
@ -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"`
|
||||||
|
}
|
||||||
129
pkg/rag/memory.go
Normal file
129
pkg/rag/memory.go
Normal file
|
|
@ -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)
|
||||||
|
}
|
||||||
230
pkg/rag/memory_test.go
Normal file
230
pkg/rag/memory_test.go
Normal file
|
|
@ -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")
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Reference in a new issue