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 <noreply@anthropic.com>
This commit is contained in:
parent
c656fc8a10
commit
3d5fe5bc4d
15 changed files with 1416 additions and 26 deletions
1
go.mod
1
go.mod
|
|
@ -22,6 +22,7 @@ require (
|
||||||
|
|
||||||
require (
|
require (
|
||||||
github.com/davecgh/go-spew v1.1.1 // indirect
|
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
|
github.com/pmezard/go-difflib v1.0.0 // indirect
|
||||||
gopkg.in/yaml.v3 v3.0.1 // indirect
|
gopkg.in/yaml.v3 v3.0.1 // indirect
|
||||||
)
|
)
|
||||||
|
|
|
||||||
2
go.sum
2
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/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 h1:6MEoNoV8sbjOVmXdvhmuX3BjVbVdcExbVyGixiyJ8ys=
|
||||||
github.com/openai/openai-go/v3 v3.22.0/go.mod h1:cdufnVK14cWcT9qA1rRtrXx4FTRsgbDPW7Ia7SS5cZo=
|
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/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 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||||
|
|
|
||||||
|
|
@ -95,6 +95,11 @@ func NewAgentInstance(
|
||||||
}
|
}
|
||||||
candidates := providers.ResolveCandidates(modelCfg, defaults.Provider)
|
candidates := providers.ResolveCandidates(modelCfg, defaults.Provider)
|
||||||
|
|
||||||
|
contextWindow := defaults.ContextWindow
|
||||||
|
if contextWindow == 0 {
|
||||||
|
contextWindow = 128000
|
||||||
|
}
|
||||||
|
|
||||||
return &AgentInstance{
|
return &AgentInstance{
|
||||||
ID: agentID,
|
ID: agentID,
|
||||||
Name: agentName,
|
Name: agentName,
|
||||||
|
|
@ -104,7 +109,7 @@ func NewAgentInstance(
|
||||||
MaxIterations: maxIter,
|
MaxIterations: maxIter,
|
||||||
MaxTokens: maxTokens,
|
MaxTokens: maxTokens,
|
||||||
Temperature: temperature,
|
Temperature: temperature,
|
||||||
ContextWindow: maxTokens,
|
ContextWindow: contextWindow,
|
||||||
Provider: provider,
|
Provider: provider,
|
||||||
Sessions: sessionsManager,
|
Sessions: sessionsManager,
|
||||||
ContextBuilder: contextBuilder,
|
ContextBuilder: contextBuilder,
|
||||||
|
|
|
||||||
|
|
@ -16,11 +16,14 @@ import (
|
||||||
"time"
|
"time"
|
||||||
"unicode/utf8"
|
"unicode/utf8"
|
||||||
|
|
||||||
|
"path/filepath"
|
||||||
|
|
||||||
"github.com/sipeed/picoclaw/pkg/bus"
|
"github.com/sipeed/picoclaw/pkg/bus"
|
||||||
"github.com/sipeed/picoclaw/pkg/channels"
|
"github.com/sipeed/picoclaw/pkg/channels"
|
||||||
"github.com/sipeed/picoclaw/pkg/config"
|
"github.com/sipeed/picoclaw/pkg/config"
|
||||||
"github.com/sipeed/picoclaw/pkg/constants"
|
"github.com/sipeed/picoclaw/pkg/constants"
|
||||||
"github.com/sipeed/picoclaw/pkg/logger"
|
"github.com/sipeed/picoclaw/pkg/logger"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/memory"
|
||||||
"github.com/sipeed/picoclaw/pkg/providers"
|
"github.com/sipeed/picoclaw/pkg/providers"
|
||||||
"github.com/sipeed/picoclaw/pkg/routing"
|
"github.com/sipeed/picoclaw/pkg/routing"
|
||||||
"github.com/sipeed/picoclaw/pkg/skills"
|
"github.com/sipeed/picoclaw/pkg/skills"
|
||||||
|
|
@ -38,6 +41,7 @@ type AgentLoop struct {
|
||||||
summarizing sync.Map
|
summarizing sync.Map
|
||||||
fallback *providers.FallbackChain
|
fallback *providers.FallbackChain
|
||||||
channelManager *channels.Manager
|
channelManager *channels.Manager
|
||||||
|
semanticStores map[string]memory.Store
|
||||||
}
|
}
|
||||||
|
|
||||||
// processOptions configures how a message is processed
|
// 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 {
|
func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers.LLMProvider) *AgentLoop {
|
||||||
registry := NewAgentRegistry(cfg, provider)
|
registry := NewAgentRegistry(cfg, provider)
|
||||||
|
|
||||||
// Register shared tools to all agents
|
// Register shared tools to all agents (including memory tools)
|
||||||
registerSharedTools(cfg, msgBus, registry, provider)
|
stores := registerSharedTools(cfg, msgBus, registry, provider)
|
||||||
|
|
||||||
// Set up shared fallback chain
|
// Set up shared fallback chain
|
||||||
cooldown := providers.NewCooldownTracker()
|
cooldown := providers.NewCooldownTracker()
|
||||||
|
|
@ -76,16 +80,20 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers
|
||||||
state: stateManager,
|
state: stateManager,
|
||||||
summarizing: sync.Map{},
|
summarizing: sync.Map{},
|
||||||
fallback: fallbackChain,
|
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(
|
func registerSharedTools(
|
||||||
cfg *config.Config,
|
cfg *config.Config,
|
||||||
msgBus *bus.MessageBus,
|
msgBus *bus.MessageBus,
|
||||||
registry *AgentRegistry,
|
registry *AgentRegistry,
|
||||||
provider providers.LLMProvider,
|
provider providers.LLMProvider,
|
||||||
) {
|
) map[string]memory.Store {
|
||||||
|
stores := make(map[string]memory.Store)
|
||||||
|
|
||||||
for _, agentID := range registry.ListAgentIDs() {
|
for _, agentID := range registry.ListAgentIDs() {
|
||||||
agent, ok := registry.GetAgent(agentID)
|
agent, ok := registry.GetAgent(agentID)
|
||||||
if !ok {
|
if !ok {
|
||||||
|
|
@ -145,9 +153,30 @@ func registerSharedTools(
|
||||||
})
|
})
|
||||||
agent.Tools.Register(spawnTool)
|
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
|
// Update context builder with the complete tools registry
|
||||||
agent.ContextBuilder.SetToolsRegistry(agent.Tools)
|
agent.ContextBuilder.SetToolsRegistry(agent.Tools)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
return stores
|
||||||
}
|
}
|
||||||
|
|
||||||
func (al *AgentLoop) Run(ctx context.Context) error {
|
func (al *AgentLoop) Run(ctx context.Context) error {
|
||||||
|
|
@ -514,14 +543,16 @@ func (al *AgentLoop) runLLMIteration(
|
||||||
var response *providers.LLMResponse
|
var response *providers.LLMResponse
|
||||||
var err error
|
var err error
|
||||||
|
|
||||||
|
llmOpts := map[string]any{
|
||||||
|
"max_tokens": agent.MaxTokens,
|
||||||
|
"temperature": agent.Temperature,
|
||||||
|
}
|
||||||
|
|
||||||
callLLM := func() (*providers.LLMResponse, error) {
|
callLLM := func() (*providers.LLMResponse, error) {
|
||||||
if len(agent.Candidates) > 1 && al.fallback != nil {
|
if len(agent.Candidates) > 1 && al.fallback != nil {
|
||||||
fbResult, fbErr := al.fallback.Execute(ctx, agent.Candidates,
|
fbResult, fbErr := al.fallback.Execute(ctx, agent.Candidates,
|
||||||
func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) {
|
func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) {
|
||||||
return agent.Provider.Chat(ctx, messages, providerToolDefs, model, map[string]any{
|
return agent.Provider.Chat(ctx, messages, providerToolDefs, model, llmOpts)
|
||||||
"max_tokens": agent.MaxTokens,
|
|
||||||
"temperature": agent.Temperature,
|
|
||||||
})
|
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
if fbErr != nil {
|
if fbErr != nil {
|
||||||
|
|
@ -534,11 +565,18 @@ func (al *AgentLoop) runLLMIteration(
|
||||||
}
|
}
|
||||||
return fbResult.Response, nil
|
return fbResult.Response, nil
|
||||||
}
|
}
|
||||||
return agent.Provider.Chat(ctx, messages, providerToolDefs, agent.Model, map[string]any{
|
// Use streaming if the provider supports it
|
||||||
"max_tokens": agent.MaxTokens,
|
if sp, ok := agent.Provider.(providers.StreamingProvider); ok {
|
||||||
"temperature": agent.Temperature,
|
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
|
// Retry loop for context/token errors
|
||||||
maxRetries := 2
|
maxRetries := 2
|
||||||
|
|
@ -783,8 +821,9 @@ func (al *AgentLoop) forceCompression(agent *AgentInstance, sessionKey string) {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
// Helper to find the mid-point of the conversation
|
// Find the mid-point of the conversation, adjusted to avoid splitting
|
||||||
mid := len(conversation) / 2
|
// tool_use/tool_result pairs.
|
||||||
|
mid := safeSplitIndex(conversation, len(conversation)/2)
|
||||||
|
|
||||||
// New history structure:
|
// New history structure:
|
||||||
// 1. System Prompt (with compression note appended)
|
// 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.
|
// GetStartupInfo returns information about loaded tools and skills for logging.
|
||||||
func (al *AgentLoop) GetStartupInfo() map[string]any {
|
func (al *AgentLoop) GetStartupInfo() map[string]any {
|
||||||
info := make(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)
|
history := agent.Sessions.GetHistory(sessionKey)
|
||||||
summary := agent.Sessions.GetSummary(sessionKey)
|
summary := agent.Sessions.GetSummary(sessionKey)
|
||||||
|
|
||||||
// Keep last 4 messages for continuity
|
// Keep last few messages for continuity, adjusting the split point so we
|
||||||
if len(history) <= 4 {
|
// never orphan a tool_use from its tool_result(s).
|
||||||
|
if len(history) <= 6 {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
toSummarize := history[:len(history)-4]
|
keepIdx := safeSplitIndex(history, len(history)-6)
|
||||||
|
if keepIdx <= 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
toSummarize := history[:keepIdx]
|
||||||
|
|
||||||
// Oversized Message Guard
|
// Oversized Message Guard
|
||||||
maxMessageTokens := agent.ContextWindow / 2
|
maxMessageTokens := agent.ContextWindow / 2
|
||||||
|
|
@ -920,7 +991,7 @@ func (al *AgentLoop) summarizeSession(agent *AgentInstance, sessionKey string) {
|
||||||
omitted := false
|
omitted := false
|
||||||
|
|
||||||
for _, m := range toSummarize {
|
for _, m := range toSummarize {
|
||||||
if m.Role != "user" && m.Role != "assistant" {
|
if m.Role != "user" && m.Role != "assistant" && m.Role != "tool" {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
msgTokens := len(m.Content) / 2
|
msgTokens := len(m.Content) / 2
|
||||||
|
|
@ -975,8 +1046,94 @@ func (al *AgentLoop) summarizeSession(agent *AgentInstance, sessionKey string) {
|
||||||
|
|
||||||
if finalSummary != "" {
|
if finalSummary != "" {
|
||||||
agent.Sessions.SetSummary(sessionKey, finalSummary)
|
agent.Sessions.SetSummary(sessionKey, finalSummary)
|
||||||
agent.Sessions.TruncateHistory(sessionKey, 4)
|
agent.Sessions.TruncateHistory(sessionKey, len(history)-keepIdx)
|
||||||
agent.Sessions.Save(sessionKey)
|
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})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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
|
// failFirstMockProvider fails on the first N calls with a specific error
|
||||||
type failFirstMockProvider struct {
|
type failFirstMockProvider struct {
|
||||||
failures int
|
failures int
|
||||||
|
|
|
||||||
|
|
@ -175,6 +175,7 @@ type AgentDefaults struct {
|
||||||
ImageModel string `json:"image_model,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_IMAGE_MODEL"`
|
ImageModel string `json:"image_model,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_IMAGE_MODEL"`
|
||||||
ImageModelFallbacks []string `json:"image_model_fallbacks,omitempty"`
|
ImageModelFallbacks []string `json:"image_model_fallbacks,omitempty"`
|
||||||
MaxTokens int `json:"max_tokens" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOKENS"`
|
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"`
|
Temperature *float64 `json:"temperature,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_TEMPERATURE"`
|
||||||
MaxToolIterations int `json:"max_tool_iterations" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOOL_ITERATIONS"`
|
MaxToolIterations int `json:"max_tool_iterations" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOOL_ITERATIONS"`
|
||||||
}
|
}
|
||||||
|
|
@ -449,6 +450,14 @@ type ToolsConfig struct {
|
||||||
Cron CronToolsConfig `json:"cron"`
|
Cron CronToolsConfig `json:"cron"`
|
||||||
Exec ExecConfig `json:"exec"`
|
Exec ExecConfig `json:"exec"`
|
||||||
Skills SkillsToolsConfig `json:"skills"`
|
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 {
|
type SkillsToolsConfig struct {
|
||||||
|
|
|
||||||
|
|
@ -15,6 +15,7 @@ func DefaultConfig() *Config {
|
||||||
Provider: "",
|
Provider: "",
|
||||||
Model: "glm-4.7",
|
Model: "glm-4.7",
|
||||||
MaxTokens: 8192,
|
MaxTokens: 8192,
|
||||||
|
ContextWindow: 128000,
|
||||||
Temperature: nil, // nil means use provider default
|
Temperature: nil, // nil means use provider default
|
||||||
MaxToolIterations: 20,
|
MaxToolIterations: 20,
|
||||||
},
|
},
|
||||||
|
|
@ -303,6 +304,11 @@ func DefaultConfig() *Config {
|
||||||
TTLSeconds: 300,
|
TTLSeconds: 300,
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
|
Memory: MemoryConfig{
|
||||||
|
OllamaURL: "http://localhost:11434",
|
||||||
|
EmbeddingModel: "nomic-embed-text",
|
||||||
|
AutoExtract: true,
|
||||||
|
},
|
||||||
},
|
},
|
||||||
Heartbeat: HeartbeatConfig{
|
Heartbeat: HeartbeatConfig{
|
||||||
Enabled: true,
|
Enabled: true,
|
||||||
|
|
|
||||||
67
pkg/memory/ollama.go
Normal file
67
pkg/memory/ollama.go
Normal file
|
|
@ -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)
|
||||||
|
}
|
||||||
113
pkg/memory/ollama_test.go
Normal file
113
pkg/memory/ollama_test.go
Normal file
|
|
@ -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
|
||||||
|
}
|
||||||
206
pkg/memory/store.go
Normal file
206
pkg/memory/store.go
Normal file
|
|
@ -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
|
||||||
|
}
|
||||||
262
pkg/memory/store_test.go
Normal file
262
pkg/memory/store_test.go
Normal file
|
|
@ -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")
|
||||||
|
}
|
||||||
|
}
|
||||||
83
pkg/tools/recall.go
Normal file
83
pkg/tools/recall.go
Normal file
|
|
@ -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())
|
||||||
|
}
|
||||||
121
pkg/tools/recall_test.go
Normal file
121
pkg/tools/recall_test.go
Normal file
|
|
@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
86
pkg/tools/remember.go
Normal file
86
pkg/tools/remember.go
Normal file
|
|
@ -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)
|
||||||
|
}
|
||||||
115
pkg/tools/remember_test.go
Normal file
115
pkg/tools/remember_test.go
Normal file
|
|
@ -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")
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Add table
Reference in a new issue