From f984a1129156e1cf262cfa0e81ae913007658c8a Mon Sep 17 00:00:00 2001 From: Peter Garland Date: Thu, 12 Feb 2026 16:48:31 -0600 Subject: [PATCH] feat: add vector search for memory and lifecycle hooks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace unbounded MEMORY.md dump with opt-in vector search that injects only the top-K relevant chunks into each prompt. When disabled (default), behavior is unchanged ā€” zero overhead. New packages: - pkg/vecstore: in-memory vector store with gob persistence, brute-force cosine search, markdown chunker (split by ## headers + paragraphs), and OpenAI-compatible HTTP embedder with retry/backoff - pkg/agent/hooks.go: lifecycle hook struct (OnContextBuild, OnPreTool, OnPostTool, OnPreLLM, OnPostMessage) ā€” nil checks only, no interface - pkg/tools/memory_search.go: memory_search tool for explicit semantic search over indexed memory Modified: - pkg/config: add memory.vector_search config section with sensible defaults (enabled: false, model: text-embedding-3-small, top-5) - pkg/agent/loop.go: integrate hooks at context build, tool execution, and post-message points - pkg/agent/context.go: accept optional enriched context in BuildMessages - cmd/picoclaw/main.go: wire vector search setup into agent and gateway commands with automatic initial indexing and incremental re-indexing Config example: { "memory": { "vector_search": { "enabled": true } } } API key and base URL fall back to providers.openai if not specified. Co-Authored-By: Claude Opus 4.6 --- cmd/picoclaw/main.go | 196 ++++++++++++++++++++++++++++++++ pkg/agent/context.go | 7 +- pkg/agent/hooks.go | 28 +++++ pkg/agent/loop.go | 41 ++++++- pkg/config/config.go | 24 ++++ pkg/tools/memory_search.go | 81 +++++++++++++ pkg/tools/memory_search_test.go | 109 ++++++++++++++++++ pkg/vecstore/chunker.go | 112 ++++++++++++++++++ pkg/vecstore/chunker_test.go | 104 +++++++++++++++++ pkg/vecstore/embed.go | 131 +++++++++++++++++++++ pkg/vecstore/store.go | 167 +++++++++++++++++++++++++++ pkg/vecstore/store_test.go | 150 ++++++++++++++++++++++++ 12 files changed, 1148 insertions(+), 2 deletions(-) create mode 100644 pkg/agent/hooks.go create mode 100644 pkg/tools/memory_search.go create mode 100644 pkg/tools/memory_search_test.go create mode 100644 pkg/vecstore/chunker.go create mode 100644 pkg/vecstore/chunker_test.go create mode 100644 pkg/vecstore/embed.go create mode 100644 pkg/vecstore/store.go create mode 100644 pkg/vecstore/store_test.go diff --git a/cmd/picoclaw/main.go b/cmd/picoclaw/main.go index 31d8dadeb..50ee8511f 100644 --- a/cmd/picoclaw/main.go +++ b/cmd/picoclaw/main.go @@ -31,6 +31,7 @@ import ( "github.com/sipeed/picoclaw/pkg/providers" "github.com/sipeed/picoclaw/pkg/skills" "github.com/sipeed/picoclaw/pkg/tools" + "github.com/sipeed/picoclaw/pkg/vecstore" "github.com/sipeed/picoclaw/pkg/voice" ) @@ -496,6 +497,10 @@ func agentCmd() { msgBus := bus.NewMessageBus() agentLoop := agent.NewAgentLoop(cfg, msgBus, provider) + // Setup vector search hooks (no-op when disabled) + hooks := setupVectorSearch(cfg, agentLoop) + agentLoop.SetHooks(hooks) + // Print agent startup info (only for interactive mode) startupInfo := agentLoop.GetStartupInfo() logger.InfoCF("agent", "Agent initialized", @@ -631,6 +636,10 @@ func gatewayCmd() { msgBus := bus.NewMessageBus() agentLoop := agent.NewAgentLoop(cfg, msgBus, provider) + // Setup vector search hooks (no-op when disabled) + hooks := setupVectorSearch(cfg, agentLoop) + agentLoop.SetHooks(hooks) + // Print agent startup info fmt.Println("\nšŸ“¦ Agent Status:") startupInfo := agentLoop.GetStartupInfo() @@ -1050,6 +1059,193 @@ func loadConfig() (*config.Config, error) { return config.LoadConfig(getConfigPath()) } +// setupVectorSearch initializes vector search hooks and tools when enabled. +// Returns the Hooks struct (always non-nil) and a save function to call on shutdown. +func setupVectorSearch(cfg *config.Config, agentLoop *agent.AgentLoop) *agent.Hooks { + hooks := &agent.Hooks{} + + vsCfg := cfg.Memory.VectorSearch + if !vsCfg.Enabled { + return hooks + } + + // Resolve API key and base URL, falling back to OpenAI provider config + apiKey := vsCfg.APIKey + if apiKey == "" { + apiKey = cfg.Providers.OpenAI.APIKey + } + apiBase := vsCfg.APIBase + if apiBase == "" { + apiBase = cfg.Providers.OpenAI.APIBase + } + if apiBase == "" { + apiBase = "https://api.openai.com/v1" + } + + model := vsCfg.Model + if model == "" { + model = "text-embedding-3-small" + } + + maxResults := vsCfg.MaxResults + if maxResults <= 0 { + maxResults = 5 + } + + chunkSize := vsCfg.ChunkSize + if chunkSize <= 0 { + chunkSize = 800 + } + + workspace := agentLoop.Workspace() + storePath := filepath.Join(workspace, "memory", ".vecstore.gob") + + embedder := vecstore.NewHTTPEmbedder(apiBase, apiKey, model) + store := vecstore.NewVectorStore(storePath) + + // Load existing store + if err := store.Load(); err != nil { + logger.ErrorCF("vecstore", "Failed to load vector store", + map[string]interface{}{"error": err.Error()}) + } + + // Initial indexing if store is empty + if store.Len() == 0 { + go indexMemoryFiles(context.Background(), workspace, embedder, store, chunkSize) + } + + // OnContextBuild: embed query and return top-K relevant chunks + hooks.OnContextBuild = func(ctx context.Context, query string) (string, error) { + embeddings, err := embedder.Embed(ctx, []string{query}) + if err != nil || len(embeddings) == 0 || len(embeddings[0]) == 0 { + return "", err + } + + results := store.Search(embeddings[0], maxResults) + if len(results) == 0 { + return "", nil + } + + var sb strings.Builder + for _, r := range results { + snippet := r.Text + if len(snippet) > 700 { + snippet = snippet[:700] + "..." + } + sb.WriteString(fmt.Sprintf("[%s | score: %.2f]\n%s\n\n", r.Source, r.Score, snippet)) + } + return sb.String(), nil + } + + // Track which files were touched by tool calls in this message + var touchedFiles []string + hooks.OnPostTool = func(_ context.Context, name string, result string, _ time.Duration) { + if name == "write_file" || name == "append_file" || name == "edit_file" { + // Check if the result mentions a memory/ path + if strings.Contains(result, "memory/") || strings.Contains(result, "memory\\") { + touchedFiles = append(touchedFiles, result) + } + } + } + + // OnPostMessage: re-index memory files that were modified + hooks.OnPostMessage = func(ctx context.Context, _, _, _ string) { + if len(touchedFiles) > 0 { + go indexMemoryFiles(ctx, workspace, embedder, store, chunkSize) + touchedFiles = nil + } + } + + // Register memory_search tool + agentLoop.RegisterTool(tools.NewMemorySearchTool(embedder, store, maxResults)) + + logger.InfoCF("vecstore", "Vector search enabled", + map[string]interface{}{ + "model": model, + "max_results": maxResults, + "chunk_size": chunkSize, + "store_path": storePath, + }) + + return hooks +} + +// indexMemoryFiles chunks and embeds all markdown files in the memory directory. +func indexMemoryFiles(ctx context.Context, workspace string, embedder vecstore.Embedder, store *vecstore.VectorStore, chunkSize int) { + memoryDir := filepath.Join(workspace, "memory") + + var allChunks []vecstore.Chunk + filepath.Walk(memoryDir, func(path string, info os.FileInfo, err error) error { + if err != nil || info.IsDir() { + return nil + } + // Skip non-markdown and the store file itself + if !strings.HasSuffix(path, ".md") { + return nil + } + + data, err := os.ReadFile(path) + if err != nil { + return nil + } + + relPath, _ := filepath.Rel(workspace, path) + chunks := vecstore.ChunkMarkdown(relPath, string(data), chunkSize) + allChunks = append(allChunks, chunks...) + return nil + }) + + if len(allChunks) == 0 { + return + } + + // Collect texts that need embedding (skip chunks already in store with same ID) + texts := make([]string, len(allChunks)) + for i, c := range allChunks { + texts[i] = c.Text + } + + // Embed in batches of 100 + const batchSize = 100 + for i := 0; i < len(texts); i += batchSize { + end := i + batchSize + if end > len(texts) { + end = len(texts) + } + + embeddings, err := embedder.Embed(ctx, texts[i:end]) + if err != nil { + logger.ErrorCF("vecstore", "Embedding batch failed", + map[string]interface{}{"error": err.Error(), "batch": i / batchSize}) + continue + } + + for j, emb := range embeddings { + allChunks[i+j].Embedding = emb + } + } + + // Filter out chunks that failed to embed + var valid []vecstore.Chunk + for _, c := range allChunks { + if len(c.Embedding) > 0 { + valid = append(valid, c) + } + } + + store.Upsert(valid) + if err := store.Save(); err != nil { + logger.ErrorCF("vecstore", "Failed to save vector store", + map[string]interface{}{"error": err.Error()}) + } + + logger.InfoCF("vecstore", "Memory indexed", + map[string]interface{}{ + "chunks_total": len(allChunks), + "chunks_embedded": len(valid), + }) +} + func cronCmd() { if len(os.Args) < 3 { cronHelp() diff --git a/pkg/agent/context.go b/pkg/agent/context.go index e32e456f9..448bb5901 100644 --- a/pkg/agent/context.go +++ b/pkg/agent/context.go @@ -157,11 +157,16 @@ func (cb *ContextBuilder) LoadBootstrapFiles() string { return result } -func (cb *ContextBuilder) BuildMessages(history []providers.Message, summary string, currentMessage string, media []string, channel, chatID string) []providers.Message { +func (cb *ContextBuilder) BuildMessages(history []providers.Message, summary string, currentMessage string, media []string, channel, chatID string, enrichedContext ...string) []providers.Message { messages := []providers.Message{} systemPrompt := cb.BuildSystemPrompt() + // Inject enriched context (e.g. vector search results) before session info + if len(enrichedContext) > 0 && enrichedContext[0] != "" { + systemPrompt += "\n\n## Relevant Memory Context\n\n" + enrichedContext[0] + } + // Add Current Session info if provided if channel != "" && chatID != "" { systemPrompt += fmt.Sprintf("\n\n## Current Session\nChannel: %s\nChat ID: %s", channel, chatID) diff --git a/pkg/agent/hooks.go b/pkg/agent/hooks.go new file mode 100644 index 000000000..f163d3104 --- /dev/null +++ b/pkg/agent/hooks.go @@ -0,0 +1,28 @@ +package agent + +import ( + "context" + "time" + + "github.com/sipeed/picoclaw/pkg/providers" +) + +// Hooks provides lifecycle integration points for the agent loop. +// All fields are optional ā€” nil checks only, zero cost when unused. +type Hooks struct { + // OnContextBuild is called before building messages to inject extra context. + // Returns additional context string to include in the system prompt. + OnContextBuild func(ctx context.Context, query string) (string, error) + + // OnPreTool is called before each tool execution. + OnPreTool func(ctx context.Context, name string, args map[string]interface{}) error + + // OnPostTool is called after each tool execution with the result and duration. + OnPostTool func(ctx context.Context, name string, result string, dur time.Duration) + + // OnPreLLM is called before each LLM call, allowing message mutation. + OnPreLLM func(ctx context.Context, messages []providers.Message) []providers.Message + + // OnPostMessage is called after a complete message exchange is saved. + OnPostMessage func(ctx context.Context, sessionKey, userMsg, response string) +} diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index fac2856e9..c0c2dad43 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -36,6 +36,7 @@ type AgentLoop struct { sessions *session.SessionManager contextBuilder *ContextBuilder tools *tools.ToolRegistry + hooks *Hooks running atomic.Bool summarizing sync.Map // Tracks which sessions are currently being summarized } @@ -148,6 +149,16 @@ func (al *AgentLoop) RegisterTool(tool tools.Tool) { al.tools.Register(tool) } +// SetHooks sets lifecycle hooks on the agent loop. +func (al *AgentLoop) SetHooks(h *Hooks) { + al.hooks = h +} + +// Workspace returns the agent's workspace path (for wiring hooks). +func (al *AgentLoop) Workspace() string { + return al.workspace +} + func (al *AgentLoop) ProcessDirect(ctx context.Context, content, sessionKey string) (string, error) { return al.ProcessDirectWithChannel(ctx, content, sessionKey, "cli", "direct") } @@ -236,7 +247,14 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, opts processOptions) (str // 1. Update tool contexts al.updateToolContexts(opts.Channel, opts.ChatID) - // 2. Build messages + // 2. Build messages (with optional enriched context from hooks) + var enrichedCtx string + if al.hooks != nil && al.hooks.OnContextBuild != nil { + if extra, err := al.hooks.OnContextBuild(ctx, opts.UserMessage); err == nil && extra != "" { + enrichedCtx = extra + } + } + history := al.sessions.GetHistory(opts.SessionKey) summary := al.sessions.GetSummary(opts.SessionKey) messages := al.contextBuilder.BuildMessages( @@ -246,6 +264,7 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, opts processOptions) (str nil, opts.Channel, opts.ChatID, + enrichedCtx, ) // 3. Save user message to session @@ -266,6 +285,11 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, opts processOptions) (str al.sessions.AddMessage(opts.SessionKey, "assistant", finalContent) al.sessions.Save(al.sessions.GetOrCreate(opts.SessionKey)) + // 6b. OnPostMessage hook + if al.hooks != nil && al.hooks.OnPostMessage != nil { + al.hooks.OnPostMessage(ctx, opts.SessionKey, opts.UserMessage, finalContent) + } + // 7. Optional: summarization if opts.EnableSummary { al.maybeSummarize(opts.SessionKey) @@ -411,10 +435,25 @@ func (al *AgentLoop) runLLMIteration(ctx context.Context, messages []providers.M "iteration": iteration, }) + // OnPreTool hook + if al.hooks != nil && al.hooks.OnPreTool != nil { + if err := al.hooks.OnPreTool(ctx, tc.Name, tc.Arguments); err != nil { + logger.DebugCF("agent", "OnPreTool hook error", + map[string]interface{}{"tool": tc.Name, "error": err.Error()}) + } + } + + toolStart := time.Now() result, err := al.tools.ExecuteWithContext(ctx, tc.Name, tc.Arguments, opts.Channel, opts.ChatID) if err != nil { result = fmt.Sprintf("Error: %v", err) } + toolDur := time.Since(toolStart) + + // OnPostTool hook + if al.hooks != nil && al.hooks.OnPostTool != nil { + al.hooks.OnPostTool(ctx, tc.Name, result, toolDur) + } toolResultMsg := providers.Message{ Role: "tool", diff --git a/pkg/config/config.go b/pkg/config/config.go index 56f1e1958..1b09aa919 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -49,9 +49,24 @@ type Config struct { Providers ProvidersConfig `json:"providers"` Gateway GatewayConfig `json:"gateway"` Tools ToolsConfig `json:"tools"` + Memory MemoryConfig `json:"memory"` mu sync.RWMutex } +type MemoryConfig struct { + VectorSearch VectorSearchConfig `json:"vector_search"` +} + +type VectorSearchConfig struct { + Enabled bool `json:"enabled"` + Provider string `json:"provider"` // default "openai" + Model string `json:"model"` // default "text-embedding-3-small" + APIKey string `json:"api_key"` // falls back to providers config + APIBase string `json:"api_base"` // falls back to providers config + MaxResults int `json:"max_results"` // default 5 + ChunkSize int `json:"chunk_size"` // default 800 +} + type AgentsConfig struct { Defaults AgentDefaults `json:"defaults"` } @@ -255,6 +270,15 @@ func DefaultConfig() *Config { }, }, }, + Memory: MemoryConfig{ + VectorSearch: VectorSearchConfig{ + Enabled: false, + Provider: "openai", + Model: "text-embedding-3-small", + MaxResults: 5, + ChunkSize: 800, + }, + }, } } diff --git a/pkg/tools/memory_search.go b/pkg/tools/memory_search.go new file mode 100644 index 000000000..27d462108 --- /dev/null +++ b/pkg/tools/memory_search.go @@ -0,0 +1,81 @@ +package tools + +import ( + "context" + "fmt" + "strings" + + "github.com/sipeed/picoclaw/pkg/vecstore" +) + +// MemorySearchTool searches memory using vector similarity. +type MemorySearchTool struct { + embedder vecstore.Embedder + store *vecstore.VectorStore + maxResults int +} + +// NewMemorySearchTool creates a memory search tool. +func NewMemorySearchTool(embedder vecstore.Embedder, store *vecstore.VectorStore, maxResults int) *MemorySearchTool { + if maxResults <= 0 { + maxResults = 5 + } + return &MemorySearchTool{ + embedder: embedder, + store: store, + maxResults: maxResults, + } +} + +func (t *MemorySearchTool) Name() string { return "memory_search" } + +func (t *MemorySearchTool) Description() string { + return "Search long-term memory for relevant information using semantic similarity. Use this to find specific memories, notes, or facts." +} + +func (t *MemorySearchTool) Parameters() map[string]interface{} { + return map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "query": map[string]interface{}{ + "type": "string", + "description": "The search query to find relevant memories", + }, + }, + "required": []string{"query"}, + } +} + +func (t *MemorySearchTool) Execute(ctx context.Context, args map[string]interface{}) (string, error) { + query, _ := args["query"].(string) + if query == "" { + return "", fmt.Errorf("query is required") + } + + // Embed the query + embeddings, err := t.embedder.Embed(ctx, []string{query}) + if err != nil { + return "", fmt.Errorf("embed query: %w", err) + } + if len(embeddings) == 0 || len(embeddings[0]) == 0 { + return "No results found.", nil + } + + // Search + results := t.store.Search(embeddings[0], t.maxResults) + if len(results) == 0 { + return "No relevant memories found.", nil + } + + // Format results + var sb strings.Builder + sb.WriteString(fmt.Sprintf("Found %d relevant memories:\n\n", len(results))) + for i, r := range results { + snippet := r.Text + if len(snippet) > 700 { + snippet = snippet[:700] + "..." + } + sb.WriteString(fmt.Sprintf("--- Result %d (score: %.2f, source: %s) ---\n%s\n\n", i+1, r.Score, r.Source, snippet)) + } + return sb.String(), nil +} diff --git a/pkg/tools/memory_search_test.go b/pkg/tools/memory_search_test.go new file mode 100644 index 000000000..3f7cc777d --- /dev/null +++ b/pkg/tools/memory_search_test.go @@ -0,0 +1,109 @@ +package tools + +import ( + "context" + "fmt" + "strings" + "testing" + "time" + + "github.com/sipeed/picoclaw/pkg/vecstore" +) + +// mockEmbedder returns a fixed embedding for any input. +type mockEmbedder struct { + embedding []float32 + err error +} + +func (m *mockEmbedder) Embed(_ context.Context, texts []string) ([][]float32, error) { + if m.err != nil { + return nil, m.err + } + result := make([][]float32, len(texts)) + for i := range texts { + result[i] = m.embedding + } + return result, nil +} + +func TestMemorySearchExecute(t *testing.T) { + store := vecstore.NewVectorStore("") + now := time.Now() + store.Upsert([]vecstore.Chunk{ + {ID: "a", Text: "The user prefers dark mode", Source: "memory/MEMORY.md", Embedding: []float32{1, 0, 0}, UpdatedAt: now}, + {ID: "b", Text: "Meeting notes from Monday", Source: "memory/202601/20260112.md", Embedding: []float32{0, 1, 0}, UpdatedAt: now}, + {ID: "c", Text: "User timezone is PST", Source: "memory/MEMORY.md", Embedding: []float32{0.9, 0.1, 0}, UpdatedAt: now}, + }) + + embedder := &mockEmbedder{embedding: []float32{1, 0, 0}} + tool := NewMemorySearchTool(embedder, store, 2) + + result, err := tool.Execute(context.Background(), map[string]interface{}{ + "query": "user preferences", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if !strings.Contains(result, "dark mode") { + t.Error("expected top result to contain 'dark mode'") + } + if !strings.Contains(result, "timezone") { + t.Error("expected second result to contain 'timezone'") + } + if strings.Contains(result, "Meeting notes") { + t.Error("should not contain third result (maxResults=2)") + } + if !strings.Contains(result, "score:") { + t.Error("result should include scores") + } + if !strings.Contains(result, "memory/MEMORY.md") { + t.Error("result should include source path") + } +} + +func TestMemorySearchEmptyQuery(t *testing.T) { + store := vecstore.NewVectorStore("") + embedder := &mockEmbedder{embedding: []float32{1, 0}} + tool := NewMemorySearchTool(embedder, store, 5) + + _, err := tool.Execute(context.Background(), map[string]interface{}{ + "query": "", + }) + if err == nil { + t.Error("expected error for empty query") + } +} + +func TestMemorySearchEmbedError(t *testing.T) { + store := vecstore.NewVectorStore("") + embedder := &mockEmbedder{err: fmt.Errorf("API unavailable")} + tool := NewMemorySearchTool(embedder, store, 5) + + _, err := tool.Execute(context.Background(), map[string]interface{}{ + "query": "test", + }) + if err == nil { + t.Error("expected error when embedder fails") + } + if !strings.Contains(err.Error(), "embed query") { + t.Errorf("error should wrap embed failure, got: %v", err) + } +} + +func TestMemorySearchNoResults(t *testing.T) { + store := vecstore.NewVectorStore("") // empty store + embedder := &mockEmbedder{embedding: []float32{1, 0}} + tool := NewMemorySearchTool(embedder, store, 5) + + result, err := tool.Execute(context.Background(), map[string]interface{}{ + "query": "anything", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !strings.Contains(result, "No relevant memories") { + t.Errorf("expected no-results message, got: %s", result) + } +} diff --git a/pkg/vecstore/chunker.go b/pkg/vecstore/chunker.go new file mode 100644 index 000000000..8cb3b96b2 --- /dev/null +++ b/pkg/vecstore/chunker.go @@ -0,0 +1,112 @@ +package vecstore + +import ( + "crypto/sha256" + "fmt" + "strings" + "time" +) + +// ChunkMarkdown splits markdown text into chunks at semantic boundaries. +// Splits first by ## headers, then sub-splits long sections by paragraphs. +// Each chunk gets a deterministic ID: sha256(source + ":" + text)[:12]. +func ChunkMarkdown(source, text string, maxChars int) []Chunk { + if maxChars <= 0 { + maxChars = 800 + } + + sections := splitByHeaders(text) + now := time.Now() + + var chunks []Chunk + for _, section := range sections { + section = strings.TrimSpace(section) + if section == "" { + continue + } + + if len(section) <= maxChars { + chunks = append(chunks, makeChunk(source, section, now)) + continue + } + + // Sub-split long sections by paragraphs + for _, part := range splitByParagraphs(section, maxChars) { + part = strings.TrimSpace(part) + if part == "" { + continue + } + chunks = append(chunks, makeChunk(source, part, now)) + } + } + return chunks +} + +// splitByHeaders splits text at ## header boundaries, keeping the header with its content. +func splitByHeaders(text string) []string { + lines := strings.Split(text, "\n") + var sections []string + var current strings.Builder + + for _, line := range lines { + if strings.HasPrefix(line, "## ") && current.Len() > 0 { + sections = append(sections, current.String()) + current.Reset() + } + current.WriteString(line) + current.WriteByte('\n') + } + if current.Len() > 0 { + sections = append(sections, current.String()) + } + return sections +} + +// splitByParagraphs splits text at double-newline boundaries, respecting maxChars. +func splitByParagraphs(text string, maxChars int) []string { + paragraphs := strings.Split(text, "\n\n") + var parts []string + var current strings.Builder + + for _, p := range paragraphs { + p = strings.TrimSpace(p) + if p == "" { + continue + } + + // If adding this paragraph would exceed max, flush current + if current.Len() > 0 && current.Len()+len(p)+2 > maxChars { + parts = append(parts, current.String()) + current.Reset() + } + + // If a single paragraph exceeds max, just add it as-is + if current.Len() == 0 && len(p) > maxChars { + parts = append(parts, p) + continue + } + + if current.Len() > 0 { + current.WriteString("\n\n") + } + current.WriteString(p) + } + if current.Len() > 0 { + parts = append(parts, current.String()) + } + return parts +} + +func makeChunk(source, text string, now time.Time) Chunk { + return Chunk{ + ID: chunkID(source, text), + Text: text, + Source: source, + UpdatedAt: now, + } +} + +func chunkID(source, text string) string { + h := sha256.Sum256([]byte(source + ":" + text)) + return fmt.Sprintf("%x", h[:6]) // 12 hex chars +} diff --git a/pkg/vecstore/chunker_test.go b/pkg/vecstore/chunker_test.go new file mode 100644 index 000000000..8f29a69d0 --- /dev/null +++ b/pkg/vecstore/chunker_test.go @@ -0,0 +1,104 @@ +package vecstore + +import ( + "strings" + "testing" +) + +func TestChunkMarkdownByHeaders(t *testing.T) { + md := `# Title + +Intro paragraph. + +## Section A + +Content A here. + +## Section B + +Content B here. +` + chunks := ChunkMarkdown("test.md", md, 800) + if len(chunks) < 2 { + t.Fatalf("expected at least 2 chunks, got %d", len(chunks)) + } + + // First chunk should contain "Title" and "Intro" + if !strings.Contains(chunks[0].Text, "Title") { + t.Errorf("first chunk should contain Title: %q", chunks[0].Text) + } + + // Should have sections A and B as separate chunks + foundA, foundB := false, false + for _, c := range chunks { + if strings.Contains(c.Text, "Section A") { + foundA = true + } + if strings.Contains(c.Text, "Section B") { + foundB = true + } + } + if !foundA || !foundB { + t.Errorf("expected sections A and B in separate chunks, foundA=%v foundB=%v", foundA, foundB) + } +} + +func TestChunkMarkdownLongSection(t *testing.T) { + // Create a long section that exceeds maxChars + long := "## Big Section\n\n" + for i := 0; i < 20; i++ { + long += "This is paragraph number " + string(rune('A'+i)) + ". It has some content.\n\n" + } + + chunks := ChunkMarkdown("test.md", long, 200) + if len(chunks) < 2 { + t.Fatalf("expected multiple chunks for long section, got %d", len(chunks)) + } + + for _, c := range chunks { + if c.Source != "test.md" { + t.Errorf("expected source 'test.md', got %q", c.Source) + } + if c.ID == "" { + t.Error("chunk ID should not be empty") + } + } +} + +func TestChunkMarkdownDeterministicIDs(t *testing.T) { + md := "## Hello\n\nWorld" + c1 := ChunkMarkdown("src.md", md, 800) + c2 := ChunkMarkdown("src.md", md, 800) + + if len(c1) != len(c2) { + t.Fatalf("chunk counts differ: %d vs %d", len(c1), len(c2)) + } + for i := range c1 { + if c1[i].ID != c2[i].ID { + t.Errorf("chunk %d: IDs differ %q vs %q", i, c1[i].ID, c2[i].ID) + } + } +} + +func TestChunkMarkdownEmpty(t *testing.T) { + chunks := ChunkMarkdown("test.md", "", 800) + if len(chunks) != 0 { + t.Errorf("expected 0 chunks for empty text, got %d", len(chunks)) + } +} + +func TestChunkIDUniqueness(t *testing.T) { + // Same text, different source ā†’ different ID + id1 := chunkID("a.md", "hello") + id2 := chunkID("b.md", "hello") + if id1 == id2 { + t.Error("IDs should differ for different sources") + } + + // Same source, different text ā†’ different ID + id3 := chunkID("a.md", "hello") + id4 := chunkID("a.md", "world") + if id3 == id4 { + t.Error("IDs should differ for different text") + } +} diff --git a/pkg/vecstore/embed.go b/pkg/vecstore/embed.go new file mode 100644 index 000000000..51a5d376c --- /dev/null +++ b/pkg/vecstore/embed.go @@ -0,0 +1,131 @@ +package vecstore + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "math" + "net/http" + "time" +) + +// Embedder generates embedding vectors from text. +type Embedder interface { + Embed(ctx context.Context, texts []string) ([][]float32, error) +} + +// HTTPEmbedder calls an OpenAI-compatible /v1/embeddings endpoint. +type HTTPEmbedder struct { + apiBase string + apiKey string + model string + client *http.Client +} + +// NewHTTPEmbedder creates an embedder targeting an OpenAI-compatible API. +func NewHTTPEmbedder(apiBase, apiKey, model string) *HTTPEmbedder { + return &HTTPEmbedder{ + apiBase: apiBase, + apiKey: apiKey, + model: model, + client: &http.Client{ + Timeout: 60 * time.Second, + }, + } +} + +type embeddingRequest struct { + Input []string `json:"input"` + Model string `json:"model"` +} + +type embeddingResponse struct { + Data []struct { + Embedding []float32 `json:"embedding"` + Index int `json:"index"` + } `json:"data"` + Error *struct { + Message string `json:"message"` + } `json:"error,omitempty"` +} + +// Embed sends all texts in one batch request and returns their embeddings. +// Retries up to 3 times with exponential backoff on transient errors. +func (e *HTTPEmbedder) Embed(ctx context.Context, texts []string) ([][]float32, error) { + if len(texts) == 0 { + return nil, nil + } + + body, err := json.Marshal(embeddingRequest{ + Input: texts, + Model: e.model, + }) + if err != nil { + return nil, fmt.Errorf("marshal embedding request: %w", err) + } + + url := e.apiBase + "/embeddings" + + const maxRetries = 3 + var lastErr error + for attempt := range maxRetries { + result, err := e.doRequest(ctx, url, body, len(texts)) + if err == nil { + return result, nil + } + lastErr = err + // Exponential backoff: 500ms, 2s, 8s + backoff := time.Duration(math.Pow(4, float64(attempt))) * 500 * time.Millisecond + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(backoff): + } + } + return nil, fmt.Errorf("embedding failed after %d retries: %w", maxRetries, lastErr) +} + +func (e *HTTPEmbedder) doRequest(ctx context.Context, url string, body []byte, n int) ([][]float32, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body)) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + if e.apiKey != "" { + req.Header.Set("Authorization", "Bearer "+e.apiKey) + } + + resp, err := e.client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("read response: %w", err) + } + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("embedding API %d: %s", resp.StatusCode, string(respBody)) + } + + var result embeddingResponse + if err := json.Unmarshal(respBody, &result); err != nil { + return nil, fmt.Errorf("decode response: %w", err) + } + if result.Error != nil { + return nil, fmt.Errorf("embedding API error: %s", result.Error.Message) + } + + // Order by index + embeddings := make([][]float32, n) + for _, d := range result.Data { + if d.Index < len(embeddings) { + embeddings[d.Index] = d.Embedding + } + } + return embeddings, nil +} diff --git a/pkg/vecstore/store.go b/pkg/vecstore/store.go new file mode 100644 index 000000000..da85cad9f --- /dev/null +++ b/pkg/vecstore/store.go @@ -0,0 +1,167 @@ +package vecstore + +import ( + "encoding/gob" + "math" + "os" + "path/filepath" + "sort" + "sync" + "time" +) + +// Chunk represents a text chunk with its embedding vector. +type Chunk struct { + ID string + Text string + Source string // file path the chunk came from + Embedding []float32 + UpdatedAt time.Time +} + +// Result is a search result with similarity score. +type Result struct { + Chunk + Score float32 +} + +// VectorStore is an in-memory vector store with gob persistence. +type VectorStore struct { + path string + chunks []Chunk + mu sync.RWMutex +} + +// NewVectorStore creates a store that persists to the given path. +func NewVectorStore(path string) *VectorStore { + return &VectorStore{path: path} +} + +// Load reads the store from disk. Returns nil if file doesn't exist. +func (vs *VectorStore) Load() error { + vs.mu.Lock() + defer vs.mu.Unlock() + + f, err := os.Open(vs.path) + if err != nil { + if os.IsNotExist(err) { + vs.chunks = nil + return nil + } + return err + } + defer f.Close() + + var chunks []Chunk + if err := gob.NewDecoder(f).Decode(&chunks); err != nil { + // Corrupt file ā€” start fresh + vs.chunks = nil + return nil + } + vs.chunks = chunks + return nil +} + +// Save writes the store to disk. +func (vs *VectorStore) Save() error { + vs.mu.RLock() + defer vs.mu.RUnlock() + + if err := os.MkdirAll(filepath.Dir(vs.path), 0755); err != nil { + return err + } + + f, err := os.Create(vs.path) + if err != nil { + return err + } + defer f.Close() + + return gob.NewEncoder(f).Encode(vs.chunks) +} + +// Search returns the top-K chunks most similar to the query embedding. +func (vs *VectorStore) Search(query []float32, topK int) []Result { + vs.mu.RLock() + defer vs.mu.RUnlock() + + if len(vs.chunks) == 0 { + return nil + } + + results := make([]Result, 0, len(vs.chunks)) + for _, c := range vs.chunks { + if len(c.Embedding) == 0 { + continue + } + score := cosine(query, c.Embedding) + results = append(results, Result{Chunk: c, Score: score}) + } + + sort.Slice(results, func(i, j int) bool { + return results[i].Score > results[j].Score + }) + + if topK > len(results) { + topK = len(results) + } + return results[:topK] +} + +// Upsert adds or replaces chunks by ID. +func (vs *VectorStore) Upsert(chunks []Chunk) { + vs.mu.Lock() + defer vs.mu.Unlock() + + idx := make(map[string]int, len(vs.chunks)) + for i, c := range vs.chunks { + idx[c.ID] = i + } + + for _, c := range chunks { + if i, ok := idx[c.ID]; ok { + vs.chunks[i] = c + } else { + vs.chunks = append(vs.chunks, c) + } + } +} + +// DeleteBySource removes all chunks from a given source. +func (vs *VectorStore) DeleteBySource(source string) { + vs.mu.Lock() + defer vs.mu.Unlock() + + filtered := vs.chunks[:0] + for _, c := range vs.chunks { + if c.Source != source { + filtered = append(filtered, c) + } + } + vs.chunks = filtered +} + +// Len returns the number of chunks in the store. +func (vs *VectorStore) Len() int { + vs.mu.RLock() + defer vs.mu.RUnlock() + return len(vs.chunks) +} + +// cosine computes cosine similarity between two vectors. +func cosine(a, b []float32) float32 { + if len(a) != len(b) || len(a) == 0 { + return 0 + } + var dot, normA, normB float64 + for i := range a { + dot += float64(a[i]) * float64(b[i]) + normA += float64(a[i]) * float64(a[i]) + normB += float64(b[i]) * float64(b[i]) + } + denom := math.Sqrt(normA) * math.Sqrt(normB) + if denom == 0 { + return 0 + } + return float32(dot / denom) +} diff --git a/pkg/vecstore/store_test.go b/pkg/vecstore/store_test.go new file mode 100644 index 000000000..04f800bc6 --- /dev/null +++ b/pkg/vecstore/store_test.go @@ -0,0 +1,150 @@ +package vecstore + +import ( + "os" + "path/filepath" + "testing" + "time" +) + +func TestCosine(t *testing.T) { + tests := []struct { + name string + a, b []float32 + want float32 + tol float32 + }{ + {"identical", []float32{1, 0, 0}, []float32{1, 0, 0}, 1.0, 0.001}, + {"orthogonal", []float32{1, 0, 0}, []float32{0, 1, 0}, 0.0, 0.001}, + {"opposite", []float32{1, 0}, []float32{-1, 0}, -1.0, 0.001}, + {"similar", []float32{1, 1}, []float32{1, 0.9}, 0.998, 0.01}, + {"empty", []float32{}, []float32{}, 0.0, 0.001}, + {"mismatched", []float32{1, 2}, []float32{1, 2, 3}, 0.0, 0.001}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := cosine(tt.a, tt.b) + if diff := got - tt.want; diff > tt.tol || diff < -tt.tol { + t.Errorf("cosine(%v, %v) = %f, want %f (tol %f)", tt.a, tt.b, got, tt.want, tt.tol) + } + }) + } +} + +func TestSearchReturnsTopK(t *testing.T) { + store := NewVectorStore("") + now := time.Now() + + store.Upsert([]Chunk{ + {ID: "a", Text: "alpha", Embedding: []float32{1, 0, 0}, UpdatedAt: now}, + {ID: "b", Text: "beta", Embedding: []float32{0, 1, 0}, UpdatedAt: now}, + {ID: "c", Text: "gamma", Embedding: []float32{0.9, 0.1, 0}, UpdatedAt: now}, + }) + + results := store.Search([]float32{1, 0, 0}, 2) + if len(results) != 2 { + t.Fatalf("expected 2 results, got %d", len(results)) + } + if results[0].ID != "a" { + t.Errorf("expected first result 'a', got %q", results[0].ID) + } + if results[1].ID != "c" { + t.Errorf("expected second result 'c', got %q", results[1].ID) + } +} + +func TestUpsertReplacesExisting(t *testing.T) { + store := NewVectorStore("") + now := time.Now() + + store.Upsert([]Chunk{ + {ID: "a", Text: "original", Embedding: []float32{1, 0}, UpdatedAt: now}, + }) + store.Upsert([]Chunk{ + {ID: "a", Text: "replaced", Embedding: []float32{0, 1}, UpdatedAt: now}, + }) + + if store.Len() != 1 { + t.Fatalf("expected 1 chunk, got %d", store.Len()) + } + + results := store.Search([]float32{0, 1}, 1) + if results[0].Text != "replaced" { + t.Errorf("expected replaced text, got %q", results[0].Text) + } +} + +func TestDeleteBySource(t *testing.T) { + store := NewVectorStore("") + now := time.Now() + + store.Upsert([]Chunk{ + {ID: "a", Text: "a", Source: "file1.md", Embedding: []float32{1, 0}, UpdatedAt: now}, + {ID: "b", Text: "b", Source: "file2.md", Embedding: []float32{0, 1}, UpdatedAt: now}, + {ID: "c", Text: "c", Source: "file1.md", Embedding: []float32{1, 1}, UpdatedAt: now}, + }) + + store.DeleteBySource("file1.md") + if store.Len() != 1 { + t.Fatalf("expected 1 chunk after delete, got %d", store.Len()) + } + + results := store.Search([]float32{0, 1}, 10) + if results[0].Source != "file2.md" { + t.Errorf("expected remaining chunk from file2.md, got %q", results[0].Source) + } +} + +func TestLoadSave(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "test.gob") + now := time.Now() + + // Save + store1 := NewVectorStore(path) + store1.Upsert([]Chunk{ + {ID: "x", Text: "hello", Source: "src", Embedding: []float32{0.5, 0.5}, UpdatedAt: now}, + }) + if err := store1.Save(); err != nil { + t.Fatalf("save: %v", err) + } + + // Load into new store + store2 := NewVectorStore(path) + if err := store2.Load(); err != nil { + t.Fatalf("load: %v", err) + } + if store2.Len() != 1 { + t.Fatalf("expected 1 chunk after load, got %d", store2.Len()) + } + + results := store2.Search([]float32{0.5, 0.5}, 1) + if results[0].Text != "hello" { + t.Errorf("expected 'hello', got %q", results[0].Text) + } +} + +func TestLoadMissingFile(t *testing.T) { + store := NewVectorStore(filepath.Join(t.TempDir(), "nonexistent.gob")) + if err := store.Load(); err != nil { + t.Fatalf("load missing file should not error: %v", err) + } + if store.Len() != 0 { + t.Fatalf("expected 0 chunks, got %d", store.Len()) + } +} + +func TestLoadCorruptFile(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "corrupt.gob") + os.WriteFile(path, []byte("not valid gob"), 0644) + + store := NewVectorStore(path) + if err := store.Load(); err != nil { + t.Fatalf("load corrupt file should not error: %v", err) + } + if store.Len() != 0 { + t.Fatalf("expected 0 chunks after corrupt load, got %d", store.Len()) + } +}