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
This commit is contained in:
viettranx committed 2026-03-09 17:21:13 +07:00
1 parent 63eff188ad
commit bee305b0cd
1 file changed
+69 -5
+69 -5
View File
@@ -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())
}