From 876104950bb2880da0fe03f1683f137ba3ecf8f2 Mon Sep 17 00:00:00 2001 From: ntduc Date: Wed, 22 Jul 2026 19:01:16 +0700 Subject: [PATCH] fix: recover missing embeddings safely --- cmd/gateway_agents.go | 6 +- cmd/gateway_agents_test.go | 69 +++- cmd/gateway_setup.go | 45 +++ docs/00-architecture-overview.md | 13 +- internal/memory/embeddings.go | 46 ++- internal/memory/embeddings_provider_test.go | 52 +++ internal/store/pg/agents.go | 132 +++++-- internal/store/pg/embedding_backfill.go | 262 ++++++++++++++ .../integration/v3_embedding_backfill_test.go | 341 ++++++++++++++++++ 9 files changed, 914 insertions(+), 52 deletions(-) create mode 100644 internal/memory/embeddings_provider_test.go create mode 100644 internal/store/pg/embedding_backfill.go create mode 100644 tests/integration/v3_embedding_backfill_test.go diff --git a/cmd/gateway_agents.go b/cmd/gateway_agents.go index 2efec76a..13382311 100644 --- a/cmd/gateway_agents.go +++ b/cmd/gateway_agents.go @@ -51,8 +51,10 @@ func resolveEmbeddingProvider( } } - // 2. Auto-detect: scan DB providers for first with settings.embedding.enabled - allProviders, err := providerStore.ListAllProviders(context.Background()) + // 2. Auto-detect only from the master tenant. The selected provider receives + // cross-tenant semantic content during startup maintenance, so a tenant-owned + // endpoint must never be selected for this process-wide responsibility. + allProviders, err := providerStore.ListProviders(masterCtx) if err != nil { slog.Warn("failed to list providers for embedding auto-detect", "error", err) return nil diff --git a/cmd/gateway_agents_test.go b/cmd/gateway_agents_test.go index 9b867625..a72a3d9b 100644 --- a/cmd/gateway_agents_test.go +++ b/cmd/gateway_agents_test.go @@ -98,7 +98,12 @@ func captureEmbeddingRequest(t *testing.T, es *store.EmbeddingSettings) map[stri t.Fatalf("Decode() error = %v", err) } w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"data":[{"embedding":[0.1,0.2]}]}`)) + _ = json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{{ + "embedding": make([]float32, store.RequiredMemoryEmbeddingDimensions), + "index": 0, + }}, + }) })) defer server.Close() @@ -137,3 +142,65 @@ func TestBuildEmbeddingProviderIgnoresIncompatibleStoredDimensions(t *testing.T) t.Fatalf("dimensions = %v, want fallback 1536", got) } } + +type embeddingProviderStoreStub struct { + masterProviders []store.LLMProviderData + allProviders []store.LLMProviderData + listTenant uuid.UUID + listAllCalled bool +} + +func (s *embeddingProviderStoreStub) CreateProvider(context.Context, *store.LLMProviderData) error { + return nil +} +func (s *embeddingProviderStoreStub) GetProvider(context.Context, uuid.UUID) (*store.LLMProviderData, error) { + return nil, nil +} +func (s *embeddingProviderStoreStub) GetProviderByName(context.Context, string) (*store.LLMProviderData, error) { + return nil, nil +} +func (s *embeddingProviderStoreStub) ListProviders(ctx context.Context) ([]store.LLMProviderData, error) { + s.listTenant = store.TenantIDFromContext(ctx) + return s.masterProviders, nil +} +func (s *embeddingProviderStoreStub) ListAllProviders(context.Context) ([]store.LLMProviderData, error) { + s.listAllCalled = true + return s.allProviders, nil +} +func (s *embeddingProviderStoreStub) UpdateProvider(context.Context, uuid.UUID, map[string]any) error { + return nil +} +func (s *embeddingProviderStoreStub) DeleteProvider(context.Context, uuid.UUID) error { return nil } + +func TestResolveEmbeddingProviderAutoDetectUsesMasterTenantOnly(t *testing.T) { + embeddingSettings := json.RawMessage(`{"embedding":{"enabled":true}}`) + providerStore := &embeddingProviderStoreStub{ + masterProviders: []store.LLMProviderData{{ + TenantID: store.MasterTenantID, + Name: "master-embedding", + ProviderType: store.ProviderOpenAICompat, + APIBase: "http://master.invalid/v1", + APIKey: "master-key", + Enabled: true, + Settings: embeddingSettings, + }}, + allProviders: []store.LLMProviderData{{ + TenantID: uuid.New(), + Name: "tenant-controlled-endpoint", + ProviderType: store.ProviderOpenAICompat, + Enabled: true, + Settings: embeddingSettings, + }}, + } + + provider := resolveEmbeddingProvider(providerStore, nil, nil) + if provider == nil || provider.Name() != "master-embedding" { + t.Fatalf("resolveEmbeddingProvider() = %v, want master-embedding", provider) + } + if providerStore.listTenant != store.MasterTenantID { + t.Fatalf("ListProviders tenant = %s, want master tenant", providerStore.listTenant) + } + if providerStore.listAllCalled { + t.Fatal("ListAllProviders was called; cross-tenant auto-detect must stay disabled") + } +} diff --git a/cmd/gateway_setup.go b/cmd/gateway_setup.go index f1337f5c..cddb4481 100644 --- a/cmd/gateway_setup.go +++ b/cmd/gateway_setup.go @@ -420,16 +420,61 @@ func setupMemoryEmbeddings( } // Wire embedding provider into vault store for semantic document search. + var vaultStore *pg.PGVaultStore if pgStores.Vault != nil { pgStores.Vault.SetEmbeddingProvider(embProvider) slog.Info("vault embeddings enabled", "provider", embProvider.Name()) + vaultStore, _ = pgStores.Vault.(*pg.PGVaultStore) } // V3: Wire embedding provider into episodic store for semantic search. + var episodicStore *pg.PGEpisodicStore if pgStores.Episodic != nil { pgStores.Episodic.SetEmbeddingProvider(embProvider) slog.Info("episodic embeddings enabled", "provider", embProvider.Name()) + episodicStore, _ = pgStores.Episodic.(*pg.PGEpisodicStore) } + + // Agent create/update embedding hooks require the provider to be wired. + var agentStore *pg.PGAgentStore + if pgAgentStore, ok := pgStores.Agents.(*pg.PGAgentStore); ok { + agentStore = pgAgentStore + agentStore.SetEmbeddingProvider(embProvider) + slog.Info("agent embeddings enabled", "provider", embProvider.Name()) + } + + // Recover the remaining semantic indexes sequentially to avoid a burst + // of concurrent batch requests during gateway startup. Each surface gets + // its own deadline so a large agent backlog cannot starve later stores. + go func() { + if agentStore != nil { + bgCtx, cancel := context.WithTimeout(context.Background(), 10*time.Minute) + if count, err := agentStore.BackfillAgentEmbeddings(bgCtx); err != nil { + slog.Warn("agent embeddings backfill failed", "error", err) + } else if count > 0 { + slog.Info("agent embeddings recovery complete", "agents_updated", count) + } + cancel() + } + if episodicStore != nil { + bgCtx, cancel := context.WithTimeout(context.Background(), 10*time.Minute) + if count, err := episodicStore.BackfillEpisodicEmbeddings(bgCtx); err != nil { + slog.Warn("episodic embeddings backfill failed", "error", err) + } else if count > 0 { + slog.Info("episodic embeddings backfill complete", "summaries_updated", count) + } + cancel() + } + if vaultStore != nil { + bgCtx, cancel := context.WithTimeout(context.Background(), 10*time.Minute) + if count, err := vaultStore.BackfillVaultEmbeddings(bgCtx); err != nil { + slog.Warn("vault embeddings backfill failed", "error", err) + } else if count > 0 { + slog.Info("vault embeddings backfill complete", "documents_updated", count) + } + cancel() + } + }() } else { slog.Warn("memory embeddings disabled (no API key), chunks stored without vectors") } diff --git a/docs/00-architecture-overview.md b/docs/00-architecture-overview.md index 989b6d23..6b91091d 100644 --- a/docs/00-architecture-overview.md +++ b/docs/00-architecture-overview.md @@ -282,8 +282,8 @@ sequenceDiagram PG-->>GW: PG stores created GW->>GW: 5. Start tracing collector GW->>PG: 6. Register providers from DB - GW->>PG: 7. Wire embedding provider to PGMemoryStore - GW->>PG: 8. Backfill memory embeddings (background) + GW->>PG: 7. Wire embedding provider to vector-enabled PG stores + GW->>PG: 8. Backfill missing embeddings in bounded background batches GW->>GW: 9. Register config-based providers GW->>GW: 10. Create tool registry (filesystem, exec, web, memory, browser, TTS, subagent, MCP) @@ -303,6 +303,15 @@ sequenceDiagram GW->>Engine: 22. Listen on host:port ``` +Embedding recovery covers memory chunks, knowledge graph entities, active +skills, non-cancelled team tasks, active agents, episodic summaries, and vault +documents. The process-wide provider is resolved from the master tenant because +startup recovery may process semantic content from multiple tenants. Agent, +episodic, and vault backfills run sequentially with independent deadlines and +bounded batches. They only select rows whose embedding is `NULL`, re-check the +source fields before writing, isolate malformed rows, and can retry remaining +rows on the next startup. + --- ## 7. Database Wiring diff --git a/internal/memory/embeddings.go b/internal/memory/embeddings.go index d98e47b0..362873e1 100644 --- a/internal/memory/embeddings.go +++ b/internal/memory/embeddings.go @@ -10,6 +10,7 @@ import ( "math" "net/http" "strings" + "time" ) // ContentHash returns a short SHA256 hex digest of the content (first 16 bytes). @@ -144,6 +145,7 @@ type OpenAIEmbeddingProvider struct { apiKey string apiURL string dimensions int // optional: truncate output to this many dimensions (0 = use model default) + httpClient *http.Client } // NewOpenAIEmbeddingProvider creates a provider for OpenAI-compatible embedding APIs. @@ -156,10 +158,11 @@ func NewOpenAIEmbeddingProvider(name, apiKey, apiURL, model string) *OpenAIEmbed } return &OpenAIEmbeddingProvider{ - name: name, - model: model, - apiKey: apiKey, - apiURL: apiURL, + name: name, + model: model, + apiKey: apiKey, + apiURL: apiURL, + httpClient: &http.Client{Timeout: 60 * time.Second}, } } @@ -194,7 +197,7 @@ func (p *OpenAIEmbeddingProvider) Embed(ctx context.Context, texts []string) ([] req.Header.Set("Content-Type", "application/json") req.Header.Set("Authorization", "Bearer "+p.apiKey) - resp, err := http.DefaultClient.Do(req) + resp, err := p.httpClient.Do(req) if err != nil { return nil, fmt.Errorf("embedding request: %w", err) } @@ -208,6 +211,7 @@ func (p *OpenAIEmbeddingProvider) Embed(ctx context.Context, texts []string) ([] var result struct { Data []struct { Embedding []float32 `json:"embedding"` + Index *int `json:"index"` } `json:"data"` } @@ -215,9 +219,37 @@ func (p *OpenAIEmbeddingProvider) Embed(ctx context.Context, texts []string) ([] return nil, fmt.Errorf("decode response: %w", err) } + if len(result.Data) != len(texts) { + return nil, fmt.Errorf("embedding response count: got %d vectors for %d inputs", len(result.Data), len(texts)) + } + embeddings := make([][]float32, len(result.Data)) - for i, d := range result.Data { - embeddings[i] = d.Embedding + indexed := false + for _, d := range result.Data { + indexed = indexed || d.Index != nil + } + seen := make([]bool, len(result.Data)) + for responseIndex, d := range result.Data { + targetIndex := responseIndex + if indexed { + if d.Index == nil || *d.Index < 0 || *d.Index >= len(result.Data) || seen[*d.Index] { + return nil, fmt.Errorf("embedding response has invalid index at position %d", responseIndex) + } + targetIndex = *d.Index + } + seen[targetIndex] = true + if len(d.Embedding) == 0 { + return nil, fmt.Errorf("embedding response has empty vector at index %d", targetIndex) + } + if p.dimensions > 0 && len(d.Embedding) != p.dimensions { + return nil, fmt.Errorf("embedding response dimension at index %d: got %d, want %d", targetIndex, len(d.Embedding), p.dimensions) + } + for _, value := range d.Embedding { + if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) { + return nil, fmt.Errorf("embedding response has non-finite value at index %d", targetIndex) + } + } + embeddings[targetIndex] = d.Embedding } return embeddings, nil diff --git a/internal/memory/embeddings_provider_test.go b/internal/memory/embeddings_provider_test.go new file mode 100644 index 00000000..3b991b30 --- /dev/null +++ b/internal/memory/embeddings_provider_test.go @@ -0,0 +1,52 @@ +package memory + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" +) + +func TestOpenAIEmbeddingProviderRestoresResponseOrder(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + {"index": 1, "embedding": []float32{2, 2, 2}}, + {"index": 0, "embedding": []float32{1, 1, 1}}, + }, + }) + })) + defer server.Close() + + provider := NewOpenAIEmbeddingProvider("test", "key", server.URL, "model").WithDimensions(3) + embeddings, err := provider.Embed(context.Background(), []string{"first", "second"}) + if err != nil { + t.Fatalf("Embed() error = %v", err) + } + if embeddings[0][0] != 1 || embeddings[1][0] != 2 { + t.Fatalf("embeddings returned in response order: %v", embeddings) + } +} + +func TestOpenAIEmbeddingProviderRejectsWrongDimension(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"data":[{"index":0,"embedding":[1,2]}]}`)) + })) + defer server.Close() + + provider := NewOpenAIEmbeddingProvider("test", "key", server.URL, "model").WithDimensions(3) + if _, err := provider.Embed(context.Background(), []string{"text"}); err == nil { + t.Fatal("Embed() error = nil, want dimension validation error") + } +} + +func TestOpenAIEmbeddingProviderUsesBoundedHTTPClient(t *testing.T) { + provider := NewOpenAIEmbeddingProvider("test", "key", "http://example.invalid", "model") + if provider.httpClient == nil || provider.httpClient.Timeout != 60*time.Second { + t.Fatalf("HTTP timeout = %v, want 60s", provider.httpClient) + } +} diff --git a/internal/store/pg/agents.go b/internal/store/pg/agents.go index 7ee8040f..7cfd09a0 100644 --- a/internal/store/pg/agents.go +++ b/internal/store/pg/agents.go @@ -4,6 +4,7 @@ import ( "context" "database/sql" "encoding/json" + "errors" "fmt" "log/slog" "strings" @@ -18,10 +19,11 @@ import ( type PGAgentStore struct { db *sql.DB embProvider store.EmbeddingProvider // optional: for agent frontmatter embeddings + embJobs chan struct{} } func NewPGAgentStore(db *sql.DB) *PGAgentStore { - return &PGAgentStore{db: db} + return &PGAgentStore{db: db, embJobs: make(chan struct{}, 4)} } // SetEmbeddingProvider sets the embedding provider for agent frontmatter vectors. @@ -44,46 +46,96 @@ func (s *PGAgentStore) generateAgentEmbedding(ctx context.Context, agentID uuid. return } vecStr := vectorToString(embeddings[0]) - if _, err := s.db.ExecContext(ctx, `UPDATE agents SET embedding = $1::vector WHERE id = $2`, vecStr, agentID); err != nil { + if _, err := s.db.ExecContext(ctx, ` + UPDATE agents SET embedding = $1::vector + WHERE id = $2 AND deleted_at IS NULL AND status = 'active' + AND COALESCE(display_name, '') = $3 + AND COALESCE(frontmatter, '') = $4`, + vecStr, agentID, displayName, frontmatter); err != nil { slog.Warn("agent embedding update failed", "agent", agentID, "error", err) } } +// queueAgentEmbedding bounds background provider calls. Updates clear the old +// vector before enqueueing, so overload leaves a retryable NULL instead of a +// stale semantic index. +func (s *PGAgentStore) queueAgentEmbedding(agentID uuid.UUID) { + if s.embProvider == nil { + return + } + select { + case s.embJobs <- struct{}{}: + go func() { + defer func() { <-s.embJobs }() + ctx, cancel := context.WithTimeout(store.WithCrossTenant(context.Background()), 75*time.Second) + defer cancel() + agent, err := s.GetByID(ctx, agentID) + if err != nil { + slog.Warn("agent embedding source load failed", "agent", agentID, "error", err) + return + } + s.generateAgentEmbedding(ctx, agent.ID, agent.DisplayName, agent.Frontmatter) + }() + default: + slog.Warn("agent embedding queue full; vector remains pending", "agent", agentID) + } +} + // BackfillAgentEmbeddings generates embeddings for all active agents that have frontmatter but no embedding. func (s *PGAgentStore) BackfillAgentEmbeddings(ctx context.Context) (int, error) { if s.embProvider == nil { - return 0, nil - } - var pending []agentBackfillRow - if err := pkgSqlxDB.SelectContext(ctx, &pending, - `SELECT id, COALESCE(display_name, '') AS display_name, COALESCE(frontmatter, '') AS frontmatter - FROM agents WHERE deleted_at IS NULL AND frontmatter IS NOT NULL AND frontmatter != '' AND embedding IS NULL`, - ); err != nil { - return 0, err - } - if len(pending) == 0 { - return 0, nil + return 0, fmt.Errorf("no embedding provider configured") } - slog.Info("backfilling agent embeddings", "count", len(pending)) - updated := 0 - for _, ag := range pending { - text := ag.DisplayName - if ag.Frontmatter != "" { - text += ": " + ag.Frontmatter + const batchSize = 50 + total := 0 + var backfillErrs []error + cursor := uuid.Nil + for { + var pending []agentBackfillRow + if err := pkgSqlxDB.SelectContext(ctx, &pending, + `SELECT id, COALESCE(display_name, '') AS display_name, COALESCE(frontmatter, '') AS frontmatter + FROM agents + WHERE deleted_at IS NULL AND status = 'active' + AND frontmatter IS NOT NULL AND frontmatter != '' AND embedding IS NULL AND id > $1 + ORDER BY id ASC + LIMIT $2`, cursor, batchSize, + ); err != nil { + return total, fmt.Errorf("query agents without embeddings: %w", err) } - embeddings, err := s.embProvider.Embed(ctx, []string{text}) - if err != nil || len(embeddings) == 0 || len(embeddings[0]) == 0 { - continue + if len(pending) == 0 { + return total, errors.Join(backfillErrs...) } - vecStr := vectorToString(embeddings[0]) - if _, err := s.db.ExecContext(ctx, `UPDATE agents SET embedding = $1::vector WHERE id = $2`, vecStr, ag.ID); err != nil { - continue + + items := make([]embeddingBackfillItem[agentBackfillRow], len(pending)) + for i, agent := range pending { + items[i] = embeddingBackfillItem[agentBackfillRow]{ + row: agent, + id: agent.ID, + text: agent.DisplayName + ": " + agent.Frontmatter, + } } - updated++ + updated, err := processEmbeddingBackfillBatch(ctx, s.embProvider, "agent", items, + func(ctx context.Context, agent agentBackfillRow, embedding []float32) (int64, error) { + result, err := s.db.ExecContext(ctx, + `UPDATE agents SET embedding = $1::vector + WHERE id = $2 AND embedding IS NULL AND deleted_at IS NULL AND status = 'active' + AND COALESCE(display_name, '') = $3 AND COALESCE(frontmatter, '') = $4`, + vectorToString(embedding), agent.ID, agent.DisplayName, agent.Frontmatter) + if err != nil { + return 0, err + } + return result.RowsAffected() + }) + total += updated + if err != nil { + backfillErrs = append(backfillErrs, err) + if ctx.Err() != nil { + return total, errors.Join(backfillErrs...) + } + } + cursor = pending[len(pending)-1].ID } - slog.Info("agent embeddings backfill complete", "updated", updated) - return updated, nil } // agentSelectCols is the column list for all agent SELECT queries. @@ -155,9 +207,9 @@ func (s *PGAgentStore) Create(ctx context.Context, agent *store.AgentData) error return fmt.Errorf("commit agent create: %w", err) } - // Generate embedding for new agent with frontmatter (best-effort, outside tx) - if agent.Frontmatter != "" && s.embProvider != nil { - go s.generateAgentEmbedding(context.Background(), agent.ID, agent.DisplayName, agent.Frontmatter) + // Generate embedding for new agent with frontmatter (best-effort, outside tx). + if agent.Frontmatter != "" { + s.queueAgentEmbedding(agent.ID) } return nil } @@ -271,6 +323,13 @@ func (s *PGAgentStore) Update(ctx context.Context, id uuid.UUID, updates map[str } } + _, frontmatterChanged := updates["frontmatter"] + _, displayNameChanged := updates["display_name"] + if frontmatterChanged || displayNameChanged { + // Never keep a vector built from stale routing metadata. A bounded async + // job or the next startup backfill will fill this NULL. + updates["embedding"] = nil + } updates["updated_at"] = time.Now() if store.IsCrossTenant(ctx) { if err := execMapUpdateWhere(ctx, s.db, "agents", updates, "id = $IDX AND deleted_at IS NULL", id); err != nil { @@ -291,15 +350,8 @@ func (s *PGAgentStore) Update(ctx context.Context, id uuid.UUID, updates map[str } } - // Regenerate embedding when frontmatter changes - if _, hasFrontmatter := updates["frontmatter"]; hasFrontmatter && s.embProvider != nil { - bgCtx := store.WithTenantID(context.Background(), store.TenantIDFromContext(ctx)) - go func() { - ag, agErr := s.GetByID(bgCtx, id) - if agErr == nil { - s.generateAgentEmbedding(bgCtx, id, ag.DisplayName, ag.Frontmatter) - } - }() + if frontmatterChanged || displayNameChanged { + s.queueAgentEmbedding(id) } return nil } diff --git a/internal/store/pg/embedding_backfill.go b/internal/store/pg/embedding_backfill.go new file mode 100644 index 00000000..1b761585 --- /dev/null +++ b/internal/store/pg/embedding_backfill.go @@ -0,0 +1,262 @@ +package pg + +import ( + "context" + "errors" + "fmt" + "log/slog" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +const embeddingBackfillBatchSize = 50 + +type embeddingBackfillItem[T any] struct { + row T + id uuid.UUID + text string +} + +// processEmbeddingBackfillBatch uses an efficient batch request first, then +// isolates individual rows if the provider rejects or truncates the batch. A +// malformed row remains retryable without starving every later row forever. +func processEmbeddingBackfillBatch[T any]( + ctx context.Context, + provider store.EmbeddingProvider, + surface string, + items []embeddingBackfillItem[T], + update func(context.Context, T, []float32) (int64, error), +) (int, error) { + texts := make([]string, len(items)) + for i, item := range items { + texts[i] = item.text + } + + embeddings, batchErr := provider.Embed(ctx, texts) + if batchErr == nil { + batchErr = validateEmbeddingBatch(embeddings, len(items)) + } + if batchErr == nil { + return updateEmbeddingBackfillItems(ctx, surface, items, embeddings, update) + } + if ctx.Err() != nil { + return 0, fmt.Errorf("generate %s embeddings: %w", surface, ctx.Err()) + } + + slog.Warn("embedding backfill batch failed; retrying rows individually", + "surface", surface, "rows", len(items), "error", batchErr) + total := 0 + var rowErrs []error + for _, item := range items { + embedding, err := provider.Embed(ctx, []string{item.text}) + if err == nil { + err = validateEmbeddingBatch(embedding, 1) + } + if err != nil { + rowErrs = append(rowErrs, fmt.Errorf("generate %s embedding id=%s: %w", surface, item.id, err)) + if ctx.Err() != nil { + break + } + continue + } + updated, err := update(ctx, item.row, embedding[0]) + if err != nil { + rowErrs = append(rowErrs, fmt.Errorf("update %s embedding id=%s: %w", surface, item.id, err)) + continue + } + total += int(updated) + } + return total, errors.Join(rowErrs...) +} + +func validateEmbeddingBatch(embeddings [][]float32, expected int) error { + if len(embeddings) != expected { + return fmt.Errorf("got %d vectors for %d inputs", len(embeddings), expected) + } + for i, embedding := range embeddings { + if len(embedding) == 0 { + return fmt.Errorf("empty vector at index %d", i) + } + } + return nil +} + +func updateEmbeddingBackfillItems[T any]( + ctx context.Context, + surface string, + items []embeddingBackfillItem[T], + embeddings [][]float32, + update func(context.Context, T, []float32) (int64, error), +) (int, error) { + total := 0 + var updateErrs []error + for i, item := range items { + updated, err := update(ctx, item.row, embeddings[i]) + if err != nil { + updateErrs = append(updateErrs, fmt.Errorf("update %s embedding id=%s: %w", surface, item.id, err)) + continue + } + total += int(updated) + } + return total, errors.Join(updateErrs...) +} + +// BackfillVaultEmbeddings generates vectors for vault documents that were +// persisted while the embedding provider was unavailable. Documents without a +// summary still get a title/path vector; enrichment can later replace it with a +// richer title/path/summary vector. +func (s *PGVaultStore) BackfillVaultEmbeddings(ctx context.Context) (int, error) { + if s.embProvider == nil { + return 0, fmt.Errorf("no embedding provider configured") + } + + type vaultBackfillRow struct { + id uuid.UUID + title string + path string + summary string + } + + total := 0 + cursor := uuid.Nil + var backfillErrs []error + for { + rows, err := s.db.QueryContext(ctx, ` + SELECT id, title, path, COALESCE(summary, '') + FROM vault_documents + WHERE embedding IS NULL AND id > $1 + ORDER BY id ASC + LIMIT $2`, cursor, embeddingBackfillBatchSize) + if err != nil { + return total, fmt.Errorf("query vault documents without embeddings: %w", err) + } + + pending := make([]vaultBackfillRow, 0, embeddingBackfillBatchSize) + for rows.Next() { + var row vaultBackfillRow + if err := rows.Scan(&row.id, &row.title, &row.path, &row.summary); err != nil { + rows.Close() + return total, fmt.Errorf("scan vault document for embedding backfill: %w", err) + } + pending = append(pending, row) + } + if err := rows.Err(); err != nil { + rows.Close() + return total, fmt.Errorf("iterate vault documents for embedding backfill: %w", err) + } + if err := rows.Close(); err != nil { + return total, fmt.Errorf("close vault embedding backfill rows: %w", err) + } + if len(pending) == 0 { + return total, errors.Join(backfillErrs...) + } + + items := make([]embeddingBackfillItem[vaultBackfillRow], len(pending)) + for i, doc := range pending { + items[i] = embeddingBackfillItem[vaultBackfillRow]{ + row: doc, + id: doc.id, + text: doc.title + " " + doc.path + " " + doc.summary, + } + } + updated, err := processEmbeddingBackfillBatch(ctx, s.embProvider, "vault document", items, + func(ctx context.Context, doc vaultBackfillRow, embedding []float32) (int64, error) { + result, err := s.db.ExecContext(ctx, ` + UPDATE vault_documents SET embedding = $1::vector + WHERE id = $2 AND embedding IS NULL + AND title = $3 AND path = $4 AND COALESCE(summary, '') = $5`, + vectorToString(embedding), doc.id, doc.title, doc.path, doc.summary) + if err != nil { + return 0, err + } + return result.RowsAffected() + }) + total += updated + if err != nil { + backfillErrs = append(backfillErrs, err) + if ctx.Err() != nil { + return total, errors.Join(backfillErrs...) + } + } + cursor = pending[len(pending)-1].id + } +} + +// BackfillEpisodicEmbeddings generates vectors for unexpired summaries that +// were persisted while the embedding provider was unavailable. +func (s *PGEpisodicStore) BackfillEpisodicEmbeddings(ctx context.Context) (int, error) { + if s.embProvider == nil { + return 0, fmt.Errorf("no embedding provider configured") + } + + type episodicBackfillRow struct { + id uuid.UUID + summary string + } + + total := 0 + cursor := uuid.Nil + var backfillErrs []error + for { + rows, err := s.db.QueryContext(ctx, ` + SELECT id, summary + FROM episodic_summaries + WHERE embedding IS NULL AND summary != '' AND id > $1 + AND (expires_at IS NULL OR expires_at > NOW()) + ORDER BY id ASC + LIMIT $2`, cursor, embeddingBackfillBatchSize) + if err != nil { + return total, fmt.Errorf("query episodic summaries without embeddings: %w", err) + } + + pending := make([]episodicBackfillRow, 0, embeddingBackfillBatchSize) + for rows.Next() { + var row episodicBackfillRow + if err := rows.Scan(&row.id, &row.summary); err != nil { + rows.Close() + return total, fmt.Errorf("scan episodic summary for embedding backfill: %w", err) + } + pending = append(pending, row) + } + if err := rows.Err(); err != nil { + rows.Close() + return total, fmt.Errorf("iterate episodic summaries for embedding backfill: %w", err) + } + if err := rows.Close(); err != nil { + return total, fmt.Errorf("close episodic embedding backfill rows: %w", err) + } + if len(pending) == 0 { + return total, errors.Join(backfillErrs...) + } + + items := make([]embeddingBackfillItem[episodicBackfillRow], len(pending)) + for i, summary := range pending { + items[i] = embeddingBackfillItem[episodicBackfillRow]{ + row: summary, + id: summary.id, + text: summary.summary, + } + } + updated, err := processEmbeddingBackfillBatch(ctx, s.embProvider, "episodic summary", items, + func(ctx context.Context, summary episodicBackfillRow, embedding []float32) (int64, error) { + result, err := s.db.ExecContext(ctx, ` + UPDATE episodic_summaries SET embedding = $1::vector + WHERE id = $2 AND embedding IS NULL AND summary = $3 + AND (expires_at IS NULL OR expires_at > NOW())`, + vectorToString(embedding), summary.id, summary.summary) + if err != nil { + return 0, err + } + return result.RowsAffected() + }) + total += updated + if err != nil { + backfillErrs = append(backfillErrs, err) + if ctx.Err() != nil { + return total, errors.Join(backfillErrs...) + } + } + cursor = pending[len(pending)-1].id + } +} diff --git a/tests/integration/v3_embedding_backfill_test.go b/tests/integration/v3_embedding_backfill_test.go new file mode 100644 index 00000000..d68b0850 --- /dev/null +++ b/tests/integration/v3_embedding_backfill_test.go @@ -0,0 +1,341 @@ +//go:build integration + +package integration + +import ( + "context" + "database/sql" + "errors" + "strings" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/store" + "github.com/nextlevelbuilder/goclaw/internal/store/pg" +) + +func TestStoreVault_BackfillEmbeddings_Idempotent(t *testing.T) { + db := testDB(t) + tenantID, agentID := seedTenantAgent(t, db) + ctx := tenantCtx(tenantID) + storeWithoutProvider := pg.NewPGVaultStore(db) + + docWithSummary := makeVaultDoc(tenantID.String(), agentID.String(), "backfill/summary.md", "Recovery summary") + docWithSummary.Summary = "Semantic recovery content" + if err := storeWithoutProvider.UpsertDocument(ctx, docWithSummary); err != nil { + t.Fatalf("UpsertDocument with summary: %v", err) + } + docWithoutSummary := makeVaultDoc(tenantID.String(), agentID.String(), "backfill/path-only.md", "Recovery path") + if err := storeWithoutProvider.UpsertDocument(ctx, docWithoutSummary); err != nil { + t.Fatalf("UpsertDocument without summary: %v", err) + } + + storeWithoutProvider.SetEmbeddingProvider(newMockEmbedProvider()) + updated, err := storeWithoutProvider.BackfillVaultEmbeddings(ctx) + if err != nil { + t.Fatalf("BackfillVaultEmbeddings: %v", err) + } + if updated != 2 { + t.Fatalf("BackfillVaultEmbeddings updated = %d, want 2", updated) + } + assertEmbeddingPresent(t, db, "SELECT embedding IS NOT NULL FROM vault_documents WHERE id = $1", "vault document", docWithSummary.ID) + assertEmbeddingPresent(t, db, "SELECT embedding IS NOT NULL FROM vault_documents WHERE id = $1", "vault document", docWithoutSummary.ID) + + updated, err = storeWithoutProvider.BackfillVaultEmbeddings(ctx) + if err != nil { + t.Fatalf("BackfillVaultEmbeddings second call: %v", err) + } + if updated != 0 { + t.Fatalf("BackfillVaultEmbeddings second call updated = %d, want 0", updated) + } +} + +func TestStoreEpisodic_BackfillEmbeddings_Idempotent(t *testing.T) { + db := testDB(t) + tenantID, agentID := seedTenantAgent(t, db) + ctx := tenantCtx(tenantID) + storeWithoutProvider := pg.NewPGEpisodicStore(db) + episode := &store.EpisodicSummary{ + TenantID: tenantID, + AgentID: agentID, + UserID: "embedding-backfill-user", + SessionKey: "embedding-backfill-session", + Summary: "The user decided to restore all missing semantic vectors.", + SourceType: "session", + SourceID: "embedding-backfill-" + uuid.NewString(), + } + if err := storeWithoutProvider.Create(ctx, episode); err != nil { + t.Fatalf("Create episodic summary: %v", err) + } + expiredAt := time.Now().Add(-time.Hour) + expiredEpisode := &store.EpisodicSummary{ + TenantID: tenantID, + AgentID: agentID, + UserID: "embedding-backfill-user", + SessionKey: "expired-embedding-backfill-session", + Summary: "This expired summary must not consume provider capacity.", + SourceType: "session", + SourceID: "expired-embedding-backfill-" + uuid.NewString(), + ExpiresAt: &expiredAt, + } + if err := storeWithoutProvider.Create(ctx, expiredEpisode); err != nil { + t.Fatalf("Create expired episodic summary: %v", err) + } + + storeWithoutProvider.SetEmbeddingProvider(newMockEmbedProvider()) + updated, err := storeWithoutProvider.BackfillEpisodicEmbeddings(ctx) + if err != nil { + t.Fatalf("BackfillEpisodicEmbeddings: %v", err) + } + if updated != 1 { + t.Fatalf("BackfillEpisodicEmbeddings updated = %d, want 1", updated) + } + assertEmbeddingPresent(t, db, "SELECT embedding IS NOT NULL FROM episodic_summaries WHERE id = $1", "episodic summary", episode.ID.String()) + assertEmbeddingMissing(t, db, "SELECT embedding IS NULL FROM episodic_summaries WHERE id = $1", "expired episodic summary", expiredEpisode.ID.String()) + + updated, err = storeWithoutProvider.BackfillEpisodicEmbeddings(ctx) + if err != nil { + t.Fatalf("BackfillEpisodicEmbeddings second call: %v", err) + } + if updated != 0 { + t.Fatalf("BackfillEpisodicEmbeddings second call updated = %d, want 0", updated) + } +} + +func TestStoreAgent_BackfillEmbeddings_Idempotent(t *testing.T) { + db := testDB(t) + tenantID, agentID := seedTenantAgent(t, db) + ctx := tenantCtx(tenantID) + agentStore := pg.NewPGAgentStore(db) + + if err := agentStore.Update(ctx, agentID, map[string]any{ + "display_name": "Recovery agent", + "frontmatter": "An agent with semantic routing metadata", + }); err != nil { + t.Fatalf("Update agent frontmatter: %v", err) + } + + agentStore.SetEmbeddingProvider(newMockEmbedProvider()) + updated, err := agentStore.BackfillAgentEmbeddings(ctx) + if err != nil { + t.Fatalf("BackfillAgentEmbeddings: %v", err) + } + if updated != 1 { + t.Fatalf("BackfillAgentEmbeddings updated = %d, want 1", updated) + } + assertEmbeddingPresent(t, db, "SELECT embedding IS NOT NULL FROM agents WHERE id = $1", "agent", agentID.String()) + + updated, err = agentStore.BackfillAgentEmbeddings(ctx) + if err != nil { + t.Fatalf("BackfillAgentEmbeddings second call: %v", err) + } + if updated != 0 { + t.Fatalf("BackfillAgentEmbeddings second call updated = %d, want 0", updated) + } +} + +func TestStoreAgent_AsyncEmbeddingDoesNotOverwriteNewerSource(t *testing.T) { + db := testDB(t) + tenantID, agentID := seedTenantAgent(t, db) + ctx := tenantCtx(tenantID) + provider := newBlockingAgentEmbedProvider() + agentStore := pg.NewPGAgentStore(db) + agentStore.SetEmbeddingProvider(provider) + + if err := agentStore.Update(ctx, agentID, map[string]any{ + "display_name": "Versioned agent", + "frontmatter": "old routing metadata", + }); err != nil { + t.Fatalf("Update old agent source: %v", err) + } + select { + case <-provider.oldStarted: + case <-time.After(3 * time.Second): + t.Fatal("old embedding request did not start") + } + + if err := agentStore.Update(ctx, agentID, map[string]any{ + "frontmatter": "new routing metadata", + }); err != nil { + t.Fatalf("Update new agent source: %v", err) + } + assertAgentEmbeddingEquals(t, db, agentID, vectorLiteral(0.2)) + + close(provider.releaseOld) + select { + case <-provider.oldDone: + case <-time.After(3 * time.Second): + t.Fatal("old embedding request did not finish") + } + assertAgentEmbeddingEquals(t, db, agentID, vectorLiteral(0.2)) +} + +func TestStoreVault_BackfillEmbeddings_RejectsIncompleteBatch(t *testing.T) { + db := testDB(t) + tenantID, agentID := seedTenantAgent(t, db) + ctx := tenantCtx(tenantID) + vaultStore := pg.NewPGVaultStore(db) + doc := makeVaultDoc(tenantID.String(), agentID.String(), "backfill/incomplete.md", "Incomplete response") + doc.Summary = "This row must remain pending when the provider drops a vector." + if err := vaultStore.UpsertDocument(ctx, doc); err != nil { + t.Fatalf("UpsertDocument: %v", err) + } + + vaultStore.SetEmbeddingProvider(incompleteEmbedProvider{}) + if _, err := vaultStore.BackfillVaultEmbeddings(ctx); err == nil { + t.Fatal("BackfillVaultEmbeddings error = nil, want incomplete batch error") + } + assertEmbeddingMissing(t, db, "SELECT embedding IS NULL FROM vault_documents WHERE id = $1", "vault document", doc.ID) +} + +func TestStoreVault_BackfillEmbeddings_PoisonRowDoesNotStarveLaterRows(t *testing.T) { + db := testDB(t) + tenantID, agentID := seedTenantAgent(t, db) + ctx := tenantCtx(tenantID) + vaultStore := pg.NewPGVaultStore(db) + + poison := makeVaultDoc(tenantID.String(), agentID.String(), "backfill/poison.md", "Poison row") + poison.Summary = "poison embedding input" + valid := makeVaultDoc(tenantID.String(), agentID.String(), "backfill/valid.md", "Valid row") + valid.Summary = "valid embedding input" + if err := vaultStore.UpsertDocument(ctx, poison); err != nil { + t.Fatalf("Upsert poison document: %v", err) + } + if err := vaultStore.UpsertDocument(ctx, valid); err != nil { + t.Fatalf("Upsert valid document: %v", err) + } + + vaultStore.SetEmbeddingProvider(poisonEmbedProvider{delegate: newMockEmbedProvider()}) + updated, err := vaultStore.BackfillVaultEmbeddings(ctx) + if err == nil { + t.Fatal("BackfillVaultEmbeddings error = nil, want poison-row error") + } + if updated != 1 { + t.Fatalf("BackfillVaultEmbeddings updated = %d, want 1 valid row", updated) + } + assertEmbeddingMissing(t, db, "SELECT embedding IS NULL FROM vault_documents WHERE id = $1", "poison vault document", poison.ID) + assertEmbeddingPresent(t, db, "SELECT embedding IS NOT NULL FROM vault_documents WHERE id = $1", "valid vault document", valid.ID) +} + +func assertEmbeddingPresent(t *testing.T, db queryRower, query, surface, id string) { + t.Helper() + var present bool + if err := db.QueryRow(query, id).Scan(&present); err != nil { + t.Fatalf("query %s embedding: %v", surface, err) + } + if !present { + t.Fatalf("expected %s embedding to be present", surface) + } +} + +func assertEmbeddingMissing(t *testing.T, db queryRower, query, surface, id string) { + t.Helper() + var missing bool + if err := db.QueryRow(query, id).Scan(&missing); err != nil { + t.Fatalf("query %s embedding: %v", surface, err) + } + if !missing { + t.Fatalf("expected %s embedding to remain missing", surface) + } +} + +type queryRower interface { + QueryRow(query string, args ...any) *sql.Row +} + +type incompleteEmbedProvider struct{} + +func (incompleteEmbedProvider) Name() string { return "incomplete" } +func (incompleteEmbedProvider) Model() string { return "incomplete" } +func (incompleteEmbedProvider) Embed(context.Context, []string) ([][]float32, error) { + return nil, nil +} + +var _ store.EmbeddingProvider = incompleteEmbedProvider{} + +type poisonEmbedProvider struct { + delegate store.EmbeddingProvider +} + +func (p poisonEmbedProvider) Name() string { return "poison" } +func (p poisonEmbedProvider) Model() string { return "poison" } +func (p poisonEmbedProvider) Embed(ctx context.Context, texts []string) ([][]float32, error) { + for _, text := range texts { + if strings.Contains(text, "poison") { + return nil, errors.New("rejected poison input") + } + } + return p.delegate.Embed(ctx, texts) +} + +var _ store.EmbeddingProvider = poisonEmbedProvider{} + +type blockingAgentEmbedProvider struct { + oldStarted chan struct{} + releaseOld chan struct{} + oldDone chan struct{} + startOnce sync.Once + doneOnce sync.Once +} + +func newBlockingAgentEmbedProvider() *blockingAgentEmbedProvider { + return &blockingAgentEmbedProvider{ + oldStarted: make(chan struct{}), + releaseOld: make(chan struct{}), + oldDone: make(chan struct{}), + } +} + +func (p *blockingAgentEmbedProvider) Name() string { return "blocking-agent" } +func (p *blockingAgentEmbedProvider) Model() string { return "blocking-agent" } +func (p *blockingAgentEmbedProvider) Embed(ctx context.Context, texts []string) ([][]float32, error) { + value := float32(0.2) + if strings.Contains(texts[0], "old routing metadata") { + value = 0.1 + p.startOnce.Do(func() { close(p.oldStarted) }) + select { + case <-p.releaseOld: + case <-ctx.Done(): + return nil, ctx.Err() + } + p.doneOnce.Do(func() { close(p.oldDone) }) + } + result := make([][]float32, len(texts)) + for i := range result { + result[i] = make([]float32, store.RequiredMemoryEmbeddingDimensions) + result[i][0] = value + } + return result, nil +} + +func assertAgentEmbeddingEquals(t *testing.T, db queryRower, agentID uuid.UUID, expected string) { + t.Helper() + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + var equal bool + if err := db.QueryRow(`SELECT COALESCE(embedding = $1::vector, false) FROM agents WHERE id = $2`, expected, agentID).Scan(&equal); err != nil { + t.Fatalf("query agent embedding: %v", err) + } + if equal { + return + } + time.Sleep(10 * time.Millisecond) + } + t.Fatal("agent embedding did not match the newest source") +} + +func vectorLiteral(first float64) string { + values := make([]string, store.RequiredMemoryEmbeddingDimensions) + values[0] = "0.2" + if first == 0.1 { + values[0] = "0.1" + } + for i := 1; i < len(values); i++ { + values[i] = "0" + } + return "[" + strings.Join(values, ",") + "]" +} + +var _ store.EmbeddingProvider = (*blockingAgentEmbedProvider)(nil)