From 3d5fe5bc4d2c65bdf12060c09962eb362eba2c57 Mon Sep 17 00:00:00 2001 From: pmorgan Date: Sun, 22 Feb 2026 11:58:10 -0700 Subject: [PATCH] Add semantic memory with chromem-go vector store Introduce remember/recall tools backed by chromem-go and Ollama embeddings for persistent semantic memory. Fix ContextWindow to use config value (128k) instead of MaxTokens (8k) to prevent premature summarization. Add auto-extraction of key facts during summarization and integrate streaming into the agent loop. Co-Authored-By: Claude Opus 4.6 --- go.mod | 1 + go.sum | 2 + pkg/agent/instance.go | 7 +- pkg/agent/loop.go | 207 +++++++++++++++++++++++++---- pkg/agent/loop_test.go | 157 ++++++++++++++++++++++ pkg/config/config.go | 9 ++ pkg/config/defaults.go | 6 + pkg/memory/ollama.go | 67 ++++++++++ pkg/memory/ollama_test.go | 113 ++++++++++++++++ pkg/memory/store.go | 206 +++++++++++++++++++++++++++++ pkg/memory/store_test.go | 262 +++++++++++++++++++++++++++++++++++++ pkg/tools/recall.go | 83 ++++++++++++ pkg/tools/recall_test.go | 121 +++++++++++++++++ pkg/tools/remember.go | 86 ++++++++++++ pkg/tools/remember_test.go | 115 ++++++++++++++++ 15 files changed, 1416 insertions(+), 26 deletions(-) create mode 100644 pkg/memory/ollama.go create mode 100644 pkg/memory/ollama_test.go create mode 100644 pkg/memory/store.go create mode 100644 pkg/memory/store_test.go create mode 100644 pkg/tools/recall.go create mode 100644 pkg/tools/recall_test.go create mode 100644 pkg/tools/remember.go create mode 100644 pkg/tools/remember_test.go diff --git a/go.mod b/go.mod index 1f88639c8..24588379c 100644 --- a/go.mod +++ b/go.mod @@ -22,6 +22,7 @@ require ( require ( github.com/davecgh/go-spew v1.1.1 // indirect + github.com/philippgille/chromem-go v0.7.0 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index 0e95bf5cd..181367429 100644 --- a/go.sum +++ b/go.sum @@ -102,6 +102,8 @@ github.com/open-dingtalk/dingtalk-stream-sdk-go v0.9.1 h1:Lb/Uzkiw2Ugt2Xf03J5wmv github.com/open-dingtalk/dingtalk-stream-sdk-go v0.9.1/go.mod h1:ln3IqPYYocZbYvl9TAOrG/cxGR9xcn4pnZRLdCTEGEU= github.com/openai/openai-go/v3 v3.22.0 h1:6MEoNoV8sbjOVmXdvhmuX3BjVbVdcExbVyGixiyJ8ys= github.com/openai/openai-go/v3 v3.22.0/go.mod h1:cdufnVK14cWcT9qA1rRtrXx4FTRsgbDPW7Ia7SS5cZo= +github.com/philippgille/chromem-go v0.7.0 h1:4jfvfyKymjKNfGxBUhHUcj1kp7B17NL/I1P+vGh1RvY= +github.com/philippgille/chromem-go v0.7.0/go.mod h1:hTd+wGEm/fFPQl7ilfCwQXkgEUxceYh86iIdoKMolPo= github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index dfbef9fbc..4f8162bbe 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -95,6 +95,11 @@ func NewAgentInstance( } candidates := providers.ResolveCandidates(modelCfg, defaults.Provider) + contextWindow := defaults.ContextWindow + if contextWindow == 0 { + contextWindow = 128000 + } + return &AgentInstance{ ID: agentID, Name: agentName, @@ -104,7 +109,7 @@ func NewAgentInstance( MaxIterations: maxIter, MaxTokens: maxTokens, Temperature: temperature, - ContextWindow: maxTokens, + ContextWindow: contextWindow, Provider: provider, Sessions: sessionsManager, ContextBuilder: contextBuilder, diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index b36f4a0c4..ae72a7a5b 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -16,11 +16,14 @@ import ( "time" "unicode/utf8" + "path/filepath" + "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/channels" "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/constants" "github.com/sipeed/picoclaw/pkg/logger" + "github.com/sipeed/picoclaw/pkg/memory" "github.com/sipeed/picoclaw/pkg/providers" "github.com/sipeed/picoclaw/pkg/routing" "github.com/sipeed/picoclaw/pkg/skills" @@ -38,6 +41,7 @@ type AgentLoop struct { summarizing sync.Map fallback *providers.FallbackChain channelManager *channels.Manager + semanticStores map[string]memory.Store } // processOptions configures how a message is processed @@ -55,8 +59,8 @@ type processOptions struct { func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers.LLMProvider) *AgentLoop { registry := NewAgentRegistry(cfg, provider) - // Register shared tools to all agents - registerSharedTools(cfg, msgBus, registry, provider) + // Register shared tools to all agents (including memory tools) + stores := registerSharedTools(cfg, msgBus, registry, provider) // Set up shared fallback chain cooldown := providers.NewCooldownTracker() @@ -70,22 +74,26 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers } return &AgentLoop{ - bus: msgBus, - cfg: cfg, - registry: registry, - state: stateManager, - summarizing: sync.Map{}, - fallback: fallbackChain, + bus: msgBus, + cfg: cfg, + registry: registry, + state: stateManager, + summarizing: sync.Map{}, + fallback: fallbackChain, + semanticStores: stores, } } -// registerSharedTools registers tools that are shared across all agents (web, message, spawn). +// registerSharedTools registers tools that are shared across all agents (web, message, spawn, memory). +// Returns a map of agent ID to semantic memory store for auto-extraction. func registerSharedTools( cfg *config.Config, msgBus *bus.MessageBus, registry *AgentRegistry, provider providers.LLMProvider, -) { +) map[string]memory.Store { + stores := make(map[string]memory.Store) + for _, agentID := range registry.ListAgentIDs() { agent, ok := registry.GetAgent(agentID) if !ok { @@ -145,9 +153,30 @@ func registerSharedTools( }) agent.Tools.Register(spawnTool) + // Semantic memory tools (graceful degradation if Ollama unavailable) + vectorDir := filepath.Join(agent.Workspace, "memory", "vectors") + store, err := memory.NewSemanticStore( + vectorDir, + cfg.Tools.Memory.OllamaURL, + cfg.Tools.Memory.EmbeddingModel, + ) + if err != nil { + logger.WarnCF("agent", "Semantic memory unavailable", + map[string]any{ + "agent_id": agentID, + "error": err.Error(), + }) + } else if store.IsAvailable() { + agent.Tools.Register(tools.NewRememberTool(store)) + agent.Tools.Register(tools.NewRecallTool(store)) + stores[agentID] = store + } + // Update context builder with the complete tools registry agent.ContextBuilder.SetToolsRegistry(agent.Tools) } + + return stores } func (al *AgentLoop) Run(ctx context.Context) error { @@ -514,14 +543,16 @@ func (al *AgentLoop) runLLMIteration( var response *providers.LLMResponse var err error + llmOpts := map[string]any{ + "max_tokens": agent.MaxTokens, + "temperature": agent.Temperature, + } + callLLM := func() (*providers.LLMResponse, error) { if len(agent.Candidates) > 1 && al.fallback != nil { fbResult, fbErr := al.fallback.Execute(ctx, agent.Candidates, func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) { - return agent.Provider.Chat(ctx, messages, providerToolDefs, model, map[string]any{ - "max_tokens": agent.MaxTokens, - "temperature": agent.Temperature, - }) + return agent.Provider.Chat(ctx, messages, providerToolDefs, model, llmOpts) }, ) if fbErr != nil { @@ -534,10 +565,17 @@ func (al *AgentLoop) runLLMIteration( } return fbResult.Response, nil } - return agent.Provider.Chat(ctx, messages, providerToolDefs, agent.Model, map[string]any{ - "max_tokens": agent.MaxTokens, - "temperature": agent.Temperature, - }) + // Use streaming if the provider supports it + if sp, ok := agent.Provider.(providers.StreamingProvider); ok { + return sp.ChatStream(ctx, messages, providerToolDefs, agent.Model, llmOpts, + func(delta string) { + // Streaming delta received — currently logged only. + // Future: push partial content to user via message bus. + logger.DebugCF("agent", "Stream delta", + map[string]any{"len": len(delta)}) + }) + } + return agent.Provider.Chat(ctx, messages, providerToolDefs, agent.Model, llmOpts) } // Retry loop for context/token errors @@ -783,8 +821,9 @@ func (al *AgentLoop) forceCompression(agent *AgentInstance, sessionKey string) { return } - // Helper to find the mid-point of the conversation - mid := len(conversation) / 2 + // Find the mid-point of the conversation, adjusted to avoid splitting + // tool_use/tool_result pairs. + mid := safeSplitIndex(conversation, len(conversation)/2) // New history structure: // 1. System Prompt (with compression note appended) @@ -820,6 +859,33 @@ func (al *AgentLoop) forceCompression(agent *AgentInstance, sessionKey string) { }) } +// safeSplitIndex finds a split point in messages at or after targetIdx that +// does not break a tool_use/tool_result pair. An assistant message with +// ToolCalls must be kept together with the tool-result messages that +// immediately follow it. The returned index is the first message to KEEP +// (i.e. messages[:idx] are dropped, messages[idx:] are kept). +func safeSplitIndex(messages []providers.Message, targetIdx int) int { + if targetIdx >= len(messages) { + return len(messages) + } + idx := targetIdx + // If we landed on a "tool" message, walk forward past all consecutive + // tool results so we don't orphan them from their assistant message. + for idx < len(messages) && messages[idx].Role == "tool" { + idx++ + } + // If we ended up past the end, back up to just before the tool group. + if idx >= len(messages) { + // Walk backwards from targetIdx to find a safe point. + idx = targetIdx + for idx > 0 && messages[idx].Role == "tool" { + idx-- + } + // idx now points to the assistant message; keep it and its tool results. + } + return idx +} + // GetStartupInfo returns information about loaded tools and skills for logging. func (al *AgentLoop) GetStartupInfo() map[string]any { info := make(map[string]any) @@ -907,12 +973,17 @@ func (al *AgentLoop) summarizeSession(agent *AgentInstance, sessionKey string) { history := agent.Sessions.GetHistory(sessionKey) summary := agent.Sessions.GetSummary(sessionKey) - // Keep last 4 messages for continuity - if len(history) <= 4 { + // Keep last few messages for continuity, adjusting the split point so we + // never orphan a tool_use from its tool_result(s). + if len(history) <= 6 { return } - toSummarize := history[:len(history)-4] + keepIdx := safeSplitIndex(history, len(history)-6) + if keepIdx <= 0 { + return + } + toSummarize := history[:keepIdx] // Oversized Message Guard maxMessageTokens := agent.ContextWindow / 2 @@ -920,7 +991,7 @@ func (al *AgentLoop) summarizeSession(agent *AgentInstance, sessionKey string) { omitted := false for _, m := range toSummarize { - if m.Role != "user" && m.Role != "assistant" { + if m.Role != "user" && m.Role != "assistant" && m.Role != "tool" { continue } msgTokens := len(m.Content) / 2 @@ -975,8 +1046,94 @@ func (al *AgentLoop) summarizeSession(agent *AgentInstance, sessionKey string) { if finalSummary != "" { agent.Sessions.SetSummary(sessionKey, finalSummary) - agent.Sessions.TruncateHistory(sessionKey, 4) + agent.Sessions.TruncateHistory(sessionKey, len(history)-keepIdx) agent.Sessions.Save(sessionKey) + + // Auto-extract key facts from the summarized messages into semantic memory + if al.cfg.Tools.Memory.AutoExtract { + if store, ok := al.semanticStores[agent.ID]; ok { + go al.extractMemories(ctx, agent, store, toSummarize) + } + } + } +} + +// extractMemories uses the LLM to extract key facts from conversation messages +// and stores them in the semantic memory store for later recall. +func (al *AgentLoop) extractMemories( + ctx context.Context, + agent *AgentInstance, + store memory.Store, + messages []providers.Message, +) { + if !store.IsAvailable() { + return + } + + // Build a prompt asking the LLM to extract memorable facts + var sb strings.Builder + sb.WriteString("Extract key facts, user preferences, decisions, and important context from this conversation.\n") + sb.WriteString("Output one fact per line in the format: [category] fact\n") + sb.WriteString("Categories: preference, fact, decision, context\n") + sb.WriteString("Only extract genuinely important or reusable information. Skip trivial details.\n\n") + for _, m := range messages { + if m.Role == "user" || m.Role == "assistant" { + fmt.Fprintf(&sb, "%s: %s\n", m.Role, utils.Truncate(m.Content, 500)) + } + } + + extractCtx, cancel := context.WithTimeout(ctx, 30*time.Second) + defer cancel() + + resp, err := agent.Provider.Chat( + extractCtx, + []providers.Message{{Role: "user", Content: sb.String()}}, + nil, + agent.Model, + map[string]any{ + "max_tokens": 512, + "temperature": 0.2, + }, + ) + if err != nil { + logger.WarnCF("agent", "Memory extraction failed", + map[string]any{"error": err.Error()}) + return + } + + // Parse "[category] fact" lines + stored := 0 + for _, line := range strings.Split(resp.Content, "\n") { + line = strings.TrimSpace(line) + if line == "" || !strings.HasPrefix(line, "[") { + continue + } + closeBracket := strings.Index(line, "]") + if closeBracket < 2 { + continue + } + category := strings.TrimSpace(line[1:closeBracket]) + content := strings.TrimSpace(line[closeBracket+1:]) + if content == "" { + continue + } + + entry := memory.MemoryEntry{ + Content: content, + Category: category, + Source: "auto-extract", + } + if err := store.Remember(extractCtx, entry); err != nil { + logger.WarnCF("agent", "Failed to store extracted memory", + map[string]any{"error": err.Error()}) + continue + } + stored++ + } + + if stored > 0 { + logger.InfoCF("agent", "Auto-extracted memories", + map[string]any{"count": stored, "agent_id": agent.ID}) } } diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index 4414398b1..53f306d18 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -521,6 +521,163 @@ func TestToolResult_UserFacingToolDoesSendMessage(t *testing.T) { } } +// TestSafeSplitIndex verifies that safeSplitIndex never breaks tool_use/tool_result pairs. +func TestSafeSplitIndex(t *testing.T) { + tests := []struct { + name string + messages []providers.Message + targetIdx int + wantIdx int + }{ + { + name: "no tool messages, split at target", + messages: []providers.Message{ + {Role: "user", Content: "hi"}, + {Role: "assistant", Content: "hello"}, + {Role: "user", Content: "bye"}, + }, + targetIdx: 1, + wantIdx: 1, + }, + { + name: "target lands on tool result, skip forward", + messages: []providers.Message{ + {Role: "user", Content: "hi"}, + {Role: "assistant", Content: "", ToolCalls: []providers.ToolCall{{ID: "1", Name: "exec"}}}, + {Role: "tool", Content: "result1"}, + {Role: "user", Content: "thanks"}, + }, + targetIdx: 2, // lands on tool result + wantIdx: 3, // skip past it + }, + { + name: "target lands on first of two tool results, skip both", + messages: []providers.Message{ + {Role: "user", Content: "hi"}, + {Role: "assistant", Content: "", ToolCalls: []providers.ToolCall{{ID: "1"}, {ID: "2"}}}, + {Role: "tool", Content: "result1"}, + {Role: "tool", Content: "result2"}, + {Role: "user", Content: "thanks"}, + }, + targetIdx: 2, + wantIdx: 4, + }, + { + name: "target on non-tool message, no adjustment", + messages: []providers.Message{ + {Role: "user", Content: "hi"}, + {Role: "assistant", Content: "hello"}, + {Role: "user", Content: "bye"}, + {Role: "assistant", Content: "goodbye"}, + }, + targetIdx: 2, + wantIdx: 2, + }, + { + name: "target on tool at end, back up to assistant", + messages: []providers.Message{ + {Role: "user", Content: "hi"}, + {Role: "assistant", Content: "", ToolCalls: []providers.ToolCall{{ID: "1"}}}, + {Role: "tool", Content: "result"}, + }, + targetIdx: 2, // tool at end + wantIdx: 1, // back up to include assistant+tool pair + }, + { + name: "target beyond messages length", + messages: []providers.Message{ + {Role: "user", Content: "hi"}, + }, + targetIdx: 5, + wantIdx: 1, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := safeSplitIndex(tt.messages, tt.targetIdx) + if got != tt.wantIdx { + t.Errorf("safeSplitIndex(target=%d) = %d, want %d", tt.targetIdx, got, tt.wantIdx) + } + }) + } +} + +// TestSanitizeHistoryForProvider_MultipleToolResults verifies that consecutive +// tool results from parallel tool calls are not dropped. +func TestSanitizeHistoryForProvider_MultipleToolResults(t *testing.T) { + tests := []struct { + name string + history []providers.Message + wantLen int + wantDesc string + }{ + { + name: "single tool result kept", + history: []providers.Message{ + {Role: "user", Content: "hi"}, + {Role: "assistant", Content: "", ToolCalls: []providers.ToolCall{{ID: "call_1", Name: "exec"}}}, + {Role: "tool", Content: "result1", ToolCallID: "call_1"}, + {Role: "assistant", Content: "done"}, + }, + wantLen: 4, + wantDesc: "all messages kept", + }, + { + name: "three consecutive tool results all kept", + history: []providers.Message{ + {Role: "user", Content: "check endpoints"}, + {Role: "assistant", Content: "", ToolCalls: []providers.ToolCall{ + {ID: "call_1", Name: "web_fetch"}, + {ID: "call_2", Name: "web_fetch"}, + {ID: "call_3", Name: "web_fetch"}, + }}, + {Role: "tool", Content: `{"status":"ok"}`, ToolCallID: "call_1"}, + {Role: "tool", Content: `{"status":"ok"}`, ToolCallID: "call_2"}, + {Role: "tool", Content: `{"status":"ok"}`, ToolCallID: "call_3"}, + {Role: "assistant", Content: "All endpoints healthy"}, + }, + wantLen: 6, + wantDesc: "all 3 tool results kept", + }, + { + name: "orphaned leading tool message dropped", + history: []providers.Message{ + {Role: "tool", Content: "orphan", ToolCallID: "call_0"}, + {Role: "user", Content: "hi"}, + {Role: "assistant", Content: "hello"}, + }, + wantLen: 2, + wantDesc: "orphaned tool message dropped, user+assistant kept", + }, + { + name: "tool result after non-toolcall assistant dropped", + history: []providers.Message{ + {Role: "user", Content: "hi"}, + {Role: "assistant", Content: "hello"}, + {Role: "tool", Content: "orphan", ToolCallID: "call_0"}, + {Role: "user", Content: "bye"}, + }, + wantLen: 3, + wantDesc: "orphaned tool result dropped", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := sanitizeHistoryForProvider(tt.history) + if len(got) != tt.wantLen { + roles := make([]string, len(got)) + for i, m := range got { + roles[i] = m.Role + } + t.Errorf("sanitizeHistoryForProvider() len = %d, want %d (%s); roles = %v", + len(got), tt.wantLen, tt.wantDesc, roles) + } + }) + } +} + // failFirstMockProvider fails on the first N calls with a specific error type failFirstMockProvider struct { failures int diff --git a/pkg/config/config.go b/pkg/config/config.go index 20556011a..66108e749 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -175,6 +175,7 @@ type AgentDefaults struct { ImageModel string `json:"image_model,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_IMAGE_MODEL"` ImageModelFallbacks []string `json:"image_model_fallbacks,omitempty"` MaxTokens int `json:"max_tokens" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOKENS"` + ContextWindow int `json:"context_window,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_CONTEXT_WINDOW"` Temperature *float64 `json:"temperature,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_TEMPERATURE"` MaxToolIterations int `json:"max_tool_iterations" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOOL_ITERATIONS"` } @@ -449,6 +450,14 @@ type ToolsConfig struct { Cron CronToolsConfig `json:"cron"` Exec ExecConfig `json:"exec"` Skills SkillsToolsConfig `json:"skills"` + Memory MemoryConfig `json:"memory"` +} + +// MemoryConfig controls the semantic memory system backed by vector embeddings. +type MemoryConfig struct { + OllamaURL string `json:"ollama_url" env:"PICOCLAW_TOOLS_MEMORY_OLLAMA_URL"` + EmbeddingModel string `json:"embedding_model" env:"PICOCLAW_TOOLS_MEMORY_EMBEDDING_MODEL"` + AutoExtract bool `json:"auto_extract" env:"PICOCLAW_TOOLS_MEMORY_AUTO_EXTRACT"` } type SkillsToolsConfig struct { diff --git a/pkg/config/defaults.go b/pkg/config/defaults.go index 7654326e7..b5584d015 100644 --- a/pkg/config/defaults.go +++ b/pkg/config/defaults.go @@ -15,6 +15,7 @@ func DefaultConfig() *Config { Provider: "", Model: "glm-4.7", MaxTokens: 8192, + ContextWindow: 128000, Temperature: nil, // nil means use provider default MaxToolIterations: 20, }, @@ -303,6 +304,11 @@ func DefaultConfig() *Config { TTLSeconds: 300, }, }, + Memory: MemoryConfig{ + OllamaURL: "http://localhost:11434", + EmbeddingModel: "nomic-embed-text", + AutoExtract: true, + }, }, Heartbeat: HeartbeatConfig{ Enabled: true, diff --git a/pkg/memory/ollama.go b/pkg/memory/ollama.go new file mode 100644 index 000000000..5ce4d0685 --- /dev/null +++ b/pkg/memory/ollama.go @@ -0,0 +1,67 @@ +// ABOUTME: Ollama connectivity for embedding generation. +// ABOUTME: Provides health checks and constants for the Ollama embedding service. +package memory + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "time" +) + +const ( + DefaultOllamaURL = "http://localhost:11434" + DefaultEmbeddingModel = "nomic-embed-text" +) + +// ollamaTagsResponse is the response from GET /api/tags. +type ollamaTagsResponse struct { + Models []ollamaModel `json:"models"` +} + +type ollamaModel struct { + Name string `json:"name"` +} + +// CheckOllamaAvailable checks whether Ollama is reachable and has the +// requested embedding model installed. Returns nil on success. +func CheckOllamaAvailable(ctx context.Context, ollamaURL, model string) error { + if ollamaURL == "" { + ollamaURL = DefaultOllamaURL + } + if model == "" { + model = DefaultEmbeddingModel + } + + client := &http.Client{Timeout: 5 * time.Second} + req, err := http.NewRequestWithContext(ctx, http.MethodGet, ollamaURL+"/api/tags", nil) + if err != nil { + return fmt.Errorf("creating request: %w", err) + } + + resp, err := client.Do(req) + if err != nil { + return fmt.Errorf("ollama unreachable at %s: %w", ollamaURL, err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("ollama returned status %d", resp.StatusCode) + } + + var tags ollamaTagsResponse + if err := json.NewDecoder(resp.Body).Decode(&tags); err != nil { + return fmt.Errorf("decoding tags response: %w", err) + } + + for _, m := range tags.Models { + // Ollama model names can include ":latest" suffix + if m.Name == model || m.Name == model+":latest" { + return nil + } + } + + return fmt.Errorf("model %q not found in ollama (have %d models); run: ollama pull %s", + model, len(tags.Models), model) +} diff --git a/pkg/memory/ollama_test.go b/pkg/memory/ollama_test.go new file mode 100644 index 000000000..40252ce24 --- /dev/null +++ b/pkg/memory/ollama_test.go @@ -0,0 +1,113 @@ +// ABOUTME: Tests for Ollama health check and connectivity. +// ABOUTME: Uses httptest mock server to verify model detection behavior. +package memory + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" +) + +func TestCheckOllamaAvailable_ModelFound(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/tags" { + http.Error(w, "not found", http.StatusNotFound) + return + } + resp := ollamaTagsResponse{ + Models: []ollamaModel{ + {Name: "llama3.2:latest"}, + {Name: "nomic-embed-text:latest"}, + }, + } + json.NewEncoder(w).Encode(resp) + })) + defer server.Close() + + err := CheckOllamaAvailable(t.Context(), server.URL, "nomic-embed-text") + if err != nil { + t.Fatalf("expected no error, got: %v", err) + } +} + +func TestCheckOllamaAvailable_ModelFoundExactMatch(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + resp := ollamaTagsResponse{ + Models: []ollamaModel{ + {Name: "nomic-embed-text"}, + }, + } + json.NewEncoder(w).Encode(resp) + })) + defer server.Close() + + err := CheckOllamaAvailable(t.Context(), server.URL, "nomic-embed-text") + if err != nil { + t.Fatalf("expected no error for exact match, got: %v", err) + } +} + +func TestCheckOllamaAvailable_ModelNotFound(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + resp := ollamaTagsResponse{ + Models: []ollamaModel{ + {Name: "llama3.2:latest"}, + }, + } + json.NewEncoder(w).Encode(resp) + })) + defer server.Close() + + err := CheckOllamaAvailable(t.Context(), server.URL, "nomic-embed-text") + if err == nil { + t.Fatal("expected error when model not found") + } + if got := err.Error(); !contains(got, "not found") { + t.Errorf("error = %q, want it to contain 'not found'", got) + } +} + +func TestCheckOllamaAvailable_ServerUnreachable(t *testing.T) { + err := CheckOllamaAvailable(t.Context(), "http://127.0.0.1:1", "nomic-embed-text") + if err == nil { + t.Fatal("expected error when server unreachable") + } + if got := err.Error(); !contains(got, "unreachable") { + t.Errorf("error = %q, want it to contain 'unreachable'", got) + } +} + +func TestCheckOllamaAvailable_ServerError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "internal error", http.StatusInternalServerError) + })) + defer server.Close() + + err := CheckOllamaAvailable(t.Context(), server.URL, "nomic-embed-text") + if err == nil { + t.Fatal("expected error on server 500") + } + if got := err.Error(); !contains(got, "status 500") { + t.Errorf("error = %q, want it to contain 'status 500'", got) + } +} + +func TestCheckOllamaAvailable_DefaultValues(t *testing.T) { + // When called with empty strings, it should use defaults and likely fail + // (no local Ollama in test environment). Just verify it doesn't panic. + _ = CheckOllamaAvailable(t.Context(), "", "") +} + +func contains(s, substr string) bool { + return len(s) >= len(substr) && searchSubstring(s, substr) +} + +func searchSubstring(s, substr string) bool { + for i := 0; i <= len(s)-len(substr); i++ { + if s[i:i+len(substr)] == substr { + return true + } + } + return false +} diff --git a/pkg/memory/store.go b/pkg/memory/store.go new file mode 100644 index 000000000..221a6ab73 --- /dev/null +++ b/pkg/memory/store.go @@ -0,0 +1,206 @@ +// ABOUTME: Semantic memory store backed by chromem-go vector database. +// ABOUTME: Provides Remember/Recall operations with Ollama embeddings for persistent memory. +package memory + +import ( + "context" + "fmt" + "strings" + "time" + + chromem "github.com/philippgille/chromem-go" +) + +// MemoryEntry represents a single memory to be stored. +type MemoryEntry struct { + ID string + Content string + Category string // "preference", "fact", "decision", etc. + Tags []string + Source string // "agent", "auto-extract" + Timestamp time.Time +} + +// RecallResult is a memory entry returned from a search, with similarity score. +type RecallResult struct { + MemoryEntry + Similarity float32 +} + +// Store is the interface for semantic memory operations. +// Using an interface allows easy testing with mock implementations. +type Store interface { + IsAvailable() bool + Count() int + Remember(ctx context.Context, entry MemoryEntry) error + Recall(ctx context.Context, query string, topK int) ([]RecallResult, error) +} + +// SemanticStore implements Store using chromem-go for vector storage +// and Ollama for embedding generation. +type SemanticStore struct { + db *chromem.DB + collection *chromem.Collection + available bool +} + +// NewSemanticStore creates a persistent vector store at persistDir. +// It connects to Ollama at ollamaURL for embeddings using the given model. +// If Ollama is unreachable, the store is created in a degraded state +// where IsAvailable() returns false and all operations return clean errors. +func NewSemanticStore(persistDir, ollamaURL, embeddingModel string) (*SemanticStore, error) { + if ollamaURL == "" { + ollamaURL = DefaultOllamaURL + } + if embeddingModel == "" { + embeddingModel = DefaultEmbeddingModel + } + + // chromem-go expects the Ollama API base URL (ending in /api), not + // the server root. Normalize so users can provide either form. + ollamaAPIBase := normalizeOllamaURL(ollamaURL) + embedFn := chromem.NewEmbeddingFuncOllama(embeddingModel, ollamaAPIBase) + + db, err := chromem.NewPersistentDB(persistDir, false) + if err != nil { + return &SemanticStore{available: false}, fmt.Errorf("creating vector db: %w", err) + } + + collection, err := db.GetOrCreateCollection("memories", nil, embedFn) + if err != nil { + return &SemanticStore{available: false}, fmt.Errorf("creating collection: %w", err) + } + + return &SemanticStore{ + db: db, + collection: collection, + available: true, + }, nil +} + +// NewSemanticStoreWithEmbedding creates a store with a custom embedding function. +// This is primarily for testing without requiring Ollama. +func NewSemanticStoreWithEmbedding(persistDir string, embedFn chromem.EmbeddingFunc) (*SemanticStore, error) { + db, err := chromem.NewPersistentDB(persistDir, false) + if err != nil { + return nil, fmt.Errorf("creating vector db: %w", err) + } + + collection, err := db.GetOrCreateCollection("memories", nil, embedFn) + if err != nil { + return nil, fmt.Errorf("creating collection: %w", err) + } + + return &SemanticStore{ + db: db, + collection: collection, + available: true, + }, nil +} + +func (s *SemanticStore) IsAvailable() bool { + return s.available +} + +func (s *SemanticStore) Count() int { + if !s.available { + return 0 + } + return s.collection.Count() +} + +func (s *SemanticStore) Remember(ctx context.Context, entry MemoryEntry) error { + if !s.available { + return fmt.Errorf("semantic memory is not available") + } + + if entry.ID == "" { + entry.ID = fmt.Sprintf("mem_%d", time.Now().UnixNano()) + } + if entry.Timestamp.IsZero() { + entry.Timestamp = time.Now() + } + + metadata := map[string]string{ + "category": entry.Category, + "source": entry.Source, + "timestamp": entry.Timestamp.Format(time.RFC3339), + } + if len(entry.Tags) > 0 { + metadata["tags"] = strings.Join(entry.Tags, ",") + } + + doc := chromem.Document{ + ID: entry.ID, + Content: entry.Content, + Metadata: metadata, + } + + return s.collection.AddDocument(ctx, doc) +} + +func (s *SemanticStore) Recall(ctx context.Context, query string, topK int) ([]RecallResult, error) { + if !s.available { + return nil, fmt.Errorf("semantic memory is not available") + } + if topK <= 0 { + topK = 5 + } + if topK > 20 { + topK = 20 + } + + // chromem-go returns an error if topK > collection count, so we + // cap it to the collection size. + count := s.collection.Count() + if count == 0 { + return nil, nil + } + if topK > count { + topK = count + } + + results, err := s.collection.Query(ctx, query, topK, nil, nil) + if err != nil { + return nil, fmt.Errorf("querying memories: %w", err) + } + + var recalls []RecallResult + for _, r := range results { + entry := MemoryEntry{ + ID: r.ID, + Content: r.Content, + Category: r.Metadata["category"], + Source: r.Metadata["source"], + } + if ts, ok := r.Metadata["timestamp"]; ok { + if t, err := time.Parse(time.RFC3339, ts); err == nil { + entry.Timestamp = t + } + } + if tagStr, ok := r.Metadata["tags"]; ok && tagStr != "" { + entry.Tags = strings.Split(tagStr, ",") + } + + recalls = append(recalls, RecallResult{ + MemoryEntry: entry, + Similarity: r.Similarity, + }) + } + + return recalls, nil +} + +// normalizeOllamaURL ensures the URL ends with /api as required by +// chromem-go's NewEmbeddingFuncOllama. Users typically configure the +// Ollama server root (http://localhost:11434) without the /api suffix. +func normalizeOllamaURL(url string) string { + url = strings.TrimRight(url, "/") + if url == "" { + return "" // let chromem-go use its default + } + if !strings.HasSuffix(url, "/api") { + url += "/api" + } + return url +} diff --git a/pkg/memory/store_test.go b/pkg/memory/store_test.go new file mode 100644 index 000000000..07750fe0f --- /dev/null +++ b/pkg/memory/store_test.go @@ -0,0 +1,262 @@ +// ABOUTME: Tests for the SemanticStore vector memory implementation. +// ABOUTME: Uses a deterministic mock embedding function to test without Ollama. +package memory + +import ( + "context" + "os" + "testing" +) + +// mockEmbeddingFunc returns a deterministic embedding based on text length. +// This gives different texts different (but reproducible) embeddings, +// allowing cosine similarity to work for testing. +func mockEmbeddingFunc(text string) []float32 { + const dims = 64 + embedding := make([]float32, dims) + for i, ch := range text { + embedding[i%dims] += float32(ch) / 1000.0 + } + // Normalize to unit vector for cosine similarity + var norm float32 + for _, v := range embedding { + norm += v * v + } + if norm > 0 { + norm = sqrt32(norm) + for i := range embedding { + embedding[i] /= norm + } + } + return embedding +} + +func sqrt32(x float32) float32 { + // Newton's method for float32 + z := x / 2 + for i := 0; i < 10; i++ { + z = z - (z*z-x)/(2*z) + } + return z +} + +func newTestStore(t *testing.T) *SemanticStore { + t.Helper() + dir := t.TempDir() + + embedFn := func(_ context.Context, text string) ([]float32, error) { + return mockEmbeddingFunc(text), nil + } + + store, err := NewSemanticStoreWithEmbedding(dir, embedFn) + if err != nil { + t.Fatalf("NewSemanticStoreWithEmbedding: %v", err) + } + return store +} + +func TestSemanticStore_IsAvailable(t *testing.T) { + store := newTestStore(t) + if !store.IsAvailable() { + t.Error("expected store to be available") + } +} + +func TestSemanticStore_RememberAndRecall(t *testing.T) { + store := newTestStore(t) + ctx := t.Context() + + // Store some memories + entries := []MemoryEntry{ + {Content: "User prefers dark mode in all applications", Category: "preference", Source: "agent"}, + {Content: "The database password is rotated every 90 days", Category: "fact", Source: "agent"}, + {Content: "We decided to use PostgreSQL instead of MySQL", Category: "decision", Source: "agent"}, + } + + for _, entry := range entries { + if err := store.Remember(ctx, entry); err != nil { + t.Fatalf("Remember(%q): %v", entry.Content, err) + } + } + + // Recall with a query related to preferences + results, err := store.Recall(ctx, "What does the user prefer for themes?", 3) + if err != nil { + t.Fatalf("Recall: %v", err) + } + + if len(results) != 3 { + t.Fatalf("expected 3 results, got %d", len(results)) + } + + // All results should have content + for _, r := range results { + if r.Content == "" { + t.Error("result has empty content") + } + if r.Similarity <= 0 { + t.Errorf("expected positive similarity, got %f", r.Similarity) + } + } +} + +func TestSemanticStore_RecallEmpty(t *testing.T) { + store := newTestStore(t) + ctx := t.Context() + + results, err := store.Recall(ctx, "anything", 5) + if err != nil { + t.Fatalf("Recall on empty store: %v", err) + } + if len(results) != 0 { + t.Errorf("expected 0 results from empty store, got %d", len(results)) + } +} + +func TestSemanticStore_RememberAutoID(t *testing.T) { + store := newTestStore(t) + ctx := t.Context() + + entry := MemoryEntry{Content: "test content", Category: "fact"} + if err := store.Remember(ctx, entry); err != nil { + t.Fatalf("Remember: %v", err) + } + + results, err := store.Recall(ctx, "test", 1) + if err != nil { + t.Fatalf("Recall: %v", err) + } + if len(results) != 1 { + t.Fatalf("expected 1 result, got %d", len(results)) + } + if results[0].ID == "" { + t.Error("expected auto-generated ID, got empty") + } +} + +func TestSemanticStore_MetadataPreserved(t *testing.T) { + store := newTestStore(t) + ctx := t.Context() + + entry := MemoryEntry{ + Content: "Important fact about the project", + Category: "fact", + Tags: []string{"project", "architecture"}, + Source: "auto-extract", + } + if err := store.Remember(ctx, entry); err != nil { + t.Fatalf("Remember: %v", err) + } + + results, err := store.Recall(ctx, "project fact", 1) + if err != nil { + t.Fatalf("Recall: %v", err) + } + if len(results) != 1 { + t.Fatalf("expected 1 result, got %d", len(results)) + } + + r := results[0] + if r.Category != "fact" { + t.Errorf("category = %q, want %q", r.Category, "fact") + } + if r.Source != "auto-extract" { + t.Errorf("source = %q, want %q", r.Source, "auto-extract") + } + if len(r.Tags) != 2 || r.Tags[0] != "project" || r.Tags[1] != "architecture" { + t.Errorf("tags = %v, want [project architecture]", r.Tags) + } + if r.Timestamp.IsZero() { + t.Error("expected non-zero timestamp") + } +} + +func TestSemanticStore_RecallTopKCapped(t *testing.T) { + store := newTestStore(t) + ctx := t.Context() + + // Store 2 entries, ask for 10 + store.Remember(ctx, MemoryEntry{Content: "first memory"}) + store.Remember(ctx, MemoryEntry{Content: "second memory"}) + + results, err := store.Recall(ctx, "memory", 10) + if err != nil { + t.Fatalf("Recall: %v", err) + } + if len(results) != 2 { + t.Errorf("expected 2 results (capped to collection size), got %d", len(results)) + } +} + +func TestSemanticStore_UnavailableStore(t *testing.T) { + store := &SemanticStore{available: false} + + if store.IsAvailable() { + t.Error("expected unavailable store") + } + + err := store.Remember(t.Context(), MemoryEntry{Content: "test"}) + if err == nil { + t.Error("expected error from unavailable store Remember") + } + + _, err = store.Recall(t.Context(), "test", 5) + if err == nil { + t.Error("expected error from unavailable store Recall") + } +} + +func TestSemanticStore_Persistence(t *testing.T) { + dir := t.TempDir() + ctx := t.Context() + + embedFn := func(_ context.Context, text string) ([]float32, error) { + return mockEmbeddingFunc(text), nil + } + + // Create store and add memory + store1, err := NewSemanticStoreWithEmbedding(dir, embedFn) + if err != nil { + t.Fatalf("first store: %v", err) + } + store1.Remember(ctx, MemoryEntry{Content: "persistent memory test"}) + + // Create a second store from the same directory + store2, err := NewSemanticStoreWithEmbedding(dir, embedFn) + if err != nil { + t.Fatalf("second store: %v", err) + } + + results, err := store2.Recall(ctx, "persistent", 1) + if err != nil { + t.Fatalf("Recall from second store: %v", err) + } + if len(results) != 1 { + t.Fatalf("expected 1 persisted result, got %d", len(results)) + } + if results[0].Content != "persistent memory test" { + t.Errorf("content = %q, want %q", results[0].Content, "persistent memory test") + } +} + +func TestNewSemanticStore_BadDir(t *testing.T) { + // Try to create a store in a path that can't exist + _, err := NewSemanticStoreWithEmbedding("/dev/null/impossible", nil) + if err == nil { + t.Error("expected error for impossible path") + } +} + +func TestNewSemanticStore_WithOllamaDefaults(t *testing.T) { + // This tests the constructor with empty strings (uses defaults). + // It will fail to connect to Ollama in test env, but should return + // a store (possibly in degraded state) without panicking. + dir := t.TempDir() + defer os.RemoveAll(dir) + + // We can't actually test with Ollama here, just verify no panic + store, _ := NewSemanticStore(dir, "http://127.0.0.1:1", "fake-model") + if store == nil { + t.Error("expected non-nil store even on connection failure") + } +} diff --git a/pkg/tools/recall.go b/pkg/tools/recall.go new file mode 100644 index 000000000..45e9bfdef --- /dev/null +++ b/pkg/tools/recall.go @@ -0,0 +1,83 @@ +// ABOUTME: Recall tool for semantic search over stored memories. +// ABOUTME: Queries the vector store with natural language and returns ranked results. +package tools + +import ( + "context" + "fmt" + "strings" + + "github.com/sipeed/picoclaw/pkg/memory" +) + +// RecallTool performs semantic search over stored memories. +type RecallTool struct { + store memory.Store +} + +func NewRecallTool(store memory.Store) *RecallTool { + return &RecallTool{store: store} +} + +func (t *RecallTool) Name() string { return "recall" } + +func (t *RecallTool) Description() string { + return "Search stored memories using natural language. Returns the most relevant memories ranked by similarity. Use this to recall facts, preferences, decisions, or context from previous conversations." +} + +func (t *RecallTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "query": map[string]any{ + "type": "string", + "description": "Natural language search query describing what you want to recall.", + }, + "top_k": map[string]any{ + "type": "integer", + "description": "Number of results to return (default 5, max 20).", + }, + }, + "required": []any{"query"}, + } +} + +func (t *RecallTool) Execute(ctx context.Context, args map[string]any) *ToolResult { + if !t.store.IsAvailable() { + return ErrorResult("Semantic memory is not available. Ollama may not be running.") + } + + query, _ := args["query"].(string) + if strings.TrimSpace(query) == "" { + return ErrorResult("query is required") + } + + topK := 5 + if k, ok := args["top_k"].(float64); ok && k > 0 { + topK = int(k) + } + + results, err := t.store.Recall(ctx, query, topK) + if err != nil { + return ErrorResult("Failed to search memories: " + err.Error()) + } + + if len(results) == 0 { + return SilentResult("No memories found matching: " + query) + } + + var sb strings.Builder + fmt.Fprintf(&sb, "Found %d memories:\n\n", len(results)) + for i, r := range results { + fmt.Fprintf(&sb, "%d. [%.0f%% match] [%s] %s\n", + i+1, r.Similarity*100, r.Category, r.Content) + if len(r.Tags) > 0 { + fmt.Fprintf(&sb, " Tags: %s\n", strings.Join(r.Tags, ", ")) + } + if !r.Timestamp.IsZero() { + fmt.Fprintf(&sb, " Stored: %s\n", r.Timestamp.Format("2006-01-02 15:04")) + } + } + + return SilentResult(sb.String()) +} diff --git a/pkg/tools/recall_test.go b/pkg/tools/recall_test.go new file mode 100644 index 000000000..34c5ff00f --- /dev/null +++ b/pkg/tools/recall_test.go @@ -0,0 +1,121 @@ +// ABOUTME: Tests for the recall tool. +// ABOUTME: Uses a mock memory store to verify search and formatting behavior. +package tools + +import ( + "strings" + "testing" + "time" + + "github.com/sipeed/picoclaw/pkg/memory" +) + +func TestRecallTool_Name(t *testing.T) { + tool := NewRecallTool(&mockStore{available: true}) + if got := tool.Name(); got != "recall" { + t.Errorf("Name() = %q, want %q", got, "recall") + } +} + +func TestRecallTool_Execute_Success(t *testing.T) { + store := &mockStore{ + available: true, + entries: []memory.MemoryEntry{ + { + Content: "User prefers dark mode", + Category: "preference", + Tags: []string{"ui"}, + Timestamp: time.Date(2026, 2, 22, 10, 0, 0, 0, time.UTC), + }, + { + Content: "Database is PostgreSQL", + Category: "fact", + }, + }, + } + tool := NewRecallTool(store) + + result := tool.Execute(t.Context(), map[string]any{ + "query": "user preferences", + "top_k": float64(5), + }) + + if result.IsError { + t.Fatalf("expected success, got error: %s", result.ForLLM) + } + if !result.Silent { + t.Error("expected silent result") + } + if !strings.Contains(result.ForLLM, "Found 2 memories") { + t.Errorf("expected 'Found 2 memories' in result, got: %s", result.ForLLM) + } + if !strings.Contains(result.ForLLM, "dark mode") { + t.Errorf("expected 'dark mode' in result, got: %s", result.ForLLM) + } + if !strings.Contains(result.ForLLM, "90% match") { + t.Errorf("expected '90%% match' in result, got: %s", result.ForLLM) + } +} + +func TestRecallTool_Execute_NoResults(t *testing.T) { + store := &mockStore{available: true, entries: nil} + tool := NewRecallTool(store) + + result := tool.Execute(t.Context(), map[string]any{ + "query": "nonexistent", + }) + + if result.IsError { + t.Fatalf("expected success (no results), got error: %s", result.ForLLM) + } + if !strings.Contains(result.ForLLM, "No memories found") { + t.Errorf("expected 'No memories found' in result, got: %s", result.ForLLM) + } +} + +func TestRecallTool_Execute_EmptyQuery(t *testing.T) { + tool := NewRecallTool(&mockStore{available: true}) + result := tool.Execute(t.Context(), map[string]any{ + "query": "", + }) + if !result.IsError { + t.Error("expected error for empty query") + } +} + +func TestRecallTool_Execute_Unavailable(t *testing.T) { + tool := NewRecallTool(&mockStore{available: false}) + result := tool.Execute(t.Context(), map[string]any{ + "query": "test", + }) + if !result.IsError { + t.Error("expected error when store unavailable") + } +} + +func TestRecallTool_Execute_DefaultTopK(t *testing.T) { + // Ensure default top_k of 5 is used when not specified + store := &mockStore{ + available: true, + entries: []memory.MemoryEntry{ + {Content: "a", Category: "fact"}, + {Content: "b", Category: "fact"}, + {Content: "c", Category: "fact"}, + {Content: "d", Category: "fact"}, + {Content: "e", Category: "fact"}, + {Content: "f", Category: "fact"}, + }, + } + tool := NewRecallTool(store) + + result := tool.Execute(t.Context(), map[string]any{ + "query": "test", + }) + + if result.IsError { + t.Fatalf("expected success, got error: %s", result.ForLLM) + } + if !strings.Contains(result.ForLLM, "Found 5 memories") { + t.Errorf("expected default top_k of 5, got: %s", result.ForLLM) + } +} diff --git a/pkg/tools/remember.go b/pkg/tools/remember.go new file mode 100644 index 000000000..f074d9712 --- /dev/null +++ b/pkg/tools/remember.go @@ -0,0 +1,86 @@ +// ABOUTME: Remember tool for storing semantic memories via vector embeddings. +// ABOUTME: Accepts content, category, and tags, and persists them for later recall. +package tools + +import ( + "context" + "strings" + + "github.com/sipeed/picoclaw/pkg/memory" +) + +// RememberTool stores a memory entry in the semantic memory store. +type RememberTool struct { + store memory.Store +} + +func NewRememberTool(store memory.Store) *RememberTool { + return &RememberTool{store: store} +} + +func (t *RememberTool) Name() string { return "remember" } + +func (t *RememberTool) Description() string { + return "Store a memory for later recall. Use this to remember important facts, user preferences, decisions, or any information worth persisting across conversations." +} + +func (t *RememberTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "content": map[string]any{ + "type": "string", + "description": "The information to remember. Be specific and self-contained.", + }, + "category": map[string]any{ + "type": "string", + "description": "Category: preference, fact, decision, context, or other.", + "enum": []string{"preference", "fact", "decision", "context", "other"}, + }, + "tags": map[string]any{ + "type": "string", + "description": "Comma-separated tags for organization (e.g. 'project,database,architecture').", + }, + }, + "required": []any{"content"}, + } +} + +func (t *RememberTool) Execute(ctx context.Context, args map[string]any) *ToolResult { + if !t.store.IsAvailable() { + return ErrorResult("Semantic memory is not available. Ollama may not be running.") + } + + content, _ := args["content"].(string) + if strings.TrimSpace(content) == "" { + return ErrorResult("content is required") + } + + category, _ := args["category"].(string) + if category == "" { + category = "other" + } + + var tags []string + if tagsStr, ok := args["tags"].(string); ok && tagsStr != "" { + for _, tag := range strings.Split(tagsStr, ",") { + tag = strings.TrimSpace(tag) + if tag != "" { + tags = append(tags, tag) + } + } + } + + entry := memory.MemoryEntry{ + Content: content, + Category: category, + Tags: tags, + Source: "agent", + } + + if err := t.store.Remember(ctx, entry); err != nil { + return ErrorResult("Failed to store memory: " + err.Error()) + } + + return SilentResult("Memory stored successfully: " + content) +} diff --git a/pkg/tools/remember_test.go b/pkg/tools/remember_test.go new file mode 100644 index 000000000..d6c97b304 --- /dev/null +++ b/pkg/tools/remember_test.go @@ -0,0 +1,115 @@ +// ABOUTME: Tests for the remember tool. +// ABOUTME: Uses a mock memory store to verify tool behavior without Ollama. +package tools + +import ( + "context" + "testing" + + "github.com/sipeed/picoclaw/pkg/memory" +) + +// mockStore implements memory.Store for testing. +type mockStore struct { + available bool + entries []memory.MemoryEntry +} + +func (m *mockStore) IsAvailable() bool { return m.available } +func (m *mockStore) Count() int { return len(m.entries) } + +func (m *mockStore) Remember(_ context.Context, entry memory.MemoryEntry) error { + m.entries = append(m.entries, entry) + return nil +} + +func (m *mockStore) Recall(_ context.Context, query string, topK int) ([]memory.RecallResult, error) { + var results []memory.RecallResult + for i, e := range m.entries { + if i >= topK { + break + } + results = append(results, memory.RecallResult{ + MemoryEntry: e, + Similarity: 0.9 - float32(i)*0.1, + }) + } + return results, nil +} + +func TestRememberTool_Name(t *testing.T) { + tool := NewRememberTool(&mockStore{available: true}) + if got := tool.Name(); got != "remember" { + t.Errorf("Name() = %q, want %q", got, "remember") + } +} + +func TestRememberTool_Execute_Success(t *testing.T) { + store := &mockStore{available: true} + tool := NewRememberTool(store) + + result := tool.Execute(t.Context(), map[string]any{ + "content": "User prefers dark mode", + "category": "preference", + "tags": "ui,theme", + }) + + if result.IsError { + t.Fatalf("expected success, got error: %s", result.ForLLM) + } + if !result.Silent { + t.Error("expected silent result") + } + if len(store.entries) != 1 { + t.Fatalf("expected 1 entry, got %d", len(store.entries)) + } + entry := store.entries[0] + if entry.Content != "User prefers dark mode" { + t.Errorf("content = %q, want %q", entry.Content, "User prefers dark mode") + } + if entry.Category != "preference" { + t.Errorf("category = %q, want %q", entry.Category, "preference") + } + if len(entry.Tags) != 2 || entry.Tags[0] != "ui" || entry.Tags[1] != "theme" { + t.Errorf("tags = %v, want [ui theme]", entry.Tags) + } + if entry.Source != "agent" { + t.Errorf("source = %q, want %q", entry.Source, "agent") + } +} + +func TestRememberTool_Execute_DefaultCategory(t *testing.T) { + store := &mockStore{available: true} + tool := NewRememberTool(store) + + result := tool.Execute(t.Context(), map[string]any{ + "content": "some fact", + }) + + if result.IsError { + t.Fatalf("expected success, got error: %s", result.ForLLM) + } + if store.entries[0].Category != "other" { + t.Errorf("category = %q, want default %q", store.entries[0].Category, "other") + } +} + +func TestRememberTool_Execute_EmptyContent(t *testing.T) { + tool := NewRememberTool(&mockStore{available: true}) + result := tool.Execute(t.Context(), map[string]any{ + "content": "", + }) + if !result.IsError { + t.Error("expected error for empty content") + } +} + +func TestRememberTool_Execute_Unavailable(t *testing.T) { + tool := NewRememberTool(&mockStore{available: false}) + result := tool.Execute(t.Context(), map[string]any{ + "content": "test", + }) + if !result.IsError { + t.Error("expected error when store unavailable") + } +}