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