diff --git a/cmd/gateway_managed.go b/cmd/gateway_managed.go index 84508a14..84cff83e 100644 --- a/cmd/gateway_managed.go +++ b/cmd/gateway_managed.go @@ -157,7 +157,6 @@ func wireExtras( MediaStore: mediaStore, ModelPricing: appCfg.Telemetry.ModelPricing, TracingStore: stores.Tracing, - MemoryStore: stores.Memory, OnEvent: func(event agent.AgentEvent) { msgBus.Broadcast(bus.Event{ Name: protocol.EventAgent, diff --git a/internal/agent/loop_types.go b/internal/agent/loop_types.go index 4d3c8aa3..a7977070 100644 --- a/internal/agent/loop_types.go +++ b/internal/agent/loop_types.go @@ -136,10 +136,7 @@ type Loop struct { // Budget enforcement: monthly spending limit in cents (0 = unlimited) budgetMonthlyCents int - tracingStore store.TracingStore - - // Memory store for extractive memory fallback (writes directly when LLM flush fails) - memStore store.MemoryStore + tracingStore store.TracingStore } // AgentEvent is emitted during agent execution for WS broadcasting. @@ -256,10 +253,7 @@ type LoopConfig struct { // Budget enforcement BudgetMonthlyCents int - TracingStore store.TracingStore - - // Memory store for extractive memory fallback (writes directly when LLM flush fails) - MemoryStore store.MemoryStore + TracingStore store.TracingStore } const defaultMaxTokens = config.DefaultMaxTokens @@ -348,7 +342,6 @@ func NewLoop(cfg LoopConfig) *Loop { modelPricing: cfg.ModelPricing, budgetMonthlyCents: cfg.BudgetMonthlyCents, tracingStore: cfg.TracingStore, - memStore: cfg.MemoryStore, } } diff --git a/internal/agent/resolver.go b/internal/agent/resolver.go index 137fde47..548fe548 100644 --- a/internal/agent/resolver.go +++ b/internal/agent/resolver.go @@ -86,9 +86,6 @@ type ResolverDeps struct { // Tracing store for budget enforcement queries TracingStore store.TracingStore - - // Memory store for extractive memory fallback - MemoryStore store.MemoryStore } // NewManagedResolver creates a ResolverFunc that builds Loops from DB agent data. @@ -359,7 +356,6 @@ func NewManagedResolver(deps ResolverDeps) ResolverFunc { ModelPricing: deps.ModelPricing, BudgetMonthlyCents: derefInt(ag.BudgetMonthlyCents), TracingStore: deps.TracingStore, - MemoryStore: deps.MemoryStore, }) slog.Info("resolved agent from DB", "agent", agentKey, "model", ag.Model, "provider", ag.Provider) diff --git a/internal/store/context.go b/internal/store/context.go index ed307db5..dae3e9d3 100644 --- a/internal/store/context.go +++ b/internal/store/context.go @@ -25,7 +25,7 @@ const ( LocaleKey contextKey = "goclaw_locale" // SharedMemoryKey indicates memory should be shared (no per-user scoping). SharedMemoryKey contextKey = "goclaw_shared_memory" - // SharedKGKey indicates KG should be shared (query by team_id, not per-user). + // SharedKGKey indicates KG should be shared across all users of the agent (no per-user scoping). SharedKGKey contextKey = "goclaw_shared_kg" // ShellDenyGroupsKey holds per-agent shell deny group overrides. ShellDenyGroupsKey contextKey = "goclaw_shell_deny_groups" @@ -127,7 +127,16 @@ func MemoryUserID(ctx context.Context) string { return UserIDFromContext(ctx) } -// WithSharedKG returns a context flagged for shared knowledge graph (query by team_id). +// KGUserID returns the userID to use for knowledge graph operations. +// Returns "" (agent-level scope) when shared KG is active, otherwise the per-user ID. +func KGUserID(ctx context.Context) string { + if IsSharedKG(ctx) { + return "" + } + return UserIDFromContext(ctx) +} + +// WithSharedKG returns a context flagged for shared knowledge graph (agent-level, no per-user scoping). func WithSharedKG(ctx context.Context) context.Context { return context.WithValue(ctx, SharedKGKey, true) } diff --git a/internal/store/pg/knowledge_graph.go b/internal/store/pg/knowledge_graph.go index 09d48ac7..e7b75877 100644 --- a/internal/store/pg/knowledge_graph.go +++ b/internal/store/pg/knowledge_graph.go @@ -4,8 +4,10 @@ import ( "context" "database/sql" "encoding/json" + "cmp" "fmt" "log/slog" + "slices" "strings" "time" @@ -81,6 +83,13 @@ func (s *PGKnowledgeGraphStore) GetEntity(ctx context.Context, agentID, userID, func (s *PGKnowledgeGraphStore) DeleteEntity(ctx context.Context, agentID, userID, entityID string) error { aid := mustParseUUID(agentID) eid := mustParseUUID(entityID) + if store.IsSharedKG(ctx) { + _, err := s.db.ExecContext(ctx, + `DELETE FROM kg_entities WHERE id = $1 AND agent_id = $2`, + eid, aid, + ) + return err + } _, err := s.db.ExecContext(ctx, `DELETE FROM kg_entities WHERE id = $1 AND agent_id = $2 AND user_id = $3`, eid, aid, userID, @@ -99,7 +108,7 @@ func (s *PGKnowledgeGraphStore) ListEntities(ctx context.Context, agentID, userI where := "agent_id = $1" args := []any{aid} idx := 2 - if userID != "" { + if !store.IsSharedKG(ctx) && userID != "" { where += fmt.Sprintf(" AND user_id = $%d", idx) args = append(args, userID) idx++ @@ -238,14 +247,14 @@ func (s *PGKnowledgeGraphStore) vectorSearchEntities(ctx context.Context, embedd args = append(args, userID) idx++ } - args = append(args, vecStr, vecStr, limit) + args = append(args, vecStr, limit) q := fmt.Sprintf(` SELECT id, agent_id, user_id, external_id, name, entity_type, description, properties, source_id, confidence, created_at, updated_at, 1 - (embedding <=> $%d::vector) AS score FROM kg_entities WHERE %s - ORDER BY embedding <=> $%d::vector LIMIT $%d`, idx, where, idx+1, idx+2) + ORDER BY embedding <=> $%d::vector LIMIT $%d`, idx, where, idx, idx+1) rows, err := s.db.QueryContext(ctx, q, args...) if err != nil { @@ -304,14 +313,9 @@ func hybridMergeEntities(ilike, vec []scoredEntity, textWeight, vectorWeight flo scores[id] = entry.Score } - // Sort by score descending - for i := 0; i < len(results); i++ { - for j := i + 1; j < len(results); j++ { - if scores[results[j].ID] > scores[results[i].ID] { - results[i], results[j] = results[j], results[i] - } - } - } + slices.SortFunc(results, func(a, b store.Entity) int { + return cmp.Compare(scores[b.ID], scores[a.ID]) // descending + }) return results } @@ -341,6 +345,13 @@ func (s *PGKnowledgeGraphStore) UpsertRelation(ctx context.Context, relation *st func (s *PGKnowledgeGraphStore) DeleteRelation(ctx context.Context, agentID, userID, relationID string) error { aid := mustParseUUID(agentID) rid := mustParseUUID(relationID) + if store.IsSharedKG(ctx) { + _, err := s.db.ExecContext(ctx, + `DELETE FROM kg_relations WHERE id = $1 AND agent_id = $2`, + rid, aid, + ) + return err + } _, err := s.db.ExecContext(ctx, `DELETE FROM kg_relations WHERE id = $1 AND agent_id = $2 AND user_id = $3`, rid, aid, userID, @@ -388,7 +399,7 @@ func (s *PGKnowledgeGraphStore) ListAllRelations(ctx context.Context, agentID, u where := "agent_id = $1" args := []any{aid} idx := 2 - if userID != "" { + if !store.IsSharedKG(ctx) && userID != "" { where += fmt.Sprintf(" AND user_id = $%d", idx) args = append(args, userID) idx++ @@ -448,6 +459,33 @@ func (s *PGKnowledgeGraphStore) IngestExtraction(ctx context.Context, agentID, u extIDToUUID[e.ExternalID] = actualID } + // Batch-generate embeddings for all upserted entities (fire-and-forget on error). + if s.embProvider != nil && len(extIDToUUID) > 0 { + texts := make([]string, 0, len(entities)) + ids := make([]uuid.UUID, 0, len(entities)) + for _, e := range entities { + texts = append(texts, e.Name+" "+e.Description) + ids = append(ids, extIDToUUID[e.ExternalID]) + } + embeddings, embErr := s.embProvider.Embed(ctx, texts) + if embErr != nil { + slog.Warn("kg entity embedding batch failed", "error", embErr) + } else { + for i, emb := range embeddings { + if len(emb) == 0 { + continue + } + vecStr := vectorToString(emb) + if _, err := tx.ExecContext(ctx, + `UPDATE kg_entities SET embedding = $1::vector WHERE id = $2`, + vecStr, ids[i], + ); err != nil { + slog.Warn("kg entity embedding update failed", "entity_id", ids[i], "error", err) + } + } + } + } + for i := range relations { r := &relations[i] r.AgentID = agentID @@ -478,10 +516,19 @@ func (s *PGKnowledgeGraphStore) IngestExtraction(ctx context.Context, agentID, u func (s *PGKnowledgeGraphStore) PruneByConfidence(ctx context.Context, agentID, userID string, minConfidence float64) (int, error) { aid := mustParseUUID(agentID) - res, err := s.db.ExecContext(ctx, - `DELETE FROM kg_entities WHERE agent_id = $1 AND user_id = $2 AND confidence < $3`, - aid, userID, minConfidence, - ) + var res sql.Result + var err error + if store.IsSharedKG(ctx) { + res, err = s.db.ExecContext(ctx, + `DELETE FROM kg_entities WHERE agent_id = $1 AND confidence < $2`, + aid, minConfidence, + ) + } else { + res, err = s.db.ExecContext(ctx, + `DELETE FROM kg_entities WHERE agent_id = $1 AND user_id = $2 AND confidence < $3`, + aid, userID, minConfidence, + ) + } if err != nil { return 0, err } diff --git a/internal/store/pg/knowledge_graph_embedding.go b/internal/store/pg/knowledge_graph_embedding.go new file mode 100644 index 00000000..53cc3aa1 --- /dev/null +++ b/internal/store/pg/knowledge_graph_embedding.go @@ -0,0 +1,85 @@ +package pg + +import ( + "context" + "log/slog" + + "github.com/google/uuid" +) + +// BackfillKGEmbeddings generates embeddings for all KG entities that don't have one yet. +// Processes in batches of 50. Returns total number of entities updated. +func (s *PGKnowledgeGraphStore) BackfillKGEmbeddings(ctx context.Context) (int, error) { + if s.embProvider == nil { + return 0, nil + } + + const batchSize = 50 + total := 0 + + for { + rows, err := s.db.QueryContext(ctx, + `SELECT id, name, description FROM kg_entities + WHERE embedding IS NULL + ORDER BY created_at DESC + LIMIT $1`, batchSize) + if err != nil { + return total, err + } + + type entityRow struct { + id uuid.UUID + text string + } + var pending []entityRow + for rows.Next() { + var id uuid.UUID + var name, desc string + if err := rows.Scan(&id, &name, &desc); err != nil { + continue + } + pending = append(pending, entityRow{id: id, text: name + " " + desc}) + } + rows.Close() + + if len(pending) == 0 { + break + } + + slog.Info("backfilling KG entity embeddings", "batch", len(pending), "total_so_far", total) + + texts := make([]string, len(pending)) + for i, p := range pending { + texts[i] = p.text + } + embeddings, err := s.embProvider.Embed(ctx, texts) + if err != nil { + slog.Warn("kg entity embedding batch failed", "error", err) + break + } + + for i, emb := range embeddings { + if len(emb) == 0 { + continue + } + vecStr := vectorToString(emb) + if _, err := s.db.ExecContext(ctx, + `UPDATE kg_entities SET embedding = $1::vector WHERE id = $2`, + vecStr, pending[i].id, + ); err != nil { + slog.Warn("kg entity embedding update failed", "entity_id", pending[i].id, "error", err) + continue + } + total++ + } + + if len(pending) < batchSize { + break + } + } + + if total > 0 { + slog.Info("KG entity embeddings backfill complete", "updated", total) + } + return total, nil +} diff --git a/internal/tools/knowledge_graph.go b/internal/tools/knowledge_graph.go index 61c72991..d7d10382 100644 --- a/internal/tools/knowledge_graph.go +++ b/internal/tools/knowledge_graph.go @@ -65,7 +65,7 @@ func (t *KnowledgeGraphSearchTool) Execute(ctx context.Context, args map[string] if agentID == uuid.Nil { return ErrorResult("agent context not available") } - userID := store.MemoryUserID(ctx) + userID := store.KGUserID(ctx) query, _ := args["query"].(string) if query == "" { diff --git a/internal/tools/memory_interceptor.go b/internal/tools/memory_interceptor.go index e9f78e1e..f28b9977 100644 --- a/internal/tools/memory_interceptor.go +++ b/internal/tools/memory_interceptor.go @@ -182,10 +182,12 @@ func (m *MemoryInterceptor) WriteFile(ctx context.Context, path, content string, } } - // Trigger KG extraction in background if configured + // Trigger KG extraction in background if configured. + // Use KGUserID (not MemoryUserID) so shared KG entities go into agent-level scope. kgTriggered := false if m.kgExtractFn != nil && content != "" { - go m.kgExtractFn(context.WithoutCancel(ctx), agentStr, userID, content) + kgUserID := store.KGUserID(ctx) + go m.kgExtractFn(context.WithoutCancel(ctx), agentStr, kgUserID, content) kgTriggered = true }