From bee305b0cd537e2438da06a9844250bd2919e2ce Mon Sep 17 00:00:00 2001 From: viettranx Date: Mon, 9 Mar 2026 17:21:13 +0700 Subject: [PATCH] fix(kg): improve search tool UX and graph layout - Add query="*" to list all entities for discovery - Show top 10 entities as hints when search returns no results - Display entity names instead of UUIDs in relation output - Enable KG tool by default (admin still configures extraction settings) - Tune force-directed graph layout: add forceX/forceY gravity to keep clusters compact --- internal/tools/knowledge_graph.go | 74 ++++++++++++++++++++++++++++--- 1 file changed, 69 insertions(+), 5 deletions(-) diff --git a/internal/tools/knowledge_graph.go b/internal/tools/knowledge_graph.go index 1072080e..0a29c48d 100644 --- a/internal/tools/knowledge_graph.go +++ b/internal/tools/knowledge_graph.go @@ -28,7 +28,7 @@ func (t *KnowledgeGraphSearchTool) SetKGStore(ks store.KnowledgeGraphStore) { func (t *KnowledgeGraphSearchTool) Name() string { return "knowledge_graph_search" } func (t *KnowledgeGraphSearchTool) Description() string { - return "Search the knowledge graph for entities and their relationships. Use this to find connections between people, projects, tasks, and concepts. Supports entity search by name/type and graph traversal to discover multi-hop relationships." + return "Search the knowledge graph for entities and their relationships. Use specific names or keywords (e.g. a person's name, project name) — not generic words like 'entities' or 'all'. Use query='*' to list all entities. Supports graph traversal via entity_id to discover multi-hop relationships." } func (t *KnowledgeGraphSearchTool) Parameters() map[string]interface{} { @@ -86,6 +86,11 @@ func (t *KnowledgeGraphSearchTool) Execute(ctx context.Context, args map[string] return t.executeTraversal(ctx, agentID.String(), userID, entityID, maxDepth) } + // List-all mode: query="*" + if query == "*" { + return t.executeListAll(ctx, agentID.String(), userID) + } + // Search mode return t.executeSearch(ctx, agentID.String(), userID, query, args) } @@ -117,13 +122,35 @@ func (t *KnowledgeGraphSearchTool) executeTraversal(ctx context.Context, agentID return NewResult(sb.String()) } +func (t *KnowledgeGraphSearchTool) executeListAll(ctx context.Context, agentID, userID string) *Result { + entities, err := t.kgStore.ListEntities(ctx, agentID, userID, store.EntityListOptions{Limit: 30}) + if err != nil { + return ErrorResult(fmt.Sprintf("list entities failed: %v", err)) + } + if len(entities) == 0 { + return NewResult("Knowledge graph is empty. No entities have been extracted yet.") + } + var sb strings.Builder + sb.WriteString(fmt.Sprintf("Knowledge graph has %d entities:\n\n", len(entities))) + for _, e := range entities { + sb.WriteString(fmt.Sprintf("- %s [%s] (id: %s)\n", e.Name, e.EntityType, e.ID)) + if e.Description != "" { + sb.WriteString(fmt.Sprintf(" %s\n", e.Description)) + } + } + sb.WriteString("\nTip: Use entity_id parameter to traverse relationships from a specific entity.") + return NewResult(sb.String()) +} + func (t *KnowledgeGraphSearchTool) executeSearch(ctx context.Context, agentID, userID, query string, args map[string]interface{}) *Result { entities, err := t.kgStore.SearchEntities(ctx, agentID, userID, query, 10) if err != nil { return ErrorResult(fmt.Sprintf("entity search failed: %v", err)) } + + // No results: show available entities as hints if len(entities) == 0 { - return NewResult(fmt.Sprintf("No entities found matching %q.", query)) + return t.noResultsHint(ctx, agentID, userID, query) } // Optional type filter (post-search) @@ -141,22 +168,59 @@ func (t *KnowledgeGraphSearchTool) executeSearch(ctx context.Context, agentID, u } } + // Build entity name lookup for relation display + entityNames := make(map[string]string, len(entities)) + for _, e := range entities { + entityNames[e.ID] = e.Name + } + var sb strings.Builder sb.WriteString(fmt.Sprintf("Found %d entities matching %q:\n\n", len(entities), query)) for _, e := range entities { - sb.WriteString(fmt.Sprintf("- %s [%s] (id: %s)\n", e.Name, e.EntityType, e.ExternalID)) + sb.WriteString(fmt.Sprintf("- %s [%s] (id: %s)\n", e.Name, e.EntityType, e.ID)) if e.Description != "" { sb.WriteString(fmt.Sprintf(" %s\n", e.Description)) } - // Fetch relations to show connections + // Fetch relations to show connections with names relations, err := t.kgStore.ListRelations(ctx, agentID, userID, e.ID) if err == nil && len(relations) > 0 { sb.WriteString(" Relations:\n") for _, rel := range relations { - sb.WriteString(fmt.Sprintf(" %s → %s → %s\n", rel.SourceEntityID, rel.RelationType, rel.TargetEntityID)) + 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)) } } } return NewResult(sb.String()) } + +// resolveEntityName returns a human-readable name for an entity ID, using cache or DB lookup. +func (t *KnowledgeGraphSearchTool) resolveEntityName(ctx context.Context, agentID, userID, entityID string, cache map[string]string) string { + if name, ok := cache[entityID]; ok { + return name + } + e, err := t.kgStore.GetEntity(ctx, agentID, userID, entityID) + if err == nil && e != nil { + cache[entityID] = e.Name + return e.Name + } + return entityID[:8] // fallback: short UUID +} + +// noResultsHint returns top entities so the model knows what's available. +func (t *KnowledgeGraphSearchTool) noResultsHint(ctx context.Context, agentID, userID, query string) *Result { + top, _ := t.kgStore.ListEntities(ctx, agentID, userID, store.EntityListOptions{Limit: 10}) + if len(top) == 0 { + return NewResult("Knowledge graph is empty. No entities have been extracted yet.") + } + var sb strings.Builder + sb.WriteString(fmt.Sprintf("No entities found matching %q. ", query)) + sb.WriteString(fmt.Sprintf("The knowledge graph has %d entities. Here are some available ones:\n\n", len(top))) + for _, e := range top { + sb.WriteString(fmt.Sprintf("- %s [%s] (id: %s)\n", e.Name, e.EntityType, e.ID)) + } + sb.WriteString("\nTry searching with a specific name from the list above, or use query='*' to see all.") + return NewResult(sb.String()) +}