diff --git a/internal/tools/knowledge_graph.go b/internal/tools/knowledge_graph.go index d7d10382..ab749b99 100644 --- a/internal/tools/knowledge_graph.go +++ b/internal/tools/knowledge_graph.go @@ -3,6 +3,7 @@ package tools import ( "context" "fmt" + "log/slog" "strings" "github.com/google/uuid" @@ -80,7 +81,7 @@ func (t *KnowledgeGraphSearchTool) Execute(ctx context.Context, args map[string] // Traversal mode: entity_id provided if entityID != "" { - return t.executeTraversal(ctx, agentID.String(), userID, entityID, maxDepth) + return t.executeTraversal(ctx, agentID.String(), userID, entityID, maxDepth, query) } // List-all mode: query="*" @@ -92,31 +93,71 @@ func (t *KnowledgeGraphSearchTool) Execute(ctx context.Context, args map[string] return t.executeSearch(ctx, agentID.String(), userID, query, args) } -func (t *KnowledgeGraphSearchTool) executeTraversal(ctx context.Context, agentID, userID, entityID string, maxDepth int) *Result { +func (t *KnowledgeGraphSearchTool) executeTraversal(ctx context.Context, agentID, userID, entityID string, maxDepth int, query string) *Result { + // Tier 1: outgoing deep traversal results, err := t.kgStore.Traverse(ctx, agentID, userID, entityID, maxDepth) if err != nil { return ErrorResult(fmt.Sprintf("graph traversal failed: %v", err)) } - if len(results) == 0 { - return NewResult(fmt.Sprintf("No connected entities found from entity_id=%q within depth %d.", entityID, maxDepth)) + if len(results) > 0 { + const maxTraversalResults = 20 + totalResults := len(results) + if totalResults > maxTraversalResults { + results = results[:maxTraversalResults] + } + var sb strings.Builder + sb.WriteString(fmt.Sprintf("Graph traversal from %q (max depth %d):\n\n", entityID, maxDepth)) + for _, r := range results { + sb.WriteString(fmt.Sprintf("- [depth %d] %s (%s)", r.Depth, r.Entity.Name, r.Entity.EntityType)) + if r.Via != "" { + sb.WriteString(fmt.Sprintf(" via %q", r.Via)) + } + if r.Entity.Description != "" { + sb.WriteString(fmt.Sprintf("\n %s", r.Entity.Description)) + } + if len(r.Path) > 0 { + sb.WriteString(fmt.Sprintf("\n path: %s", strings.Join(r.Path, " → "))) + } + sb.WriteString("\n") + } + if totalResults > maxTraversalResults { + sb.WriteString(fmt.Sprintf("\n(+%d more entities reachable, use query to narrow or adjust max_depth)\n", totalResults-maxTraversalResults)) + } + return NewResult(sb.String()) } - var sb strings.Builder - sb.WriteString(fmt.Sprintf("Graph traversal from %q (max depth %d):\n\n", entityID, maxDepth)) - for _, r := range results { - sb.WriteString(fmt.Sprintf("- [depth %d] %s (%s)", r.Depth, r.Entity.Name, r.Entity.EntityType)) - if r.Via != "" { - sb.WriteString(fmt.Sprintf(" via %q", r.Via)) - } - if r.Entity.Description != "" { - sb.WriteString(fmt.Sprintf("\n %s", r.Entity.Description)) - } - if len(r.Path) > 0 { - sb.WriteString(fmt.Sprintf("\n path: %s", strings.Join(r.Path, " → "))) - } - sb.WriteString("\n") + // Tier 2: direct connections (bidirectional, 1-hop, cap 10) + relations, relErr := t.kgStore.ListRelations(ctx, agentID, userID, entityID) + if relErr != nil { + slog.Warn("kg.listRelations failed", "entity_id", entityID, "error", relErr) } - return NewResult(sb.String()) + if len(relations) > 0 { + const maxDirectConnections = 10 + totalCount := len(relations) + if totalCount > maxDirectConnections { + relations = relations[:maxDirectConnections] + } + nameCache := make(map[string]string) + entityName := t.resolveEntityName(ctx, agentID, userID, entityID, nameCache) + var sb strings.Builder + sb.WriteString(fmt.Sprintf("Direct connections of %q:\n\n", entityName)) + for _, rel := range relations { + srcName := t.resolveEntityName(ctx, agentID, userID, rel.SourceEntityID, nameCache) + tgtName := t.resolveEntityName(ctx, agentID, userID, rel.TargetEntityID, nameCache) + sb.WriteString(fmt.Sprintf(" %s —[%s]→ %s\n", srcName, rel.RelationType, tgtName)) + } + if totalCount > maxDirectConnections { + sb.WriteString(fmt.Sprintf("\n(%d more connections not shown)\n", totalCount-maxDirectConnections)) + } + return NewResult(sb.String()) + } + + // Tier 3: fallback to search if query provided + if query != "" && query != "*" { + return t.executeSearch(ctx, agentID, userID, query, nil) + } + + return NewResult(fmt.Sprintf("No connected entities found from entity_id=%q.", entityID)) } func (t *KnowledgeGraphSearchTool) executeListAll(ctx context.Context, agentID, userID string) *Result { @@ -179,15 +220,20 @@ func (t *KnowledgeGraphSearchTool) executeSearch(ctx context.Context, agentID, u sb.WriteString(fmt.Sprintf(" %s\n", e.Description)) } - // Fetch relations to show connections with names + // Fetch relations to show connections with names (cap 5 per entity) relations, err := t.kgStore.ListRelations(ctx, agentID, userID, e.ID) if err == nil && len(relations) > 0 { + const maxRelationsPerEntity = 5 sb.WriteString(" Relations:\n") - for _, rel := range relations { + shown := min(len(relations), maxRelationsPerEntity) + for _, rel := range relations[:shown] { srcName := t.resolveEntityName(ctx, agentID, userID, rel.SourceEntityID, entityNames) tgtName := t.resolveEntityName(ctx, agentID, userID, rel.TargetEntityID, entityNames) sb.WriteString(fmt.Sprintf(" %s —[%s]→ %s\n", srcName, rel.RelationType, tgtName)) } + if len(relations) > maxRelationsPerEntity { + sb.WriteString(fmt.Sprintf(" (+%d more, use entity_id=%q to see all)\n", len(relations)-maxRelationsPerEntity, e.ID)) + } } } return NewResult(sb.String()) diff --git a/internal/tools/knowledge_graph_test.go b/internal/tools/knowledge_graph_test.go new file mode 100644 index 00000000..306948fd --- /dev/null +++ b/internal/tools/knowledge_graph_test.go @@ -0,0 +1,399 @@ +package tools + +import ( + "context" + "fmt" + "strings" + "testing" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +// ── mock KG store ────────────────────────────────────────────────── + +type mockKGStore struct { + entities map[string]store.Entity // keyed by entity ID + relations []store.Relation // all relations + traversal map[string][]store.TraversalResult // keyed by start entity ID +} + +func newMockKGStore() *mockKGStore { + return &mockKGStore{ + entities: make(map[string]store.Entity), + traversal: make(map[string][]store.TraversalResult), + } +} + +func (m *mockKGStore) UpsertEntity(_ context.Context, e *store.Entity) error { + m.entities[e.ID] = *e + return nil +} + +func (m *mockKGStore) GetEntity(_ context.Context, _, _, entityID string) (*store.Entity, error) { + if e, ok := m.entities[entityID]; ok { + return &e, nil + } + return nil, fmt.Errorf("entity not found") +} + +func (m *mockKGStore) DeleteEntity(context.Context, string, string, string) error { return nil } + +func (m *mockKGStore) ListEntities(_ context.Context, _, _ string, opts store.EntityListOptions) ([]store.Entity, error) { + var out []store.Entity + for _, e := range m.entities { + out = append(out, e) + if opts.Limit > 0 && len(out) >= opts.Limit { + break + } + } + return out, nil +} + +func (m *mockKGStore) SearchEntities(_ context.Context, _, _, query string, limit int) ([]store.Entity, error) { + var out []store.Entity + q := strings.ToLower(query) + for _, e := range m.entities { + if strings.Contains(strings.ToLower(e.Name), q) || strings.Contains(strings.ToLower(e.Description), q) { + out = append(out, e) + if limit > 0 && len(out) >= limit { + break + } + } + } + return out, nil +} + +func (m *mockKGStore) UpsertRelation(_ context.Context, r *store.Relation) error { + m.relations = append(m.relations, *r) + return nil +} + +func (m *mockKGStore) DeleteRelation(context.Context, string, string, string) error { return nil } + +func (m *mockKGStore) ListRelations(_ context.Context, _, _, entityID string) ([]store.Relation, error) { + var out []store.Relation + for _, r := range m.relations { + if r.SourceEntityID == entityID || r.TargetEntityID == entityID { + out = append(out, r) + } + } + return out, nil +} + +func (m *mockKGStore) ListAllRelations(context.Context, string, string, int) ([]store.Relation, error) { + return m.relations, nil +} + +func (m *mockKGStore) Traverse(_ context.Context, _, _, startEntityID string, _ int) ([]store.TraversalResult, error) { + return m.traversal[startEntityID], nil +} + +func (m *mockKGStore) IngestExtraction(context.Context, string, string, []store.Entity, []store.Relation) error { + return nil +} + +func (m *mockKGStore) PruneByConfidence(context.Context, string, string, float64) (int, error) { + return 0, nil +} + +func (m *mockKGStore) Stats(context.Context, string, string) (*store.GraphStats, error) { + return &store.GraphStats{}, nil +} + +func (m *mockKGStore) SetEmbeddingProvider(store.EmbeddingProvider) {} +func (m *mockKGStore) Close() error { return nil } + +// ── test helpers ─────────────────────────────────────────────────── + +var ( + testAgentID = uuid.New() + testUserID = "test-user" +) + +func kgContext() context.Context { + ctx := context.Background() + ctx = store.WithAgentID(ctx, testAgentID) + ctx = store.WithUserID(ctx, testUserID) + return ctx +} + +// setupBaseGraph creates the shared test graph: +// +// A(Viettx) →[owns]→ B(GoClaw) →[implements]→ C(Dầu thô) →[related_to]→ D(Kuwait) +// A(Viettx) →[manages]→ C(Dầu thô) +// E(Chiến sự Trung Đông) — isolated, no relations +func setupBaseGraph() (*mockKGStore, map[string]string) { + ms := newMockKGStore() + ids := map[string]string{ + "A": uuid.NewString(), + "B": uuid.NewString(), + "C": uuid.NewString(), + "D": uuid.NewString(), + "E": uuid.NewString(), + } + + entities := []store.Entity{ + {ID: ids["A"], AgentID: testAgentID.String(), UserID: testUserID, Name: "Viettx", EntityType: "person"}, + {ID: ids["B"], AgentID: testAgentID.String(), UserID: testUserID, Name: "GoClaw", EntityType: "project"}, + {ID: ids["C"], AgentID: testAgentID.String(), UserID: testUserID, Name: "Dầu thô", EntityType: "concept"}, + {ID: ids["D"], AgentID: testAgentID.String(), UserID: testUserID, Name: "Kuwait", EntityType: "location"}, + {ID: ids["E"], AgentID: testAgentID.String(), UserID: testUserID, Name: "Chiến sự Trung Đông", EntityType: "event"}, + } + for i := range entities { + ms.entities[entities[i].ID] = entities[i] + } + + ms.relations = []store.Relation{ + {ID: uuid.NewString(), AgentID: testAgentID.String(), UserID: testUserID, SourceEntityID: ids["A"], RelationType: "owns", TargetEntityID: ids["B"]}, + {ID: uuid.NewString(), AgentID: testAgentID.String(), UserID: testUserID, SourceEntityID: ids["A"], RelationType: "manages", TargetEntityID: ids["C"]}, + {ID: uuid.NewString(), AgentID: testAgentID.String(), UserID: testUserID, SourceEntityID: ids["C"], RelationType: "related_to", TargetEntityID: ids["D"]}, + {ID: uuid.NewString(), AgentID: testAgentID.String(), UserID: testUserID, SourceEntityID: ids["B"], RelationType: "implements", TargetEntityID: ids["C"]}, + } + + // Pre-compute traversal results (mock outgoing-only behavior) + // A → B, C (outgoing from A) + ms.traversal[ids["A"]] = []store.TraversalResult{ + {Entity: entities[1], Depth: 1, Via: "owns"}, // B=GoClaw + {Entity: entities[2], Depth: 1, Via: "manages"}, // C=Dầu thô + } + // B → C (outgoing from B) + ms.traversal[ids["B"]] = []store.TraversalResult{ + {Entity: entities[2], Depth: 1, Via: "implements"}, // C=Dầu thô + } + // C → D (outgoing from C) + ms.traversal[ids["C"]] = []store.TraversalResult{ + {Entity: entities[3], Depth: 1, Via: "related_to"}, // D=Kuwait + } + // D has no outgoing → empty traversal + // E is isolated → empty traversal + + return ms, ids +} + +// ── tests ────────────────────────────────────────────────────────── + +func TestKGTraversal_Tier1_OutgoingEdges(t *testing.T) { + ms, ids := setupBaseGraph() + tool := NewKnowledgeGraphSearchTool() + tool.SetKGStore(ms) + + ctx := kgContext() + result := tool.executeTraversal(ctx, testAgentID.String(), testUserID, ids["A"], 2, "Viettx") + text := result.ForLLM + + if !strings.Contains(text, "GoClaw") { + t.Error("expected result to contain 'GoClaw'") + } + if !strings.Contains(text, "Dầu thô") { + t.Error("expected result to contain 'Dầu thô'") + } + if strings.Contains(text, "Direct connections") { + t.Error("tier 1 should NOT show 'Direct connections'") + } +} + +func TestKGTraversal_Tier2_OnlyIncomingEdges(t *testing.T) { + ms, ids := setupBaseGraph() + tool := NewKnowledgeGraphSearchTool() + tool.SetKGStore(ms) + + ctx := kgContext() + // D=Kuwait has 0 outgoing, 1 incoming (C→D) + result := tool.executeTraversal(ctx, testAgentID.String(), testUserID, ids["D"], 2, "Kuwait") + text := result.ForLLM + + if !strings.Contains(text, "Direct connections") { + t.Error("expected tier 2 'Direct connections' section") + } + if !strings.Contains(text, "Dầu thô") { + t.Error("expected to see 'Dầu thô' in direct connections") + } + if !strings.Contains(text, "—[related_to]→") { + t.Error("expected relation format '—[related_to]→'") + } +} + +func TestKGTraversal_Tier3_IsolatedWithQuery(t *testing.T) { + ms, ids := setupBaseGraph() + tool := NewKnowledgeGraphSearchTool() + tool.SetKGStore(ms) + + ctx := kgContext() + // E=Chiến sự TĐ: 0 outgoing, 0 incoming, but searchable by name + result := tool.executeTraversal(ctx, testAgentID.String(), testUserID, ids["E"], 2, "Chiến sự") + text := result.ForLLM + + if !strings.Contains(text, "Chiến sự Trung Đông") { + t.Error("expected tier 3 fallback to find 'Chiến sự Trung Đông' via search") + } + if !strings.Contains(text, "Found") { + t.Error("expected search result format with 'Found'") + } +} + +func TestKGTraversal_Tier3_IsolatedNoQuery(t *testing.T) { + ms, ids := setupBaseGraph() + tool := NewKnowledgeGraphSearchTool() + tool.SetKGStore(ms) + + ctx := kgContext() + // E=isolated, no query fallback + result := tool.executeTraversal(ctx, testAgentID.String(), testUserID, ids["E"], 2, "") + text := result.ForLLM + + if !strings.Contains(text, "No connected entities found") { + t.Errorf("expected 'No connected entities found', got: %s", text) + } +} + +func TestKGTraversal_Tier2_CappedAt10(t *testing.T) { + ms := newMockKGStore() + entityX := uuid.NewString() + ms.entities[entityX] = store.Entity{ID: entityX, AgentID: testAgentID.String(), UserID: testUserID, Name: "HubNode", EntityType: "concept"} + + // Create 15 incoming relations to X + for i := 0; i < 15; i++ { + srcID := uuid.NewString() + srcName := fmt.Sprintf("Source_%02d", i) + ms.entities[srcID] = store.Entity{ID: srcID, AgentID: testAgentID.String(), UserID: testUserID, Name: srcName, EntityType: "concept"} + ms.relations = append(ms.relations, store.Relation{ + ID: uuid.NewString(), AgentID: testAgentID.String(), UserID: testUserID, + SourceEntityID: srcID, RelationType: "connects_to", TargetEntityID: entityX, + }) + } + // X has no outgoing → empty traversal + + tool := NewKnowledgeGraphSearchTool() + tool.SetKGStore(ms) + + ctx := kgContext() + result := tool.executeTraversal(ctx, testAgentID.String(), testUserID, entityX, 2, "") + text := result.ForLLM + + count := strings.Count(text, "—[connects_to]→") + if count > 10 { + t.Errorf("expected max 10 direct connections, got %d", count) + } + if count == 0 { + t.Error("expected at least 1 direct connection") + } +} + +func TestKGTraversal_Tier1_SkipsTier2WhenTraversalHasResults(t *testing.T) { + ms, ids := setupBaseGraph() + tool := NewKnowledgeGraphSearchTool() + tool.SetKGStore(ms) + + ctx := kgContext() + // B=GoClaw has 1 outgoing (B→C) and 1 incoming (A→B) + result := tool.executeTraversal(ctx, testAgentID.String(), testUserID, ids["B"], 2, "GoClaw") + text := result.ForLLM + + if !strings.Contains(text, "Dầu thô") { + t.Error("expected traversal to contain 'Dầu thô' (outgoing from B)") + } + if strings.Contains(text, "Direct connections") { + t.Error("tier 1 should NOT fall through to tier 2") + } +} + +func TestKGTraversal_Tier2_RelationFormat(t *testing.T) { + ms := newMockKGStore() + idF := uuid.NewString() + idG := uuid.NewString() + idH := uuid.NewString() + + ms.entities[idF] = store.Entity{ID: idF, AgentID: testAgentID.String(), UserID: testUserID, Name: "NodeF", EntityType: "concept"} + ms.entities[idG] = store.Entity{ID: idG, AgentID: testAgentID.String(), UserID: testUserID, Name: "NodeG", EntityType: "concept"} + ms.entities[idH] = store.Entity{ID: idH, AgentID: testAgentID.String(), UserID: testUserID, Name: "NodeH", EntityType: "concept"} + + // F→G (outgoing from F) and H→F (incoming to F) + ms.relations = []store.Relation{ + {ID: uuid.NewString(), AgentID: testAgentID.String(), UserID: testUserID, SourceEntityID: idF, RelationType: "owns", TargetEntityID: idG}, + {ID: uuid.NewString(), AgentID: testAgentID.String(), UserID: testUserID, SourceEntityID: idH, RelationType: "manages", TargetEntityID: idF}, + } + // F has no traversal results (mock empty) + + tool := NewKnowledgeGraphSearchTool() + tool.SetKGStore(ms) + + ctx := kgContext() + result := tool.executeTraversal(ctx, testAgentID.String(), testUserID, idF, 2, "") + text := result.ForLLM + + // Outgoing: F —[owns]→ G + if !strings.Contains(text, "NodeF —[owns]→ NodeG") { + t.Errorf("expected 'NodeF —[owns]→ NodeG' in output, got: %s", text) + } + // Incoming: H —[manages]→ F + if !strings.Contains(text, "NodeH —[manages]→ NodeF") { + t.Errorf("expected 'NodeH —[manages]→ NodeF' in output, got: %s", text) + } +} + +func TestKGTraversal_Tier1_CappedAt20(t *testing.T) { + ms := newMockKGStore() + startID := uuid.NewString() + ms.entities[startID] = store.Entity{ID: startID, AgentID: testAgentID.String(), UserID: testUserID, Name: "Start", EntityType: "concept"} + + // Create 25 traversal results + var results []store.TraversalResult + for i := 0; i < 25; i++ { + eid := uuid.NewString() + name := fmt.Sprintf("Node_%02d", i) + ms.entities[eid] = store.Entity{ID: eid, AgentID: testAgentID.String(), UserID: testUserID, Name: name, EntityType: "concept"} + results = append(results, store.TraversalResult{Entity: ms.entities[eid], Depth: 1, Via: "links_to"}) + } + ms.traversal[startID] = results + + tool := NewKnowledgeGraphSearchTool() + tool.SetKGStore(ms) + + ctx := kgContext() + result := tool.executeTraversal(ctx, testAgentID.String(), testUserID, startID, 2, "") + text := result.ForLLM + + count := strings.Count(text, "links_to") + if count > 20 { + t.Errorf("expected max 20 traversal results, got %d", count) + } + if !strings.Contains(text, "+5 more entities reachable") { + t.Errorf("expected truncation hint, got: %s", text) + } +} + +func TestKGSearch_RelationsCappedAt5(t *testing.T) { + ms := newMockKGStore() + entityID := uuid.NewString() + ms.entities[entityID] = store.Entity{ID: entityID, AgentID: testAgentID.String(), UserID: testUserID, Name: "HubEntity", EntityType: "concept"} + + // Create 8 outgoing relations from entity + for i := 0; i < 8; i++ { + tgtID := uuid.NewString() + tgtName := fmt.Sprintf("Target_%02d", i) + ms.entities[tgtID] = store.Entity{ID: tgtID, AgentID: testAgentID.String(), UserID: testUserID, Name: tgtName, EntityType: "concept"} + ms.relations = append(ms.relations, store.Relation{ + ID: uuid.NewString(), AgentID: testAgentID.String(), UserID: testUserID, + SourceEntityID: entityID, RelationType: "connects", TargetEntityID: tgtID, + }) + } + + tool := NewKnowledgeGraphSearchTool() + tool.SetKGStore(ms) + + ctx := kgContext() + result := tool.executeSearch(ctx, testAgentID.String(), testUserID, "HubEntity", nil) + text := result.ForLLM + + relCount := strings.Count(text, "—[connects]→") + if relCount > 5 { + t.Errorf("expected max 5 relations per entity in search, got %d", relCount) + } + if !strings.Contains(text, "+3 more") { + t.Errorf("expected truncation hint '+3 more', got: %s", text) + } + if !strings.Contains(text, "use entity_id=") { + t.Error("expected hint to use entity_id for full relations") + } +}