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) <noreply@anthropic.com>
This commit is contained in:
Tyler VoandClaude Opus 4.6 authored and viettranx committed 2026-03-21 13:11:24 +07:00
1 parent c50ce7c48e
commit ab09deecd7
2 files changed
+179 -9

No files matched your search

+79 -9
View File
@@ -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)