From ab09deecd7b5d77a75ac272b8317708d1fadd6c2 Mon Sep 17 00:00:00 2001 From: Tyler Vo Date: Sat, 21 Mar 2026 13:22:51 +1100 Subject: [PATCH] perf: wire embedding_cache table into IndexDocument hot path The embedding_cache table existed in the schema but was never queried. IndexDocument() now batch-checks the cache before calling the embedding API, writes fresh results back, and merges cached + fresh vectors. Graceful fallback: if cache lookup fails, falls through to full API call. Closes #326 Co-Authored-By: Claude Opus 4.6 (1M context) --- internal/store/pg/memory_docs.go | 88 +++++++++++++++-- internal/store/pg/memory_embedding_cache.go | 100 ++++++++++++++++++++ 2 files changed, 179 insertions(+), 9 deletions(-) create mode 100644 internal/store/pg/memory_embedding_cache.go diff --git a/internal/store/pg/memory_docs.go b/internal/store/pg/memory_docs.go index ccacb5a7..48672e50 100644 --- a/internal/store/pg/memory_docs.go +++ b/internal/store/pg/memory_docs.go @@ -178,18 +178,88 @@ func (s *PGMemoryStore) IndexDocument(ctx context.Context, agentID, userID, path return nil } - // Generate embeddings + // Generate embeddings with cache var embeddings [][]float32 if s.provider != nil { - texts := make([]string, len(chunks)) + providerName := s.provider.Name() + providerModel := s.provider.Model() + + // Compute content hashes for all chunks + hashes := make([]string, len(chunks)) for i, c := range chunks { - texts[i] = c.Text + hashes[i] = memory.ContentHash(c.Text) } - var embErr error - embeddings, embErr = s.provider.Embed(ctx, texts) - if embErr != nil { - slog.Warn("memory embedding failed, storing chunks without vectors", - "path", path, "chunks", len(chunks), "error", embErr) + + // Batch lookup cached embeddings + cached, cacheErr := s.lookupEmbeddingCache(ctx, hashes, providerName, providerModel) + if cacheErr != nil { + slog.Warn("embedding cache lookup failed, falling back to full API call", + "path", path, "error", cacheErr) + cached = nil + } + + // Determine which chunks need fresh embeddings + var uncachedIdxs []int + var uncachedTexts []string + for i, c := range chunks { + if cached != nil { + if _, ok := cached[hashes[i]]; ok { + continue + } + } + uncachedIdxs = append(uncachedIdxs, i) + uncachedTexts = append(uncachedTexts, c.Text) + } + + if len(cached) > 0 { + slog.Info("embedding cache hit", + "path", path, "cached", len(cached), "uncached", len(uncachedTexts)) + } + + // Call embedding API only for uncached texts + var freshEmbeddings [][]float32 + if len(uncachedTexts) > 0 { + var embErr error + freshEmbeddings, embErr = s.provider.Embed(ctx, uncachedTexts) + if embErr != nil { + slog.Warn("memory embedding failed, storing chunks without vectors", + "path", path, "chunks", len(chunks), "error", embErr) + } + } + + // Write fresh embeddings back to cache + if len(freshEmbeddings) > 0 { + var cacheEntries []embeddingCacheEntry + for j, emb := range freshEmbeddings { + if j < len(uncachedIdxs) { + cacheEntries = append(cacheEntries, embeddingCacheEntry{ + Hash: hashes[uncachedIdxs[j]], + Embedding: emb, + }) + } + } + if writeErr := s.writeEmbeddingCache(ctx, cacheEntries, providerName, providerModel); writeErr != nil { + slog.Warn("embedding cache write failed", "path", path, "error", writeErr) + } + } + + // Merge cached + fresh embeddings into final slice + if cached != nil || freshEmbeddings != nil { + embeddings = make([][]float32, len(chunks)) + // Fill from cache + for i, h := range hashes { + if cached != nil { + if emb, ok := cached[h]; ok { + embeddings[i] = emb + } + } + } + // Fill from fresh + for j, idx := range uncachedIdxs { + if j < len(freshEmbeddings) { + embeddings[idx] = freshEmbeddings[j] + } + } } } @@ -204,7 +274,7 @@ func (s *PGMemoryStore) IndexDocument(ctx context.Context, agentID, userID, path uid = &userID } - if embeddings != nil && i < len(embeddings) { + if embeddings != nil && i < len(embeddings) && embeddings[i] != nil { // Insert with embedding via raw SQL (pgvector) s.db.ExecContext(ctx, `INSERT INTO memory_chunks (id, agent_id, document_id, user_id, path, start_line, end_line, hash, text, embedding, updated_at) diff --git a/internal/store/pg/memory_embedding_cache.go b/internal/store/pg/memory_embedding_cache.go new file mode 100644 index 00000000..fb56a12d --- /dev/null +++ b/internal/store/pg/memory_embedding_cache.go @@ -0,0 +1,100 @@ +package pg + +import ( + "context" + "fmt" + "log/slog" + "strconv" + "strings" + "time" + + "github.com/lib/pq" +) + +// embeddingCacheEntry holds data for a single cache row. +type embeddingCacheEntry struct { + Hash string + Embedding []float32 +} + +// lookupEmbeddingCache fetches cached embeddings for the given content hashes. +// Returns a map from hash -> embedding vector. Missing hashes are simply absent. +func (s *PGMemoryStore) lookupEmbeddingCache(ctx context.Context, hashes []string, provider, model string) (map[string][]float32, error) { + if len(hashes) == 0 { + return nil, nil + } + + rows, err := s.db.QueryContext(ctx, + `SELECT hash, embedding FROM embedding_cache WHERE hash = ANY($1) AND provider = $2 AND model = $3`, + pq.Array(hashes), provider, model, + ) + if err != nil { + return nil, fmt.Errorf("lookup embedding cache: %w", err) + } + defer rows.Close() + + result := make(map[string][]float32, len(hashes)) + for rows.Next() { + var hash, vecStr string + if err := rows.Scan(&hash, &vecStr); err != nil { + slog.Warn("embedding cache scan error", "error", err) + continue + } + vec, err := parseVector(vecStr) + if err != nil { + slog.Warn("embedding cache parse error", "hash", hash, "error", err) + continue + } + result[hash] = vec + } + return result, rows.Err() +} + +// writeEmbeddingCache batch-upserts embedding cache entries. +func (s *PGMemoryStore) writeEmbeddingCache(ctx context.Context, entries []embeddingCacheEntry, provider, model string) error { + if len(entries) == 0 { + return nil + } + + now := time.Now() + for _, e := range entries { + dims := len(e.Embedding) + vecStr := vectorToString(e.Embedding) + _, err := s.db.ExecContext(ctx, + `INSERT INTO embedding_cache (hash, provider, model, embedding, dims, created_at, updated_at) + VALUES ($1, $2, $3, $4::vector, $5, $6, $6) + ON CONFLICT (hash, provider, model) + DO UPDATE SET embedding = EXCLUDED.embedding, dims = EXCLUDED.dims, updated_at = EXCLUDED.updated_at`, + e.Hash, provider, model, vecStr, dims, now, + ) + if err != nil { + return fmt.Errorf("write embedding cache hash=%s: %w", e.Hash, err) + } + } + return nil +} + +// parseVector converts a pgvector string like "[0.1,0.2,0.3]" into []float32. +func parseVector(s string) ([]float32, error) { + s = strings.TrimSpace(s) + if len(s) < 2 { + return nil, fmt.Errorf("vector string too short: %q", s) + } + // Strip surrounding brackets + s = strings.TrimPrefix(s, "[") + s = strings.TrimSuffix(s, "]") + if s == "" { + return nil, nil + } + + parts := strings.Split(s, ",") + vec := make([]float32, 0, len(parts)) + for _, p := range parts { + f, err := strconv.ParseFloat(strings.TrimSpace(p), 32) + if err != nil { + return nil, fmt.Errorf("parse vector element %q: %w", p, err) + } + vec = append(vec, float32(f)) + } + return vec, nil +}