From 1c0c86de10458cff52af2a17fd4b6262807848ea Mon Sep 17 00:00:00 2001 From: Victor Vargas Date: Wed, 15 Jul 2026 14:50:31 -0700 Subject: [PATCH] feat(rag): memory taxonomy (episodic/semantic/procedural) + episodic auto-capture Fragments now carry a memory type in metadata (legacy fragments count as procedural) with SearchByType filtering, and EpisodeCapture summarizes a finished turn with the local LLM and stores it as episodic memory, so the agent can answer "what did we do yesterday?". Includes an E2E test against a live llama.cpp server (gated) and taxonomy unit tests. Co-Authored-By: Claude Fable 5 --- pkg/rag/autocapture.go | 88 ++++++++++++++ pkg/rag/e2e_local_test.go | 99 ++++++++++++++++ pkg/rag/memory.go | 80 ++++++++++++- pkg/rag/taxonomy_test.go | 234 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 497 insertions(+), 4 deletions(-) create mode 100644 pkg/rag/autocapture.go create mode 100644 pkg/rag/e2e_local_test.go create mode 100644 pkg/rag/taxonomy_test.go diff --git a/pkg/rag/autocapture.go b/pkg/rag/autocapture.go new file mode 100644 index 0000000..a343f5e --- /dev/null +++ b/pkg/rag/autocapture.go @@ -0,0 +1,88 @@ +package rag + +import ( + "context" + "fmt" + "strings" + "time" + + "github.com/VictorVargas/rony-llm-agent/pkg/llm" +) + +// DefaultCapturePrompt is the summarization instruction EpisodeCapture uses +// when Config doesn't provide one. Consumers localize it by passing their +// own (e.g. the Rony harness passes a Spanish prompt). +const DefaultCapturePrompt = "Summarize the following exchange between a user and an AI assistant " + + "in 1-2 sentences, in the past tense, focusing on what was asked and what was done or answered. " + + "Respond ONLY with the summary, no headers or extra commentary." + +// captureMaxInputChars bounds how much of the turn is sent to the +// summarizing LLM. Auto-capture runs after every successful turn, so its +// cost must stay small and constant — the start of a long reply carries the +// gist; the tail of a truncated one rarely changes the 1-2 sentence summary. +const captureMaxInputChars = 6000 + +// EpisodeCapture implements Phase 2 §3.5 auto-capture: at the end of a +// successful turn, an LLM (ideally a small/local one — this runs on every +// turn) condenses the exchange into a 1-2 sentence event and stores it as +// episodic memory, so future sessions can recall "what happened" without the +// user ever having asked to save anything. +type EpisodeCapture struct { + Memory Memory + LLM llm.LLMClient + ProjectID string + // Prompt overrides DefaultCapturePrompt (e.g. for localization). + Prompt string +} + +// Capture summarizes one finished turn and stores it as an episodic +// fragment. toolsUsed (may be empty) is recorded in metadata so a recalled +// episode also says how the work was done. Callers typically run this in a +// background goroutine with its own timeout — a capture failure should never +// block or break the turn that just finished. +func (c *EpisodeCapture) Capture(ctx context.Context, userInput, assistantReply string, toolsUsed ...string) error { + if c == nil || c.Memory == nil || c.LLM == nil { + return fmt.Errorf("episode capture: memory and llm are required") + } + if strings.TrimSpace(userInput) == "" || strings.TrimSpace(assistantReply) == "" { + return fmt.Errorf("episode capture: nothing to capture") + } + + prompt := c.Prompt + if prompt == "" { + prompt = DefaultCapturePrompt + } + + transcript := fmt.Sprintf("User: %s\n\nAssistant: %s", userInput, assistantReply) + if len(transcript) > captureMaxInputChars { + transcript = transcript[:captureMaxInputChars] + } + + resp, err := c.LLM.Generate(ctx, llm.CompletionRequest{ + Messages: []llm.Message{ + {Role: llm.RoleSystem, Content: prompt}, + {Role: llm.RoleUser, Content: transcript}, + }, + }) + if err != nil { + return fmt.Errorf("episode capture: summarize: %w", err) + } + summary := strings.TrimSpace(resp.Content) + if summary == "" { + return fmt.Errorf("episode capture: empty summary") + } + + metadata := map[string]string{ + "date": time.Now().Format("2006-01-02"), + } + if len(toolsUsed) > 0 { + metadata["tools"] = strings.Join(toolsUsed, ",") + } + + return c.Memory.Add(ctx, Fragment{ + Content: summary, + Type: MemoryEpisodic, + ProjectID: c.ProjectID, + Metadata: metadata, + }) +} diff --git a/pkg/rag/e2e_local_test.go b/pkg/rag/e2e_local_test.go new file mode 100644 index 0000000..1f94640 --- /dev/null +++ b/pkg/rag/e2e_local_test.go @@ -0,0 +1,99 @@ +package rag_test + +import ( + "context" + "net/http" + "path/filepath" + "testing" + "time" + + llm_llamacpp "github.com/VictorVargas/rony-llm-agent/pkg/llm/providers/llamacpp" + "github.com/VictorVargas/rony-llm-agent/pkg/rag" + "github.com/VictorVargas/rony-llm-agent/pkg/rag/backends/sqlitevec" + "github.com/VictorVargas/rony-llm-agent/pkg/rag/embeddings" +) + +// TestEpisodeCapture_EndToEndLocalServer exercises the full auto-capture +// path — real LLM summarization, sqlitevec storage, taxonomy-filtered +// recall — against the llama.cpp server Rony actually uses. Skipped when no +// server is listening on localhost:8080, so it never breaks CI or offline +// runs; with the server up it's the proof the feature works for real, not +// just against stubs. +func TestEpisodeCapture_EndToEndLocalServer(t *testing.T) { + probe, err := (&http.Client{Timeout: 2 * time.Second}).Get("http://localhost:8080/v1/models") + if err != nil { + t.Skipf("no local llama.cpp server on :8080: %v", err) + } + probe.Body.Close() + + client, err := llm_llamacpp.New(llm_llamacpp.Config{}) + if err != nil { + t.Fatalf("llamacpp client: %v", err) + } + + backend, err := sqlitevec.New(filepath.Join(t.TempDir(), "e2e_memory.db")) + if err != nil { + t.Fatalf("sqlitevec: %v", err) + } + defer backend.Close() + + // Same embedder wiring the harness uses: llama.cpp embeddings when the + // server exposes them, transparent FTS5 fallback otherwise. + embedder, err := embeddings.NewLlamaCpp(embeddings.LlamaCppConfig{}) + if err != nil { + t.Fatalf("embedder: %v", err) + } + mem, err := rag.New(rag.Config{Backend: backend, Embedder: embedder}) + if err != nil { + t.Fatalf("memory: %v", err) + } + + capture := &rag.EpisodeCapture{ + Memory: mem, + LLM: client, + ProjectID: "e2e-test", + Prompt: "Resume el siguiente intercambio entre un usuario y un asistente de IA en 1-2 frases, en pasado, " + + "enfocándote en qué se pidió y qué se hizo. Responde ÚNICAMENTE con el resumen. /no_think", + } + + ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second) + defer cancel() + + userInput := "¿Puedes optimizar el cliente OpenAI del proyecto? El streaming no funciona." + reply := "Encontré que Stream() enviaba stream:false y descartaba el modelo configurado. " + + "Reescribí el cliente: ahora hace streaming real, acumula tool calls y reporta el usage. Los tests pasan." + if err := capture.Capture(ctx, userInput, reply, "read", "edit", "bash"); err != nil { + t.Fatalf("capture failed against the real server: %v", err) + } + + // Recall it back, restricted to episodic memory. + episodes, err := mem.SearchByType(ctx, "optimización del cliente OpenAI streaming", 5, rag.MemoryEpisodic) + if err != nil { + t.Fatalf("search failed: %v", err) + } + if len(episodes) == 0 { + t.Fatal("expected the captured episode to be recallable via SearchByType(episodic)") + } + ep := episodes[0] + if ep.Type != rag.MemoryEpisodic { + t.Fatalf("expected episodic type, got %q", ep.Type) + } + if ep.Metadata["tools"] != "read,edit,bash" { + t.Fatalf("expected tools metadata, got %q", ep.Metadata["tools"]) + } + if ep.Metadata["date"] == "" { + t.Fatal("expected a date on the episode") + } + t.Logf("captured episode: %s", ep.Content) + + // A procedural-only search must NOT return the episode. + procs, err := mem.SearchByType(ctx, "optimización del cliente OpenAI streaming", 5, rag.MemoryProcedural) + if err != nil { + t.Fatalf("procedural search failed: %v", err) + } + for _, p := range procs { + if p.ID == ep.ID { + t.Fatal("episode leaked into a procedural-only search") + } + } +} diff --git a/pkg/rag/memory.go b/pkg/rag/memory.go index 0e0b7d1..a0cb924 100644 --- a/pkg/rag/memory.go +++ b/pkg/rag/memory.go @@ -8,10 +8,36 @@ import ( "github.com/google/uuid" ) +// MemoryType classifies a fragment within the three-tier taxonomy from the +// Phase 2 spec (docs/phase2.md §3.1). Working memory (the current session's +// messages) lives in the consuming product's own state, not here. +type MemoryType string + +const ( + // MemoryEpisodic records past events: "what happened / what I did on + // ". Typically auto-captured at the end of successful turns (see + // EpisodeCapture) rather than saved deliberately. + MemoryEpisodic MemoryType = "episodic" + // MemorySemantic records consolidated knowledge and facts: "how the + // architecture works", "the API returns X". Curated — saved when + // something is worth knowing independent of when it was learned. + MemorySemantic MemoryType = "semantic" + // MemoryProcedural records how to do things: workflows, procedures, + // user preferences about process. This is also what every fragment + // saved before the taxonomy existed is treated as — the pre-taxonomy + // tools (save_process et al.) only ever stored procedures. + MemoryProcedural MemoryType = "procedural" +) + +// metaTypeKey is the metadata key the fragment's MemoryType round-trips +// through, so backends need no schema change to support the taxonomy. +const metaTypeKey = "memory_type" + // Fragment represents a piece of content stored in the RAG system. type Fragment struct { ID string Content string + Type MemoryType // defaults to MemoryProcedural when empty (pre-taxonomy compatibility) Vector []float32 Metadata map[string]string Timestamp time.Time @@ -22,6 +48,10 @@ type Fragment struct { type Memory interface { Add(ctx context.Context, fragment Fragment) error Search(ctx context.Context, query string, topK int) ([]Fragment, error) + // SearchByType is Search restricted to the given memory types. No types + // means no restriction (same as Search). Fragments stored before the + // taxonomy existed match MemoryProcedural. + SearchByType(ctx context.Context, query string, topK int, types ...MemoryType) ([]Fragment, error) Forget(ctx context.Context, id string) error ForgetAll(ctx context.Context) error } @@ -85,6 +115,10 @@ func (m *memory) Add(ctx context.Context, fragment Fragment) error { if fragment.Metadata == nil { fragment.Metadata = make(map[string]string) } + if fragment.Type == "" { + fragment.Type = MemoryProcedural + } + fragment.Metadata[metaTypeKey] = string(fragment.Type) fragment.Metadata["project_id"] = fragment.ProjectID fragment.Timestamp = time.Now() @@ -101,9 +135,34 @@ func (m *memory) Add(ctx context.Context, fragment Fragment) error { } func (m *memory) Search(ctx context.Context, query string, topK int) ([]Fragment, error) { + return m.SearchByType(ctx, query, topK) +} + +// typeOf resolves a stored result's memory type; fragments saved before the +// taxonomy existed carry no memory_type metadata and were all procedures. +func typeOf(metadata map[string]string) MemoryType { + if t := MemoryType(metadata[metaTypeKey]); t != "" { + return t + } + return MemoryProcedural +} + +// typeFilterOverfetch is how many times topK gets requested from the backend +// when SearchByType has to post-filter by memory type: the Backend interface +// has no type predicate (deliberately — backends stay schema-agnostic), so +// filtering happens here and the extra headroom keeps a type-restricted +// search from coming back near-empty just because the top raw matches +// happened to be of other types. +const typeFilterOverfetch = 4 + +func (m *memory) SearchByType(ctx context.Context, query string, topK int, types ...MemoryType) ([]Fragment, error) { if topK <= 0 { topK = 5 } + fetchK := topK + if len(types) > 0 { + fetchK = topK * typeFilterOverfetch + } // Same fallback as Add: if embedding the query fails, search proceeds // with no vector so the backend can fall back to lexical matching. @@ -112,18 +171,31 @@ func (m *memory) Search(ctx context.Context, query string, topK int) ([]Fragment queryVector = nil } - results, err := m.backend.Search(ctx, query, queryVector, topK) + results, err := m.backend.Search(ctx, query, queryVector, fetchK) if err != nil { return nil, fmt.Errorf("search: %w", err) } - fragments := make([]Fragment, len(results)) - for i, r := range results { - fragments[i] = Fragment{ + wanted := make(map[MemoryType]bool, len(types)) + for _, t := range types { + wanted[t] = true + } + + fragments := make([]Fragment, 0, topK) + for _, r := range results { + fragType := typeOf(r.Metadata) + if len(wanted) > 0 && !wanted[fragType] { + continue + } + fragments = append(fragments, Fragment{ ID: r.ID, Content: r.Content, + Type: fragType, Metadata: r.Metadata, ProjectID: r.Metadata["project_id"], + }) + if len(fragments) == topK { + break } } return fragments, nil diff --git a/pkg/rag/taxonomy_test.go b/pkg/rag/taxonomy_test.go new file mode 100644 index 0000000..d3c3351 --- /dev/null +++ b/pkg/rag/taxonomy_test.go @@ -0,0 +1,234 @@ +package rag_test + +import ( + "context" + "fmt" + "iter" + "strings" + "testing" + + "github.com/VictorVargas/rony-llm-agent/pkg/llm" + "github.com/VictorVargas/rony-llm-agent/pkg/rag" + "github.com/VictorVargas/rony-llm-agent/pkg/rag/embeddings" +) + +func TestMemory_Add_StoresMemoryTypeInMetadata(t *testing.T) { + var gotMeta map[string]string + m, err := rag.New(rag.Config{ + Backend: &mockBackend{ + upsertFunc: func(_ context.Context, _ string, _ []float32, _ string, metadata map[string]string) error { + gotMeta = metadata + return nil + }, + }, + Embedder: &embeddings.MockEmbedder{}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if err := m.Add(context.Background(), rag.Fragment{Content: "an event", Type: rag.MemoryEpisodic}); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if gotMeta["memory_type"] != "episodic" { + t.Fatalf("expected memory_type=episodic in metadata, got %q", gotMeta["memory_type"]) + } +} + +func TestMemory_Add_DefaultsToProcedural(t *testing.T) { + var gotMeta map[string]string + m, err := rag.New(rag.Config{ + Backend: &mockBackend{ + upsertFunc: func(_ context.Context, _ string, _ []float32, _ string, metadata map[string]string) error { + gotMeta = metadata + return nil + }, + }, + Embedder: &embeddings.MockEmbedder{}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if err := m.Add(context.Background(), rag.Fragment{Content: "a process"}); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if gotMeta["memory_type"] != "procedural" { + t.Fatalf("expected untyped fragments to default to procedural, got %q", gotMeta["memory_type"]) + } +} + +func TestMemory_SearchByType_FiltersAndTreatsLegacyAsProcedural(t *testing.T) { + backendResults := []rag.SearchResult{ + {ID: "1", Content: "episode", Metadata: map[string]string{"memory_type": "episodic"}}, + {ID: "2", Content: "fact", Metadata: map[string]string{"memory_type": "semantic"}}, + {ID: "3", Content: "legacy process", Metadata: map[string]string{}}, // pre-taxonomy fragment + {ID: "4", Content: "typed process", Metadata: map[string]string{"memory_type": "procedural"}}, + } + m, err := rag.New(rag.Config{ + Backend: &mockBackend{ + searchFunc: func(_ context.Context, _ string, _ []float32, _ int) ([]rag.SearchResult, error) { + return backendResults, nil + }, + }, + Embedder: &embeddings.MockEmbedder{}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + got, err := m.SearchByType(context.Background(), "q", 10, rag.MemoryProcedural) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(got) != 2 || got[0].ID != "3" || got[1].ID != "4" { + t.Fatalf("expected legacy + typed procedural fragments, got %+v", got) + } + if got[0].Type != rag.MemoryProcedural { + t.Fatalf("expected legacy fragment to surface as procedural, got %q", got[0].Type) + } + + episodes, err := m.SearchByType(context.Background(), "q", 10, rag.MemoryEpisodic) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(episodes) != 1 || episodes[0].ID != "1" { + t.Fatalf("expected only the episodic fragment, got %+v", episodes) + } + + all, err := m.SearchByType(context.Background(), "q", 10) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(all) != 4 { + t.Fatalf("expected no type restriction to return everything, got %d", len(all)) + } +} + +func TestMemory_SearchByType_OverfetchesWhenFiltering(t *testing.T) { + var gotTopK int + m, err := rag.New(rag.Config{ + Backend: &mockBackend{ + searchFunc: func(_ context.Context, _ string, _ []float32, topK int) ([]rag.SearchResult, error) { + gotTopK = topK + return nil, nil + }, + }, + Embedder: &embeddings.MockEmbedder{}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if _, err := m.SearchByType(context.Background(), "q", 5, rag.MemoryEpisodic); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if gotTopK <= 5 { + t.Fatalf("expected the backend to be asked for more than topK candidates when filtering, got %d", gotTopK) + } + + if _, err := m.Search(context.Background(), "q", 5); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if gotTopK != 5 { + t.Fatalf("expected unfiltered search to request exactly topK, got %d", gotTopK) + } +} + +// captureLLM is a minimal llm.LLMClient stub for capture tests. +type captureLLM struct { + response string + err error + gotUser string +} + +func (c *captureLLM) Generate(_ context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + for _, m := range req.Messages { + if m.Role == llm.RoleUser { + c.gotUser = m.Content + } + } + if c.err != nil { + return llm.CompletionResponse{}, c.err + } + return llm.CompletionResponse{Content: c.response}, nil +} + +func (c *captureLLM) Stream(_ context.Context, _ llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] { + return func(func(llm.StreamChunk, error) bool) {} +} +func (c *captureLLM) Name() string { return "capture-stub" } +func (c *captureLLM) Capabilities() llm.ProviderCapabilities { return llm.ProviderCapabilities{} } + +func TestEpisodeCapture_SavesEpisodicSummary(t *testing.T) { + var saved rag.Fragment + backend := &mockBackend{ + upsertFunc: func(_ context.Context, id string, _ []float32, content string, metadata map[string]string) error { + saved = rag.Fragment{ID: id, Content: content, Metadata: metadata} + return nil + }, + } + mem, err := rag.New(rag.Config{Backend: backend, Embedder: &embeddings.MockEmbedder{}}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + stub := &captureLLM{response: " The user asked how to deploy and the assistant explained the release steps. "} + cap := &rag.EpisodeCapture{Memory: mem, LLM: stub, ProjectID: "proj1"} + + err = cap.Capture(context.Background(), "how do I deploy?", "You run make release...", "read", "bash") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if saved.Content != "The user asked how to deploy and the assistant explained the release steps." { + t.Fatalf("expected trimmed summary as content, got %q", saved.Content) + } + if saved.Metadata["memory_type"] != "episodic" { + t.Fatalf("expected episodic type, got %q", saved.Metadata["memory_type"]) + } + if saved.Metadata["tools"] != "read,bash" { + t.Fatalf("expected tools metadata, got %q", saved.Metadata["tools"]) + } + if saved.Metadata["project_id"] != "proj1" { + t.Fatalf("expected project_id metadata, got %q", saved.Metadata["project_id"]) + } + if !strings.Contains(stub.gotUser, "how do I deploy?") { + t.Fatalf("expected the turn transcript to reach the LLM, got %q", stub.gotUser) + } +} + +func TestEpisodeCapture_SkipsEmptyTurnsAndFailures(t *testing.T) { + upserts := 0 + backend := &mockBackend{ + upsertFunc: func(_ context.Context, _ string, _ []float32, _ string, _ map[string]string) error { + upserts++ + return nil + }, + } + mem, err := rag.New(rag.Config{Backend: backend, Embedder: &embeddings.MockEmbedder{}}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + cap := &rag.EpisodeCapture{Memory: mem, LLM: &captureLLM{response: "summary"}, ProjectID: "p"} + if err := cap.Capture(context.Background(), "", "reply"); err == nil { + t.Fatal("expected error for empty user input") + } + if err := cap.Capture(context.Background(), "input", " "); err == nil { + t.Fatal("expected error for empty assistant reply") + } + + failing := &rag.EpisodeCapture{Memory: mem, LLM: &captureLLM{err: fmt.Errorf("llm down")}, ProjectID: "p"} + if err := failing.Capture(context.Background(), "input", "reply"); err == nil { + t.Fatal("expected error when the LLM fails") + } + empty := &rag.EpisodeCapture{Memory: mem, LLM: &captureLLM{response: " "}, ProjectID: "p"} + if err := empty.Capture(context.Background(), "input", "reply"); err == nil { + t.Fatal("expected error for an empty summary") + } + + if upserts != 0 { + t.Fatalf("expected nothing to be saved on failures, got %d upserts", upserts) + } +}