feat(kg): add knowledge graph with LLM extraction, traversal, and graph visualization

- KnowledgeGraphStore interface + PostgreSQL implementation (recursive CTE traversal, 5s timeout)
- LLM entity extraction pipeline triggered on memory writes (background goroutine)
- knowledge_graph_search agent tool with search + traversal modes
- HTTP API: CRUD entities, traverse, extract, stats, graph endpoints
- Web UI: KG tab on memory page with table/graph toggle, entity detail, manual extraction
- Force-directed graph visualization using @xyflow/react + d3-force
- Builtin tool seed with configurable provider/model/confidence settings
This commit is contained in:
viettranx committed 2026-03-09 17:11:20 +07:00
1 parent 9249d9e358
commit 63eff188ad
31 files changed
+2625 -11

No files matched your search

+7 -1
View File
@@ -118,7 +118,8 @@ func runGateway() {
// Memory tools — PG-backed; always registered (PG memory is always available)
toolsReg.Register(tools.NewMemorySearchTool())
toolsReg.Register(tools.NewMemoryGetTool())
slog.Info("memory tools registered (PG-backed)")
toolsReg.Register(tools.NewKnowledgeGraphSearchTool())
slog.Info("memory + knowledge graph tools registered (PG-backed)")
// Browser automation tool
var browserMgr *browser.Manager
@@ -623,6 +624,11 @@ func runGateway() {
server.SetMemoryHandler(httpapi.NewMemoryHandler(pgStores.Memory, cfg.Gateway.Token))
}
// Knowledge graph API
if pgStores != nil && pgStores.KnowledgeGraph != nil {
server.SetKnowledgeGraphHandler(httpapi.NewKnowledgeGraphHandler(pgStores.KnowledgeGraph, providerRegistry, cfg.Gateway.Token))
}
// Workspace file serving endpoint — serves files by absolute path, auth-token protected.
// Supports media from any agent workspace (each agent has its own workspace from DB).
server.SetFilesHandler(httpapi.NewFilesHandler(cfg.Gateway.Token))
+4
View File
@@ -37,6 +37,10 @@ func builtinToolSeedData() []store.BuiltinToolDef {
{Name: "memory_get", DisplayName: "Memory Get", Description: "Retrieve a specific memory entry by key", Category: "memory", Enabled: true,
Requires: []string{"memory"},
},
{Name: "knowledge_graph_search", DisplayName: "Knowledge Graph Search", Description: "Search entities and traverse relationships in the knowledge graph", Category: "memory", Enabled: true,
Settings: json.RawMessage(`{"extract_on_memory_write":false,"extraction_provider":"","extraction_model":"","min_confidence":0.75}`),
Requires: []string{"knowledge_graph"},
},
// media
{Name: "read_image", DisplayName: "Read Image", Description: "Analyze images using a vision-capable LLM provider", Category: "media", Enabled: true,
+84 -4
View File
@@ -2,6 +2,7 @@ package cmd
import (
"context"
"encoding/json"
"log/slog"
"path/filepath"
@@ -11,6 +12,7 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/bus"
"github.com/nextlevelbuilder/goclaw/internal/config"
"github.com/nextlevelbuilder/goclaw/internal/hooks"
kg "github.com/nextlevelbuilder/goclaw/internal/knowledgegraph"
mcpbridge "github.com/nextlevelbuilder/goclaw/internal/mcp"
"github.com/nextlevelbuilder/goclaw/internal/media"
"github.com/nextlevelbuilder/goclaw/internal/providers"
@@ -162,6 +164,15 @@ func wireExtras(
// Wire virtual FS interceptors: route context + memory file reads/writes to DB.
// Share ONE ContextFileInterceptor instance between read_file and write_file
// so they share the same cache.
// Write-capable tools share a memory interceptor with optional KG extraction hook.
var writeMemIntc *tools.MemoryInterceptor
if stores.Memory != nil {
writeMemIntc = tools.NewMemoryInterceptor(stores.Memory, workspace)
// Hook KG extraction on memory writes if KG store is available
if stores.KnowledgeGraph != nil && stores.BuiltinTools != nil {
writeMemIntc.SetKGExtractFunc(buildKGExtractFunc(stores.KnowledgeGraph, stores.BuiltinTools, providerReg))
}
}
if readTool, ok := toolsReg.Get("read_file"); ok {
if ia, ok := readTool.(tools.InterceptorAware); ok {
if contextFileInterceptor != nil {
@@ -177,8 +188,8 @@ func wireExtras(
if contextFileInterceptor != nil {
ia.SetContextFileInterceptor(contextFileInterceptor)
}
if stores.Memory != nil {
ia.SetMemoryInterceptor(tools.NewMemoryInterceptor(stores.Memory, workspace))
if writeMemIntc != nil {
ia.SetMemoryInterceptor(writeMemIntc)
}
}
}
@@ -187,8 +198,8 @@ func wireExtras(
if contextFileInterceptor != nil {
ia.SetContextFileInterceptor(contextFileInterceptor)
}
if stores.Memory != nil {
ia.SetMemoryInterceptor(tools.NewMemoryInterceptor(stores.Memory, workspace))
if writeMemIntc != nil {
ia.SetMemoryInterceptor(writeMemIntc)
}
}
}
@@ -229,6 +240,16 @@ func wireExtras(
slog.Info("memory layering enabled (Postgres)")
}
// Wire knowledge graph store on KG tool
if stores.KnowledgeGraph != nil {
if kgTool, ok := toolsReg.Get("knowledge_graph_search"); ok {
if kgt, ok := kgTool.(*tools.KnowledgeGraphSearchTool); ok {
kgt.SetKGStore(stores.KnowledgeGraph)
}
}
slog.Info("knowledge graph tool wired (Postgres)")
}
// --- Cache invalidation event subscribers ---
// Context file cache: invalidate on agent/context data changes
@@ -514,3 +535,62 @@ func wireExtras(
return contextFileInterceptor, delegateMgr, mcpPool, mediaStore
}
// kgSettings holds KG extraction settings from the builtin_tools table.
type kgSettings struct {
ExtractOnMemoryWrite bool `json:"extract_on_memory_write"`
ExtractionProvider string `json:"extraction_provider"`
ExtractionModel string `json:"extraction_model"`
MinConfidence float64 `json:"min_confidence"`
}
// buildKGExtractFunc returns a callback that extracts entities from memory content.
// Settings are read from the builtin_tools table on each invocation (not cached),
// so changes take effect immediately without restart.
func buildKGExtractFunc(kgStore store.KnowledgeGraphStore, bts store.BuiltinToolStore, providerReg *providers.Registry) tools.KGExtractFunc {
return func(ctx context.Context, agentID, userID, content string) {
slog.Info("kg extract: triggered", "agent", agentID, "user", userID, "content_len", len(content))
// Read settings from DB on each call so admin changes take effect immediately
raw, err := bts.GetSettings(ctx, "knowledge_graph_search")
if err != nil || raw == nil {
slog.Warn("kg extract: no settings found", "error", err)
return
}
var settings kgSettings
if err := json.Unmarshal(raw, &settings); err != nil {
slog.Warn("kg extract: invalid settings", "error", err)
return
}
if !settings.ExtractOnMemoryWrite || settings.ExtractionProvider == "" || settings.ExtractionModel == "" {
return
}
p, err := providerReg.Get(settings.ExtractionProvider)
if err != nil {
slog.Warn("kg extract: provider not found", "provider", settings.ExtractionProvider, "error", err)
return
}
extractor := kg.NewExtractor(p, settings.ExtractionModel, settings.MinConfidence)
result, err := extractor.Extract(ctx, content)
if err != nil {
slog.Warn("kg extract: extraction failed", "agent", agentID, "error", err)
return
}
if len(result.Entities) == 0 && len(result.Relations) == 0 {
return
}
for i := range result.Entities {
result.Entities[i].AgentID = agentID
result.Entities[i].UserID = userID
}
for i := range result.Relations {
result.Relations[i].AgentID = agentID
result.Relations[i].UserID = userID
}
if err := kgStore.IngestExtraction(ctx, agentID, userID, result.Entities, result.Relations); err != nil {
slog.Warn("kg extract: ingest failed", "agent", agentID, "error", err)
return
}
slog.Info("kg extract: ingested from memory write", "agent", agentID, "entities", len(result.Entities), "relations", len(result.Relations))
}
}
+9
View File
@@ -48,6 +48,7 @@ type Server struct {
builtinToolsHandler *httpapi.BuiltinToolsHandler // builtin tool management API
pendingMessagesHandler *httpapi.PendingMessagesHandler // pending messages API
memoryHandler *httpapi.MemoryHandler // memory management API
kgHandler *httpapi.KnowledgeGraphHandler // knowledge graph API
oauthHandler *httpapi.OAuthHandler // OAuth endpoints
filesHandler *httpapi.FilesHandler // workspace file serving
storageHandler *httpapi.StorageHandler // storage file management
@@ -214,6 +215,11 @@ func (s *Server) BuildMux() *http.ServeMux {
s.memoryHandler.RegisterRoutes(mux)
}
// Knowledge graph API
if s.kgHandler != nil {
s.kgHandler.RegisterRoutes(mux)
}
// Workspace file serving (available in all modes)
if s.filesHandler != nil {
s.filesHandler.RegisterRoutes(mux)
@@ -449,6 +455,9 @@ func (s *Server) SetMediaServeHandler(h *httpapi.MediaServeHandler) { s.mediaSer
// SetMemoryHandler sets the memory management handler.
func (s *Server) SetMemoryHandler(h *httpapi.MemoryHandler) { s.memoryHandler = h }
// SetKnowledgeGraphHandler sets the knowledge graph handler.
func (s *Server) SetKnowledgeGraphHandler(h *httpapi.KnowledgeGraphHandler) { s.kgHandler = h }
// SetAgentStore sets the agent store for context injection in tools_invoke.
func (s *Server) SetAgentStore(as store.AgentStore) { s.agentStore = as }
+57
View File
@@ -0,0 +1,57 @@
package http
import (
"net/http"
kg "github.com/nextlevelbuilder/goclaw/internal/knowledgegraph"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
)
// KnowledgeGraphHandler handles KG entity/relation management endpoints.
type KnowledgeGraphHandler struct {
store store.KnowledgeGraphStore
providerReg *providers.Registry
token string
}
// NewKnowledgeGraphHandler creates a handler for KG management endpoints.
func NewKnowledgeGraphHandler(s store.KnowledgeGraphStore, providerReg *providers.Registry, token string) *KnowledgeGraphHandler {
return &KnowledgeGraphHandler{store: s, providerReg: providerReg, token: token}
}
// NewExtractor creates an Extractor from the given provider name and model.
func (h *KnowledgeGraphHandler) NewExtractor(providerName, model string, minConfidence float64) *kg.Extractor {
if h.providerReg == nil || providerName == "" || model == "" {
return nil
}
p, err := h.providerReg.Get(providerName)
if err != nil {
return nil
}
return kg.NewExtractor(p, model, minConfidence)
}
// RegisterRoutes registers all KG routes on the given mux.
func (h *KnowledgeGraphHandler) RegisterRoutes(mux *http.ServeMux) {
mux.HandleFunc("GET /v1/agents/{agentID}/kg/entities", h.auth(h.handleListEntities))
mux.HandleFunc("GET /v1/agents/{agentID}/kg/entities/{entityID}", h.auth(h.handleGetEntity))
mux.HandleFunc("POST /v1/agents/{agentID}/kg/entities", h.auth(h.handleUpsertEntity))
mux.HandleFunc("DELETE /v1/agents/{agentID}/kg/entities/{entityID}", h.auth(h.handleDeleteEntity))
mux.HandleFunc("POST /v1/agents/{agentID}/kg/traverse", h.auth(h.handleTraverse))
mux.HandleFunc("POST /v1/agents/{agentID}/kg/extract", h.auth(h.handleExtract))
mux.HandleFunc("GET /v1/agents/{agentID}/kg/stats", h.auth(h.handleStats))
mux.HandleFunc("GET /v1/agents/{agentID}/kg/graph", h.auth(h.handleGraph))
}
func (h *KnowledgeGraphHandler) auth(next http.HandlerFunc) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
if h.token != "" {
if extractBearerToken(r) != h.token {
writeJSON(w, http.StatusUnauthorized, map[string]string{"error": "unauthorized"})
return
}
}
next(w, r)
}
}
+259
View File
@@ -0,0 +1,259 @@
package http
import (
"encoding/json"
"log/slog"
"net/http"
"strconv"
"github.com/nextlevelbuilder/goclaw/internal/store"
)
func (h *KnowledgeGraphHandler) handleListEntities(w http.ResponseWriter, r *http.Request) {
agentID := r.PathValue("agentID")
userID := r.URL.Query().Get("user_id")
entityType := r.URL.Query().Get("type")
query := r.URL.Query().Get("q")
limit, _ := strconv.Atoi(r.URL.Query().Get("limit"))
offset, _ := strconv.Atoi(r.URL.Query().Get("offset"))
if limit <= 0 {
limit = 50
}
// If query is provided, use search
if query != "" {
entities, err := h.store.SearchEntities(r.Context(), agentID, userID, query, limit)
if err != nil {
slog.Warn("kg.search_entities failed", "error", err)
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()})
return
}
if entities == nil {
entities = []store.Entity{}
}
writeJSON(w, http.StatusOK, entities)
return
}
entities, err := h.store.ListEntities(r.Context(), agentID, userID, store.EntityListOptions{
EntityType: entityType,
Limit: limit,
Offset: offset,
})
if err != nil {
slog.Warn("kg.list_entities failed", "error", err)
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()})
return
}
if entities == nil {
entities = []store.Entity{}
}
writeJSON(w, http.StatusOK, entities)
}
func (h *KnowledgeGraphHandler) handleGetEntity(w http.ResponseWriter, r *http.Request) {
agentID := r.PathValue("agentID")
entityID := r.PathValue("entityID")
userID := r.URL.Query().Get("user_id")
entity, err := h.store.GetEntity(r.Context(), agentID, userID, entityID)
if err != nil {
writeJSON(w, http.StatusNotFound, map[string]string{"error": "entity not found"})
return
}
relations, err := h.store.ListRelations(r.Context(), agentID, userID, entityID)
if err != nil {
relations = []store.Relation{}
}
if relations == nil {
relations = []store.Relation{}
}
writeJSON(w, http.StatusOK, map[string]interface{}{
"entity": entity,
"relations": relations,
})
}
func (h *KnowledgeGraphHandler) handleUpsertEntity(w http.ResponseWriter, r *http.Request) {
agentID := r.PathValue("agentID")
var entity store.Entity
if err := json.NewDecoder(r.Body).Decode(&entity); err != nil {
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid request body"})
return
}
entity.AgentID = agentID
if entity.ExternalID == "" || entity.Name == "" || entity.EntityType == "" {
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "external_id, name, and entity_type are required"})
return
}
if entity.Confidence <= 0 {
entity.Confidence = 1.0
}
if err := h.store.UpsertEntity(r.Context(), &entity); err != nil {
slog.Warn("kg.upsert_entity failed", "error", err)
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()})
return
}
writeJSON(w, http.StatusOK, map[string]string{"status": "ok"})
}
func (h *KnowledgeGraphHandler) handleDeleteEntity(w http.ResponseWriter, r *http.Request) {
agentID := r.PathValue("agentID")
entityID := r.PathValue("entityID")
userID := r.URL.Query().Get("user_id")
if err := h.store.DeleteEntity(r.Context(), agentID, userID, entityID); err != nil {
slog.Warn("kg.delete_entity failed", "error", err)
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()})
return
}
writeJSON(w, http.StatusOK, map[string]string{"status": "deleted"})
}
func (h *KnowledgeGraphHandler) handleTraverse(w http.ResponseWriter, r *http.Request) {
agentID := r.PathValue("agentID")
var body struct {
EntityID string `json:"entity_id"`
UserID string `json:"user_id"`
MaxDepth int `json:"max_depth"`
}
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid request body"})
return
}
if body.EntityID == "" {
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "entity_id is required"})
return
}
if body.MaxDepth <= 0 {
body.MaxDepth = 2
}
if body.MaxDepth > 3 {
body.MaxDepth = 3
}
results, err := h.store.Traverse(r.Context(), agentID, body.UserID, body.EntityID, body.MaxDepth)
if err != nil {
slog.Warn("kg.traverse failed", "error", err)
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()})
return
}
if results == nil {
results = []store.TraversalResult{}
}
writeJSON(w, http.StatusOK, results)
}
func (h *KnowledgeGraphHandler) handleExtract(w http.ResponseWriter, r *http.Request) {
agentID := r.PathValue("agentID")
var body struct {
Text string `json:"text"`
UserID string `json:"user_id"`
Provider string `json:"provider"`
Model string `json:"model"`
MinConf float64 `json:"min_confidence"`
}
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid request body"})
return
}
if body.Text == "" {
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "text is required"})
return
}
if body.Provider == "" || body.Model == "" {
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "provider and model are required"})
return
}
extractor := h.NewExtractor(body.Provider, body.Model, body.MinConf)
if extractor == nil {
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid provider or model"})
return
}
result, err := extractor.Extract(r.Context(), body.Text)
if err != nil {
slog.Warn("kg.extract failed", "error", err)
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()})
return
}
// Ingest extracted entities and relations into the store
for i := range result.Entities {
result.Entities[i].AgentID = agentID
result.Entities[i].UserID = body.UserID
}
for i := range result.Relations {
result.Relations[i].AgentID = agentID
result.Relations[i].UserID = body.UserID
}
if err := h.store.IngestExtraction(r.Context(), agentID, body.UserID, result.Entities, result.Relations); err != nil {
slog.Warn("kg.ingest_extraction failed", "error", err)
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()})
return
}
writeJSON(w, http.StatusOK, map[string]interface{}{
"entities": len(result.Entities),
"relations": len(result.Relations),
})
}
func (h *KnowledgeGraphHandler) handleStats(w http.ResponseWriter, r *http.Request) {
agentID := r.PathValue("agentID")
userID := r.URL.Query().Get("user_id")
stats, err := h.store.Stats(r.Context(), agentID, userID)
if err != nil {
slog.Warn("kg.stats failed", "error", err)
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()})
return
}
writeJSON(w, http.StatusOK, stats)
}
// handleGraph returns all entities + relations for graph visualization.
func (h *KnowledgeGraphHandler) handleGraph(w http.ResponseWriter, r *http.Request) {
agentID := r.PathValue("agentID")
userID := r.URL.Query().Get("user_id")
limit, _ := strconv.Atoi(r.URL.Query().Get("limit"))
if limit <= 0 {
limit = 200
}
entities, err := h.store.ListEntities(r.Context(), agentID, userID, store.EntityListOptions{Limit: limit})
if err != nil {
slog.Warn("kg.graph entities failed", "error", err)
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()})
return
}
if entities == nil {
entities = []store.Entity{}
}
relations, err := h.store.ListAllRelations(r.Context(), agentID, userID, limit*3)
if err != nil {
slog.Warn("kg.graph relations failed", "error", err)
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()})
return
}
if relations == nil {
relations = []store.Relation{}
}
writeJSON(w, http.StatusOK, map[string]interface{}{
"entities": entities,
"relations": relations,
})
}
+108
View File
@@ -0,0 +1,108 @@
package knowledgegraph
import (
"context"
"encoding/json"
"fmt"
"log/slog"
"strings"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
)
// ExtractionResult holds entities and relations extracted from text.
type ExtractionResult struct {
Entities []store.Entity `json:"entities"`
Relations []store.Relation `json:"relations"`
}
// Extractor extracts entities and relations from text using an LLM.
type Extractor struct {
provider providers.Provider
model string
minConfidence float64
}
// NewExtractor creates a new Extractor with the given provider, model, and confidence threshold.
func NewExtractor(provider providers.Provider, model string, minConfidence float64) *Extractor {
if minConfidence <= 0 {
minConfidence = 0.75
}
return &Extractor{provider: provider, model: model, minConfidence: minConfidence}
}
// Extract calls the LLM to extract entities and relations from text.
func (e *Extractor) Extract(ctx context.Context, text string) (*ExtractionResult, error) {
// Truncate very long texts to avoid overwhelming the LLM
const maxInputChars = 6000
if len(text) > maxInputChars {
text = text[:maxInputChars] + "\n\n[...truncated]"
}
req := providers.ChatRequest{
Messages: []providers.Message{
{Role: "system", Content: extractionSystemPrompt},
{Role: "user", Content: text},
},
Model: e.model,
Options: map[string]interface{}{
"max_tokens": 8192,
"temperature": 0.0,
},
}
resp, err := e.provider.Chat(ctx, req)
if err != nil {
return nil, fmt.Errorf("kg extraction LLM call: %w", err)
}
// Parse JSON response
var result ExtractionResult
content := strings.TrimSpace(resp.Content)
// Handle markdown code blocks
content = stripCodeBlock(content)
if err := json.Unmarshal([]byte(content), &result); err != nil {
preview := content
if len(preview) > 300 {
preview = preview[:300] + "..."
}
slog.Warn("kg extraction: failed to parse LLM response", "error", err, "content_len", len(content), "finish_reason", resp.FinishReason, "preview", preview)
return nil, fmt.Errorf("parse extraction result: %w", err)
}
// Filter by confidence threshold
filtered := &ExtractionResult{}
for _, ent := range result.Entities {
if ent.Confidence >= e.minConfidence {
ent.ExternalID = strings.ToLower(strings.TrimSpace(ent.ExternalID))
ent.Name = strings.TrimSpace(ent.Name)
ent.EntityType = strings.ToLower(strings.TrimSpace(ent.EntityType))
filtered.Entities = append(filtered.Entities, ent)
}
}
for _, rel := range result.Relations {
if rel.Confidence >= e.minConfidence {
rel.SourceEntityID = strings.ToLower(strings.TrimSpace(rel.SourceEntityID))
rel.TargetEntityID = strings.ToLower(strings.TrimSpace(rel.TargetEntityID))
rel.RelationType = strings.ToLower(strings.TrimSpace(rel.RelationType))
filtered.Relations = append(filtered.Relations, rel)
}
}
return filtered, nil
}
// stripCodeBlock removes ```json ... ``` wrapper if present.
func stripCodeBlock(s string) string {
s = strings.TrimSpace(s)
if strings.HasPrefix(s, "```") {
if idx := strings.Index(s, "\n"); idx >= 0 {
s = s[idx+1:]
}
if idx := strings.LastIndex(s, "```"); idx >= 0 {
s = s[:idx]
}
}
return strings.TrimSpace(s)
}
@@ -0,0 +1,34 @@
package knowledgegraph
const extractionSystemPrompt = `You are an entity and relationship extractor. Given text, extract structured knowledge.
Output valid JSON with this schema:
{
"entities": [
{
"external_id": "unique-lowercase-id",
"name": "Display Name",
"entity_type": "person|project|task|event|concept|location|organization",
"description": "Brief description (max 50 chars)",
"confidence": 0.0-1.0
}
],
"relations": [
{
"source_entity_id": "external_id of source entity",
"relation_type": "verb_phrase",
"target_entity_id": "external_id of target entity",
"confidence": 0.0-1.0
}
]
}
Rules:
- Extract only the TOP 15 most important entities and their relations
- Keep descriptions very short (under 50 characters)
- external_id: lowercase, no spaces, use hyphens (e.g. "john-doe", "project-alpha")
- entity_type: person, project, task, event, concept, location, or organization
- relation_type: lowercase verb phrase with underscores (e.g. works_on, reported_to)
- confidence: 1.0 = explicitly stated, 0.5 = inferred
- Keep names in original language
- Output ONLY the JSON object, no markdown, no explanation, no code blocks`
+78
View File
@@ -0,0 +1,78 @@
package store
import "context"
// Entity represents a node in the knowledge graph.
type Entity struct {
ID string `json:"id"`
AgentID string `json:"agent_id"`
UserID string `json:"user_id,omitempty"`
ExternalID string `json:"external_id"`
Name string `json:"name"`
EntityType string `json:"entity_type"`
Description string `json:"description,omitempty"`
Properties map[string]string `json:"properties,omitempty"`
SourceID string `json:"source_id,omitempty"`
Confidence float64 `json:"confidence"`
CreatedAt int64 `json:"created_at"`
UpdatedAt int64 `json:"updated_at"`
}
// Relation represents an edge between two entities.
type Relation struct {
ID string `json:"id"`
AgentID string `json:"agent_id"`
UserID string `json:"user_id,omitempty"`
SourceEntityID string `json:"source_entity_id"`
RelationType string `json:"relation_type"`
TargetEntityID string `json:"target_entity_id"`
Confidence float64 `json:"confidence"`
Properties map[string]string `json:"properties,omitempty"`
CreatedAt int64 `json:"created_at"`
}
// TraversalResult is a connected entity with path info.
type TraversalResult struct {
Entity Entity `json:"entity"`
Depth int `json:"depth"`
Path []string `json:"path"`
Via string `json:"via"`
}
// EntityListOptions configures a list query for entities.
type EntityListOptions struct {
EntityType string
Limit int
Offset int
}
// GraphStats contains aggregate counts for a scoped graph.
type GraphStats struct {
EntityCount int `json:"entity_count"`
RelationCount int `json:"relation_count"`
EntityTypes map[string]int `json:"entity_types"`
}
// KnowledgeGraphStore manages entity-relationship graphs.
type KnowledgeGraphStore interface {
UpsertEntity(ctx context.Context, entity *Entity) error
GetEntity(ctx context.Context, agentID, userID, entityID string) (*Entity, error)
DeleteEntity(ctx context.Context, agentID, userID, entityID string) error
ListEntities(ctx context.Context, agentID, userID string, opts EntityListOptions) ([]Entity, error)
SearchEntities(ctx context.Context, agentID, userID, query string, limit int) ([]Entity, error)
UpsertRelation(ctx context.Context, relation *Relation) error
DeleteRelation(ctx context.Context, agentID, userID, relationID string) error
ListRelations(ctx context.Context, agentID, userID, entityID string) ([]Relation, error)
// ListAllRelations returns all relations for an agent (optionally scoped by user).
ListAllRelations(ctx context.Context, agentID, userID string, limit int) ([]Relation, error)
Traverse(ctx context.Context, agentID, userID, startEntityID string, maxDepth int) ([]TraversalResult, error)
IngestExtraction(ctx context.Context, agentID, userID string, entities []Entity, relations []Relation) error
PruneByConfidence(ctx context.Context, agentID, userID string, minConfidence float64) (int, error)
Stats(ctx context.Context, agentID, userID string) (*GraphStats, error)
Close() error
}
+5
View File
@@ -116,6 +116,11 @@ func (s *PGBuiltinToolStore) Seed(ctx context.Context, tools []store.BuiltinTool
category = EXCLUDED.category,
requires = EXCLUDED.requires,
metadata = EXCLUDED.metadata,
settings = CASE
WHEN builtin_tools.settings IS NULL OR builtin_tools.settings::text IN ('{}', 'null')
THEN EXCLUDED.settings
ELSE builtin_tools.settings
END,
updated_at = EXCLUDED.updated_at`)
if err != nil {
return fmt.Errorf("prepare seed stmt: %w", err)
+1
View File
@@ -40,5 +40,6 @@ func NewPGStores(cfg store.StoreConfig) (*store.Stores, error) {
Teams: NewPGTeamStore(db),
BuiltinTools: NewPGBuiltinToolStore(db),
PendingMessages: NewPGPendingMessageStore(db),
KnowledgeGraph: NewPGKnowledgeGraphStore(db),
}, nil
}
+400
View File
@@ -0,0 +1,400 @@
package pg
import (
"context"
"database/sql"
"encoding/json"
"fmt"
"time"
"github.com/google/uuid"
"github.com/nextlevelbuilder/goclaw/internal/store"
)
// PGKnowledgeGraphStore implements store.KnowledgeGraphStore backed by Postgres.
type PGKnowledgeGraphStore struct {
db *sql.DB
}
// NewPGKnowledgeGraphStore creates a new PG-backed knowledge graph store.
func NewPGKnowledgeGraphStore(db *sql.DB) *PGKnowledgeGraphStore {
return &PGKnowledgeGraphStore{db: db}
}
func (s *PGKnowledgeGraphStore) UpsertEntity(ctx context.Context, entity *store.Entity) error {
aid := mustParseUUID(entity.AgentID)
props, err := json.Marshal(entity.Properties)
if err != nil {
props = []byte("{}")
}
now := time.Now()
id := uuid.Must(uuid.NewV7())
_, err = s.db.ExecContext(ctx, `
INSERT INTO kg_entities
(id, agent_id, user_id, external_id, name, entity_type, description, properties, source_id, confidence, created_at, updated_at)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $11)
ON CONFLICT (agent_id, user_id, external_id) DO UPDATE SET
name = EXCLUDED.name,
entity_type = EXCLUDED.entity_type,
description = EXCLUDED.description,
properties = EXCLUDED.properties,
source_id = EXCLUDED.source_id,
confidence = EXCLUDED.confidence,
updated_at = EXCLUDED.updated_at`,
id, aid, entity.UserID, entity.ExternalID, entity.Name, entity.EntityType,
entity.Description, props, entity.SourceID, entity.Confidence, now,
)
return err
}
func (s *PGKnowledgeGraphStore) GetEntity(ctx context.Context, agentID, userID, entityID string) (*store.Entity, error) {
aid := mustParseUUID(agentID)
eid := mustParseUUID(entityID)
row := s.db.QueryRowContext(ctx, `
SELECT id, agent_id, user_id, external_id, name, entity_type, description,
properties, source_id, confidence, created_at, updated_at
FROM kg_entities WHERE id = $1 AND agent_id = $2 AND user_id = $3`,
eid, aid, userID,
)
return scanEntity(row)
}
func (s *PGKnowledgeGraphStore) DeleteEntity(ctx context.Context, agentID, userID, entityID string) error {
aid := mustParseUUID(agentID)
eid := mustParseUUID(entityID)
_, err := s.db.ExecContext(ctx,
`DELETE FROM kg_entities WHERE id = $1 AND agent_id = $2 AND user_id = $3`,
eid, aid, userID,
)
return err
}
func (s *PGKnowledgeGraphStore) ListEntities(ctx context.Context, agentID, userID string, opts store.EntityListOptions) ([]store.Entity, error) {
aid := mustParseUUID(agentID)
limit := opts.Limit
if limit <= 0 {
limit = 50
}
// Build dynamic WHERE clause: always filter by agent_id, optionally by user_id and entity_type
where := "agent_id = $1"
args := []any{aid}
idx := 2
if userID != "" {
where += fmt.Sprintf(" AND user_id = $%d", idx)
args = append(args, userID)
idx++
}
if opts.EntityType != "" {
where += fmt.Sprintf(" AND entity_type = $%d", idx)
args = append(args, opts.EntityType)
idx++
}
args = append(args, limit, opts.Offset)
query := fmt.Sprintf(`
SELECT id, agent_id, user_id, external_id, name, entity_type, description,
properties, source_id, confidence, created_at, updated_at
FROM kg_entities WHERE %s
ORDER BY updated_at DESC LIMIT $%d OFFSET $%d`, where, idx, idx+1)
rows, err := s.db.QueryContext(ctx, query, args...)
if err != nil {
return nil, err
}
defer rows.Close()
return scanEntities(rows)
}
func (s *PGKnowledgeGraphStore) SearchEntities(ctx context.Context, agentID, userID, query string, limit int) ([]store.Entity, error) {
aid := mustParseUUID(agentID)
if limit <= 0 {
limit = 20
}
pattern := "%" + query + "%"
where := "agent_id = $1"
args := []any{aid}
idx := 2
if userID != "" {
where += fmt.Sprintf(" AND user_id = $%d", idx)
args = append(args, userID)
idx++
}
args = append(args, pattern, limit)
q := fmt.Sprintf(`
SELECT id, agent_id, user_id, external_id, name, entity_type, description,
properties, source_id, confidence, created_at, updated_at
FROM kg_entities
WHERE %s AND (name ILIKE $%d OR description ILIKE $%d)
ORDER BY confidence DESC, updated_at DESC LIMIT $%d`, where, idx, idx, idx+1)
rows, err := s.db.QueryContext(ctx, q, args...)
if err != nil {
return nil, err
}
defer rows.Close()
return scanEntities(rows)
}
func (s *PGKnowledgeGraphStore) UpsertRelation(ctx context.Context, relation *store.Relation) error {
aid := mustParseUUID(relation.AgentID)
src := mustParseUUID(relation.SourceEntityID)
tgt := mustParseUUID(relation.TargetEntityID)
props, err := json.Marshal(relation.Properties)
if err != nil {
props = []byte("{}")
}
id := uuid.Must(uuid.NewV7())
now := time.Now()
_, err = s.db.ExecContext(ctx, `
INSERT INTO kg_relations
(id, agent_id, user_id, source_entity_id, relation_type, target_entity_id, confidence, properties, created_at)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
ON CONFLICT (agent_id, user_id, source_entity_id, relation_type, target_entity_id) DO UPDATE SET
confidence = EXCLUDED.confidence,
properties = EXCLUDED.properties`,
id, aid, relation.UserID, src, relation.RelationType, tgt, relation.Confidence, props, now,
)
return err
}
func (s *PGKnowledgeGraphStore) DeleteRelation(ctx context.Context, agentID, userID, relationID string) error {
aid := mustParseUUID(agentID)
rid := mustParseUUID(relationID)
_, err := s.db.ExecContext(ctx,
`DELETE FROM kg_relations WHERE id = $1 AND agent_id = $2 AND user_id = $3`,
rid, aid, userID,
)
return err
}
func (s *PGKnowledgeGraphStore) ListRelations(ctx context.Context, agentID, userID, entityID string) ([]store.Relation, error) {
aid := mustParseUUID(agentID)
eid := mustParseUUID(entityID)
rows, err := s.db.QueryContext(ctx, `
SELECT id, agent_id, user_id, source_entity_id, relation_type, target_entity_id,
confidence, properties, created_at
FROM kg_relations
WHERE agent_id = $1 AND user_id = $2
AND (source_entity_id = $3 OR target_entity_id = $3)
ORDER BY created_at DESC`,
aid, userID, eid,
)
if err != nil {
return nil, err
}
defer rows.Close()
return scanRelations(rows)
}
func (s *PGKnowledgeGraphStore) ListAllRelations(ctx context.Context, agentID, userID string, limit int) ([]store.Relation, error) {
aid := mustParseUUID(agentID)
if limit <= 0 {
limit = 200
}
where := "agent_id = $1"
args := []any{aid}
idx := 2
if userID != "" {
where += fmt.Sprintf(" AND user_id = $%d", idx)
args = append(args, userID)
idx++
}
args = append(args, limit)
q := fmt.Sprintf(`
SELECT id, agent_id, user_id, source_entity_id, relation_type, target_entity_id,
confidence, properties, created_at
FROM kg_relations WHERE %s
ORDER BY created_at DESC LIMIT $%d`, where, idx)
rows, err := s.db.QueryContext(ctx, q, args...)
if err != nil {
return nil, err
}
defer rows.Close()
return scanRelations(rows)
}
func (s *PGKnowledgeGraphStore) IngestExtraction(ctx context.Context, agentID, userID string, entities []store.Entity, relations []store.Relation) error {
tx, err := s.db.BeginTx(ctx, nil)
if err != nil {
return err
}
defer tx.Rollback() //nolint:errcheck
aid := mustParseUUID(agentID)
now := time.Now()
// Upsert entities and build external_id → DB UUID lookup for relations
extIDToUUID := make(map[string]uuid.UUID, len(entities))
for i := range entities {
e := &entities[i]
e.AgentID = agentID
e.UserID = userID
props, _ := json.Marshal(e.Properties)
id := uuid.Must(uuid.NewV7())
// Use RETURNING to get the actual ID (could be existing row on conflict)
var actualID uuid.UUID
if err := tx.QueryRowContext(ctx, `
INSERT INTO kg_entities
(id, agent_id, user_id, external_id, name, entity_type, description, properties, source_id, confidence, created_at, updated_at)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $11)
ON CONFLICT (agent_id, user_id, external_id) DO UPDATE SET
name = EXCLUDED.name,
entity_type = EXCLUDED.entity_type,
description = EXCLUDED.description,
properties = EXCLUDED.properties,
source_id = EXCLUDED.source_id,
confidence = EXCLUDED.confidence,
updated_at = EXCLUDED.updated_at
RETURNING id`,
id, aid, userID, e.ExternalID, e.Name, e.EntityType,
e.Description, props, e.SourceID, e.Confidence, now,
).Scan(&actualID); err != nil {
return err
}
extIDToUUID[e.ExternalID] = actualID
}
for i := range relations {
r := &relations[i]
r.AgentID = agentID
r.UserID = userID
// Resolve external_id references to actual DB UUIDs
src, ok1 := extIDToUUID[r.SourceEntityID]
tgt, ok2 := extIDToUUID[r.TargetEntityID]
if !ok1 || !ok2 {
continue // skip relations referencing unknown entities
}
props, _ := json.Marshal(r.Properties)
id := uuid.Must(uuid.NewV7())
if _, err := tx.ExecContext(ctx, `
INSERT INTO kg_relations
(id, agent_id, user_id, source_entity_id, relation_type, target_entity_id, confidence, properties, created_at)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
ON CONFLICT (agent_id, user_id, source_entity_id, relation_type, target_entity_id) DO UPDATE SET
confidence = EXCLUDED.confidence,
properties = EXCLUDED.properties`,
id, aid, userID, src, r.RelationType, tgt, r.Confidence, props, now,
); err != nil {
return err
}
}
return tx.Commit()
}
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,
)
if err != nil {
return 0, err
}
n, _ := res.RowsAffected()
return int(n), nil
}
func (s *PGKnowledgeGraphStore) Stats(ctx context.Context, agentID, userID string) (*store.GraphStats, error) {
aid := mustParseUUID(agentID)
stats := &store.GraphStats{EntityTypes: make(map[string]int)}
userFilter := ""
args := []any{aid}
if userID != "" {
userFilter = " AND user_id = $2"
args = append(args, userID)
}
if err := s.db.QueryRowContext(ctx,
`SELECT COUNT(*) FROM kg_entities WHERE agent_id = $1`+userFilter, args...,
).Scan(&stats.EntityCount); err != nil {
return nil, err
}
if err := s.db.QueryRowContext(ctx,
`SELECT COUNT(*) FROM kg_relations WHERE agent_id = $1`+userFilter, args...,
).Scan(&stats.RelationCount); err != nil {
return nil, err
}
rows, err := s.db.QueryContext(ctx,
`SELECT entity_type, COUNT(*) FROM kg_entities WHERE agent_id = $1`+userFilter+` GROUP BY entity_type`, args...,
)
if err != nil {
return nil, err
}
defer rows.Close()
for rows.Next() {
var t string
var c int
if err := rows.Scan(&t, &c); err != nil {
continue
}
stats.EntityTypes[t] = c
}
return stats, nil
}
func (s *PGKnowledgeGraphStore) Close() error { return nil }
// --- scan helpers ---
type rowScanner interface {
Scan(dest ...any) error
}
func scanEntity(row rowScanner) (*store.Entity, error) {
var e store.Entity
var props []byte
var createdAt, updatedAt time.Time
if err := row.Scan(
&e.ID, &e.AgentID, &e.UserID, &e.ExternalID, &e.Name, &e.EntityType,
&e.Description, &props, &e.SourceID, &e.Confidence, &createdAt, &updatedAt,
); err != nil {
return nil, err
}
json.Unmarshal(props, &e.Properties) //nolint:errcheck
e.CreatedAt = createdAt.UnixMilli()
e.UpdatedAt = updatedAt.UnixMilli()
return &e, nil
}
func scanEntities(rows *sql.Rows) ([]store.Entity, error) {
var result []store.Entity
for rows.Next() {
var e store.Entity
var props []byte
var createdAt, updatedAt time.Time
if err := rows.Scan(
&e.ID, &e.AgentID, &e.UserID, &e.ExternalID, &e.Name, &e.EntityType,
&e.Description, &props, &e.SourceID, &e.Confidence, &createdAt, &updatedAt,
); err != nil {
continue
}
json.Unmarshal(props, &e.Properties) //nolint:errcheck
e.CreatedAt = createdAt.UnixMilli()
e.UpdatedAt = updatedAt.UnixMilli()
result = append(result, e)
}
return result, rows.Err()
}
func scanRelations(rows *sql.Rows) ([]store.Relation, error) {
var result []store.Relation
for rows.Next() {
var r store.Relation
var props []byte
var createdAt time.Time
if err := rows.Scan(
&r.ID, &r.AgentID, &r.UserID, &r.SourceEntityID, &r.RelationType,
&r.TargetEntityID, &r.Confidence, &props, &createdAt,
); err != nil {
continue
}
json.Unmarshal(props, &r.Properties) //nolint:errcheck
r.CreatedAt = createdAt.UnixMilli()
result = append(result, r)
}
return result, rows.Err()
}
@@ -0,0 +1,112 @@
package pg
import (
"context"
"encoding/json"
"time"
"github.com/lib/pq"
"github.com/nextlevelbuilder/goclaw/internal/store"
)
// Traverse walks the knowledge graph from startEntityID up to maxDepth hops
// using a recursive CTE. Returns all reachable entities (excluding the start node).
// A 5-second statement timeout is applied for safety.
func (s *PGKnowledgeGraphStore) Traverse(ctx context.Context, agentID, userID, startEntityID string, maxDepth int) ([]store.TraversalResult, error) {
if maxDepth <= 0 {
maxDepth = 3
}
aid := mustParseUUID(agentID)
startID := mustParseUUID(startEntityID)
tx, err := s.db.BeginTx(ctx, nil)
if err != nil {
return nil, err
}
defer tx.Rollback() //nolint:errcheck
if _, err := tx.ExecContext(ctx, `SET LOCAL statement_timeout = '5000'`); err != nil {
return nil, err
}
rows, err := tx.QueryContext(ctx, `
WITH RECURSIVE paths AS (
SELECT
e.id, e.agent_id, e.user_id, e.external_id,
e.name, e.entity_type, e.description,
e.properties, e.source_id, e.confidence,
e.created_at, e.updated_at,
1 AS depth,
ARRAY[e.id::text] AS path,
''::text AS via
FROM kg_entities e
WHERE e.id = $1 AND e.agent_id = $2 AND e.user_id = $3
UNION ALL
SELECT
e.id, e.agent_id, e.user_id, e.external_id,
e.name, e.entity_type, e.description,
e.properties, e.source_id, e.confidence,
e.created_at, e.updated_at,
p.depth + 1,
p.path || e.id::text,
r.relation_type
FROM paths p
JOIN kg_relations r ON p.id = r.source_entity_id
JOIN kg_entities e ON r.target_entity_id = e.id
WHERE p.depth < $4
AND NOT e.id::text = ANY(p.path)
)
SELECT
id, agent_id, user_id, external_id,
name, entity_type, description,
properties, source_id, confidence,
created_at, updated_at,
depth, path, via
FROM paths WHERE depth > 1`,
startID, aid, userID, maxDepth,
)
if err != nil {
return nil, err
}
defer rows.Close()
var results []store.TraversalResult
for rows.Next() {
var e store.Entity
var props []byte
var createdAt, updatedAt time.Time
var depth int
var path []string
var via string
if err := rows.Scan(
&e.ID, &e.AgentID, &e.UserID, &e.ExternalID,
&e.Name, &e.EntityType, &e.Description,
&props, &e.SourceID, &e.Confidence,
&createdAt, &updatedAt,
&depth, pq.Array(&path), &via,
); err != nil {
continue
}
if len(props) > 0 {
json.Unmarshal(props, &e.Properties) //nolint:errcheck
}
e.CreatedAt = createdAt.UnixMilli()
e.UpdatedAt = updatedAt.UnixMilli()
results = append(results, store.TraversalResult{
Entity: e,
Depth: depth,
Path: path,
Via: via,
})
}
if err := rows.Err(); err != nil {
return nil, err
}
return results, tx.Commit()
}
+1
View File
@@ -21,4 +21,5 @@ type Stores struct {
Teams TeamStore
BuiltinTools BuiltinToolStore
PendingMessages PendingMessageStore
KnowledgeGraph KnowledgeGraphStore
}
+162
View File
@@ -0,0 +1,162 @@
package tools
import (
"context"
"fmt"
"strings"
"github.com/google/uuid"
"github.com/nextlevelbuilder/goclaw/internal/store"
)
// KnowledgeGraphSearchTool provides graph-based search for agents.
type KnowledgeGraphSearchTool struct {
kgStore store.KnowledgeGraphStore
}
// NewKnowledgeGraphSearchTool creates a new KnowledgeGraphSearchTool.
func NewKnowledgeGraphSearchTool() *KnowledgeGraphSearchTool {
return &KnowledgeGraphSearchTool{}
}
// SetKGStore sets the KnowledgeGraphStore for this tool.
func (t *KnowledgeGraphSearchTool) SetKGStore(ks store.KnowledgeGraphStore) {
t.kgStore = ks
}
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."
}
func (t *KnowledgeGraphSearchTool) Parameters() map[string]interface{} {
return map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"query": map[string]interface{}{
"type": "string",
"description": "Search query for entity names or descriptions",
},
"entity_type": map[string]interface{}{
"type": "string",
"description": "Filter by entity type (person, project, task, event, concept, location, organization)",
},
"entity_id": map[string]interface{}{
"type": "string",
"description": "Entity ID to traverse from (for relationship discovery)",
},
"max_depth": map[string]interface{}{
"type": "number",
"description": "Maximum traversal depth (default 2, max 3)",
},
},
"required": []string{"query"},
}
}
func (t *KnowledgeGraphSearchTool) Execute(ctx context.Context, args map[string]interface{}) *Result {
if t.kgStore == nil {
return NewResult("Knowledge graph is not enabled for this agent.")
}
agentID := store.AgentIDFromContext(ctx)
if agentID == uuid.Nil {
return ErrorResult("agent context not available")
}
userID := store.UserIDFromContext(ctx)
query, _ := args["query"].(string)
if query == "" {
return ErrorResult("query parameter is required")
}
entityID, _ := args["entity_id"].(string)
maxDepth := 2
if md, ok := args["max_depth"].(float64); ok && md > 0 {
maxDepth = int(md)
if maxDepth > 3 {
maxDepth = 3
}
}
// Traversal mode: entity_id provided
if entityID != "" {
return t.executeTraversal(ctx, agentID.String(), userID, entityID, maxDepth)
}
// Search mode
return t.executeSearch(ctx, agentID.String(), userID, query, args)
}
func (t *KnowledgeGraphSearchTool) executeTraversal(ctx context.Context, agentID, userID, entityID string, maxDepth int) *Result {
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))
}
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")
}
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))
}
if len(entities) == 0 {
return NewResult(fmt.Sprintf("No entities found matching %q.", query))
}
// Optional type filter (post-search)
entityType, _ := args["entity_type"].(string)
if entityType != "" {
filtered := entities[:0]
for _, e := range entities {
if e.EntityType == entityType {
filtered = append(filtered, e)
}
}
entities = filtered
if len(entities) == 0 {
return NewResult(fmt.Sprintf("No entities of type %q found matching %q.", entityType, query))
}
}
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))
if e.Description != "" {
sb.WriteString(fmt.Sprintf(" %s\n", e.Description))
}
// Fetch relations to show connections
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))
}
}
}
return NewResult(sb.String())
}
+17 -2
View File
@@ -59,11 +59,16 @@ func isMemoryPath(path, workspace string) bool {
return false
}
// KGExtractFunc is a callback invoked after a memory write to extract KG entities.
// agentID, userID, content are passed from the write context.
type KGExtractFunc func(ctx context.Context, agentID, userID, content string)
// MemoryInterceptor routes memory file reads/writes to the MemoryStore.
// Keeps MEMORY.md and memory/* in Postgres.
type MemoryInterceptor struct {
memStore store.MemoryStore
workspace string
memStore store.MemoryStore
workspace string
kgExtractFn KGExtractFunc
}
// NewMemoryInterceptor creates an interceptor backed by the given memory store.
@@ -71,6 +76,11 @@ func NewMemoryInterceptor(ms store.MemoryStore, workspace string) *MemoryInterce
return &MemoryInterceptor{memStore: ms, workspace: workspace}
}
// SetKGExtractFunc sets the callback for KG extraction after memory writes.
func (m *MemoryInterceptor) SetKGExtractFunc(fn KGExtractFunc) {
m.kgExtractFn = fn
}
// ReadFile attempts to read a memory file from the DB.
// Returns (content, true, nil) if handled, or ("", false, nil) if not a memory path.
func (m *MemoryInterceptor) ReadFile(ctx context.Context, path string) (string, bool, error) {
@@ -137,6 +147,11 @@ func (m *MemoryInterceptor) WriteFile(ctx context.Context, path, content string)
}
}
// Trigger KG extraction in background if configured
if m.kgExtractFn != nil && content != "" {
go m.kgExtractFn(context.WithoutCancel(ctx), agentStr, userID, content)
}
return true, nil
}
+1 -1
View File
@@ -2,4 +2,4 @@ package upgrade
// RequiredSchemaVersion is the schema migration version this binary requires.
// Bump this whenever adding a new SQL migration file.
const RequiredSchemaVersion uint = 12
const RequiredSchemaVersion uint = 13
@@ -0,0 +1,2 @@
DROP TABLE IF EXISTS kg_relations;
DROP TABLE IF EXISTS kg_entities;
+33
View File
@@ -0,0 +1,33 @@
CREATE TABLE kg_entities (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
agent_id UUID NOT NULL REFERENCES agents(id) ON DELETE CASCADE,
user_id VARCHAR(255) NOT NULL DEFAULT '',
external_id VARCHAR(255) NOT NULL,
name TEXT NOT NULL,
entity_type VARCHAR(100) NOT NULL,
description TEXT DEFAULT '',
properties JSONB DEFAULT '{}',
source_id VARCHAR(255) DEFAULT '',
confidence FLOAT NOT NULL DEFAULT 1.0,
created_at TIMESTAMPTZ DEFAULT NOW(),
updated_at TIMESTAMPTZ DEFAULT NOW(),
UNIQUE(agent_id, user_id, external_id)
);
CREATE TABLE kg_relations (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
agent_id UUID NOT NULL REFERENCES agents(id) ON DELETE CASCADE,
user_id VARCHAR(255) NOT NULL DEFAULT '',
source_entity_id UUID NOT NULL REFERENCES kg_entities(id) ON DELETE CASCADE,
relation_type VARCHAR(200) NOT NULL,
target_entity_id UUID NOT NULL REFERENCES kg_entities(id) ON DELETE CASCADE,
confidence FLOAT NOT NULL DEFAULT 1.0,
properties JSONB DEFAULT '{}',
created_at TIMESTAMPTZ DEFAULT NOW(),
UNIQUE(agent_id, user_id, source_entity_id, relation_type, target_entity_id)
);
CREATE INDEX idx_kg_entities_scope ON kg_entities(agent_id, user_id);
CREATE INDEX idx_kg_entities_type ON kg_entities(agent_id, user_id, entity_type);
CREATE INDEX idx_kg_relations_source ON kg_relations(source_entity_id, relation_type);
CREATE INDEX idx_kg_relations_target ON kg_relations(target_entity_id);
+3
View File
@@ -16,8 +16,10 @@
"@tailwindcss/typography": "^0.5.19",
"@tanstack/react-query": "^5.90.21",
"@tanstack/react-table": "^8.21.3",
"@xyflow/react": "^12.10.1",
"class-variance-authority": "^0.7.1",
"clsx": "^2.1.1",
"d3-force": "^3.0.0",
"framer-motion": "^12.34.3",
"highlight.js": "^11.11.1",
"lucide-react": "^0.468.0",
@@ -33,6 +35,7 @@
},
"devDependencies": {
"@tailwindcss/vite": "^4.0.0",
"@types/d3-force": "^3.0.10",
"@types/react": "^19.0.0",
"@types/react-dom": "^19.0.0",
"@types/ws": "^8.18.1",
+202
View File
@@ -26,12 +26,18 @@ importers:
'@tanstack/react-table':
specifier: ^8.21.3
version: 8.21.3(react-dom@19.2.4(react@19.2.4))(react@19.2.4)
'@xyflow/react':
specifier: ^12.10.1
version: 12.10.1(@types/react@19.2.14)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)
class-variance-authority:
specifier: ^0.7.1
version: 0.7.1
clsx:
specifier: ^2.1.1
version: 2.1.1
d3-force:
specifier: ^3.0.0
version: 3.0.0
framer-motion:
specifier: ^12.34.3
version: 12.34.3(react-dom@19.2.4(react@19.2.4))(react@19.2.4)
@@ -72,6 +78,9 @@ importers:
'@tailwindcss/vite':
specifier: ^4.0.0
version: 4.2.0(vite@6.4.1(@types/node@25.3.0)(jiti@2.6.1)(lightningcss@1.31.1))
'@types/d3-force':
specifier: ^3.0.10
version: 3.0.10
'@types/react':
specifier: ^19.0.0
version: 19.2.14
@@ -1352,6 +1361,27 @@ packages:
'@types/babel__traverse@7.28.0':
resolution: {integrity: sha512-8PvcXf70gTDZBgt9ptxJ8elBeBjcLOAcOtoO/mPJjtji1+CdGbHgm77om1GrsPxsiE+uXIpNSK64UYaIwQXd4Q==}
'@types/d3-color@3.1.3':
resolution: {integrity: sha512-iO90scth9WAbmgv7ogoq57O9YpKmFBbmoEoCHDB2xMBY0+/KVrqAaCDyCE16dUspeOvIxFFRI+0sEtqDqy2b4A==}
'@types/d3-drag@3.0.7':
resolution: {integrity: sha512-HE3jVKlzU9AaMazNufooRJ5ZpWmLIoc90A37WU2JMmeq28w1FQqCZswHZ3xR+SuxYftzHq6WU6KJHvqxKzTxxQ==}
'@types/d3-force@3.0.10':
resolution: {integrity: sha512-ZYeSaCF3p73RdOKcjj+swRlZfnYpK1EbaDiYICEEp5Q6sUiqFaFQ9qgoshp5CzIyyb/yD09kD9o2zEltCexlgw==}
'@types/d3-interpolate@3.0.4':
resolution: {integrity: sha512-mgLPETlrpVV1YRJIglr4Ez47g7Yxjl1lj7YKsiMCb27VJH9W8NVM6Bb9d8kkpG/uAQS5AmbA48q2IAolKKo1MA==}
'@types/d3-selection@3.0.11':
resolution: {integrity: sha512-bhAXu23DJWsrI45xafYpkQ4NtcKMwWnAC/vKrd2l+nxMFuvOT3XMYTIj2opv8vq8AO5Yh7Qac/nSeP/3zjTK0w==}
'@types/d3-transition@3.0.9':
resolution: {integrity: sha512-uZS5shfxzO3rGlu0cC3bjmMFKsXv+SmZZcgp0KD22ts4uGXp5EVYGzu/0YdwZeKmddhcAccYtREJKkPfXkZuCg==}
'@types/d3-zoom@3.0.8':
resolution: {integrity: sha512-iqMC4/YlFCSlO8+2Ii1GGGliCAY4XdeG748w5vQUbevlbDu0zSjH/+jojorQVBK/se0j6DUFNPBGSqD3YWYnDw==}
'@types/debug@4.1.12':
resolution: {integrity: sha512-vIChWdVG3LG1SMxEvI/AK+FWJthlrqlTu7fbrlywTkkaONwk/UAGaULXRlf8vkzFBLVm0zkMdCquhL5aOjhXPQ==}
@@ -1399,6 +1429,15 @@ packages:
peerDependencies:
vite: ^4.2.0 || ^5.0.0 || ^6.0.0 || ^7.0.0
'@xyflow/react@12.10.1':
resolution: {integrity: sha512-5eSWtIK/+rkldOuFbOOz44CRgQRjtS9v5nufk77DV+XBnfCGL9HAQ8PG00o2ZYKqkEU/Ak6wrKC95Tu+2zuK3Q==}
peerDependencies:
react: '>=17'
react-dom: '>=17'
'@xyflow/system@0.0.75':
resolution: {integrity: sha512-iXs+AGFLi8w/VlAoc/iSxk+CxfT6o64Uw/k0CKASOPqjqz6E0rb5jFZgJtXGZCpfQI6OQpu5EnumP5fGxQheaQ==}
aria-hidden@1.2.6:
resolution: {integrity: sha512-ik3ZgC9dY/lYVVM++OISsaYDeg1tb0VtP5uL3ouh1koGOaUMDPpbFIei4JkFimWUFPn90sbMNMXQAIVOlnYKJA==}
engines: {node: '>=10'}
@@ -1437,6 +1476,9 @@ packages:
class-variance-authority@0.7.1:
resolution: {integrity: sha512-Ka+9Trutv7G8M6WT6SeiRWz792K5qEqIGEGzXKhAE6xOWAY6pPH8U+9IY3oCMv6kqTmLsv7Xh/2w2RigkePMsg==}
classcat@5.0.5:
resolution: {integrity: sha512-JhZUT7JFcQy/EzW605k/ktHtncoo9vnyW/2GspNYwFlN1C/WmjuV/xtS04e9SOkL2sTdw0VAZ2UGCcQ9lR6p6w==}
clsx@2.1.1:
resolution: {integrity: sha512-eYm0QWBtUrBWZWG0d386OGAw16Z995PiOVo2B7bjWSbHedGl5e0ZWaq65kOGgUSNesEIDkB9ISbTg/JK9dhCZA==}
engines: {node: '>=6'}
@@ -1459,6 +1501,52 @@ packages:
csstype@3.2.3:
resolution: {integrity: sha512-z1HGKcYy2xA8AGQfwrn0PAy+PB7X/GSj3UVJW9qKyn43xWa+gl5nXmU4qqLMRzWVLFC8KusUX8T/0kCiOYpAIQ==}
d3-color@3.1.0:
resolution: {integrity: sha512-zg/chbXyeBtMQ1LbD/WSoW2DpC3I0mpmPdW+ynRTj/x2DAWYrIY7qeZIHidozwV24m4iavr15lNwIwLxRmOxhA==}
engines: {node: '>=12'}
d3-dispatch@3.0.1:
resolution: {integrity: sha512-rzUyPU/S7rwUflMyLc1ETDeBj0NRuHKKAcvukozwhshr6g6c5d8zh4c2gQjY2bZ0dXeGLWc1PF174P2tVvKhfg==}
engines: {node: '>=12'}
d3-drag@3.0.0:
resolution: {integrity: sha512-pWbUJLdETVA8lQNJecMxoXfH6x+mO2UQo8rSmZ+QqxcbyA3hfeprFgIT//HW2nlHChWeIIMwS2Fq+gEARkhTkg==}
engines: {node: '>=12'}
d3-ease@3.0.1:
resolution: {integrity: sha512-wR/XK3D3XcLIZwpbvQwQ5fK+8Ykds1ip7A2Txe0yxncXSdq1L9skcG7blcedkOX+ZcgxGAmLX1FrRGbADwzi0w==}
engines: {node: '>=12'}
d3-force@3.0.0:
resolution: {integrity: sha512-zxV/SsA+U4yte8051P4ECydjD/S+qeYtnaIyAs9tgHCqfguma/aAQDjo85A9Z6EKhBirHRJHXIgJUlffT4wdLg==}
engines: {node: '>=12'}
d3-interpolate@3.0.1:
resolution: {integrity: sha512-3bYs1rOD33uo8aqJfKP3JWPAibgw8Zm2+L9vBKEHJ2Rg+viTR7o5Mmv5mZcieN+FRYaAOWX5SJATX6k1PWz72g==}
engines: {node: '>=12'}
d3-quadtree@3.0.1:
resolution: {integrity: sha512-04xDrxQTDTCFwP5H6hRhsRcb9xxv2RzkcsygFzmkSIOJy3PeRJP7sNk3VRIbKXcog561P9oU0/rVH6vDROAgUw==}
engines: {node: '>=12'}
d3-selection@3.0.0:
resolution: {integrity: sha512-fmTRWbNMmsmWq6xJV8D19U/gw/bwrHfNXxrIN+HfZgnzqTHp9jOmKMhsTUjXOJnZOdZY9Q28y4yebKzqDKlxlQ==}
engines: {node: '>=12'}
d3-timer@3.0.1:
resolution: {integrity: sha512-ndfJ/JxxMd3nw31uyKoY2naivF+r29V+Lc0svZxe1JvvIRmi8hUsrMvdOwgS1o6uBHmiz91geQ0ylPP0aj1VUA==}
engines: {node: '>=12'}
d3-transition@3.0.1:
resolution: {integrity: sha512-ApKvfjsSR6tg06xrL434C0WydLr7JewBB3V+/39RMHsaXTOG0zmt/OAXeng5M5LBm0ojmxJrpomQVZ1aPvBL4w==}
engines: {node: '>=12'}
peerDependencies:
d3-selection: 2 - 3
d3-zoom@3.0.0:
resolution: {integrity: sha512-b8AmV3kfQaqWAuacbPuNbL6vahnOJflOhexLzMMNLga62+/nh0JzvJ0aO/5a5MVgUFGS7Hu1P9P03o3fJkDCyw==}
engines: {node: '>=12'}
debug@4.4.3:
resolution: {integrity: sha512-RGwwWnwQvkVfavKVt22FGLw+xYSdzARwm0ru6DhTVA3umU5hZc28V3kO4stgYryrTlLpuvgI9GiijltAjNbcqA==}
engines: {node: '>=6.0'}
@@ -2131,6 +2219,21 @@ packages:
yallist@3.1.1:
resolution: {integrity: sha512-a4UGQaWPH59mOXUYnAG2ewncQS4i4F43Tv3JoAM+s2VDAmS9NsK8GpDMLrCHPksFT7h3K6TOoUNn2pb7RoXx4g==}
zustand@4.5.7:
resolution: {integrity: sha512-CHOUy7mu3lbD6o6LJLfllpjkzhHXSBlX8B9+qPddUsIfeF5S/UZ5q0kmCsnRqT1UHFQZchNFDDzMbQsuesHWlw==}
engines: {node: '>=12.7.0'}
peerDependencies:
'@types/react': '>=16.8'
immer: '>=9.0.6'
react: '>=16.8'
peerDependenciesMeta:
'@types/react':
optional: true
immer:
optional: true
react:
optional: true
zustand@5.0.11:
resolution: {integrity: sha512-fdZY+dk7zn/vbWNCYmzZULHRrss0jx5pPFiOuMZ/5HJN6Yv3u+1Wswy/4MpZEkEGhtNH+pwxZB8OKgUBPzYAGg==}
engines: {node: '>=12.20.0'}
@@ -3338,6 +3441,29 @@ snapshots:
dependencies:
'@babel/types': 7.29.0
'@types/d3-color@3.1.3': {}
'@types/d3-drag@3.0.7':
dependencies:
'@types/d3-selection': 3.0.11
'@types/d3-force@3.0.10': {}
'@types/d3-interpolate@3.0.4':
dependencies:
'@types/d3-color': 3.1.3
'@types/d3-selection@3.0.11': {}
'@types/d3-transition@3.0.9':
dependencies:
'@types/d3-selection': 3.0.11
'@types/d3-zoom@3.0.8':
dependencies:
'@types/d3-interpolate': 3.0.4
'@types/d3-selection': 3.0.11
'@types/debug@4.1.12':
dependencies:
'@types/ms': 2.1.0
@@ -3392,6 +3518,29 @@ snapshots:
transitivePeerDependencies:
- supports-color
'@xyflow/react@12.10.1(@types/react@19.2.14)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)':
dependencies:
'@xyflow/system': 0.0.75
classcat: 5.0.5
react: 19.2.4
react-dom: 19.2.4(react@19.2.4)
zustand: 4.5.7(@types/react@19.2.14)(react@19.2.4)
transitivePeerDependencies:
- '@types/react'
- immer
'@xyflow/system@0.0.75':
dependencies:
'@types/d3-drag': 3.0.7
'@types/d3-interpolate': 3.0.4
'@types/d3-selection': 3.0.11
'@types/d3-transition': 3.0.9
'@types/d3-zoom': 3.0.8
d3-drag: 3.0.0
d3-interpolate: 3.0.1
d3-selection: 3.0.0
d3-zoom: 3.0.0
aria-hidden@1.2.6:
dependencies:
tslib: 2.8.1
@@ -3424,6 +3573,8 @@ snapshots:
dependencies:
clsx: 2.1.1
classcat@5.0.5: {}
clsx@2.1.1: {}
comma-separated-tokens@2.0.3: {}
@@ -3436,6 +3587,50 @@ snapshots:
csstype@3.2.3: {}
d3-color@3.1.0: {}
d3-dispatch@3.0.1: {}
d3-drag@3.0.0:
dependencies:
d3-dispatch: 3.0.1
d3-selection: 3.0.0
d3-ease@3.0.1: {}
d3-force@3.0.0:
dependencies:
d3-dispatch: 3.0.1
d3-quadtree: 3.0.1
d3-timer: 3.0.1
d3-interpolate@3.0.1:
dependencies:
d3-color: 3.1.0
d3-quadtree@3.0.1: {}
d3-selection@3.0.0: {}
d3-timer@3.0.1: {}
d3-transition@3.0.1(d3-selection@3.0.0):
dependencies:
d3-color: 3.1.0
d3-dispatch: 3.0.1
d3-ease: 3.0.1
d3-interpolate: 3.0.1
d3-selection: 3.0.0
d3-timer: 3.0.1
d3-zoom@3.0.0:
dependencies:
d3-dispatch: 3.0.1
d3-drag: 3.0.0
d3-interpolate: 3.0.1
d3-selection: 3.0.0
d3-transition: 3.0.1(d3-selection@3.0.0)
debug@4.4.3:
dependencies:
ms: 2.1.3
@@ -4371,6 +4566,13 @@ snapshots:
yallist@3.1.1: {}
zustand@4.5.7(@types/react@19.2.14)(react@19.2.4):
dependencies:
use-sync-external-store: 1.6.0(react@19.2.4)
optionalDependencies:
'@types/react': 19.2.14
react: 19.2.4
zustand@5.0.11(@types/react@19.2.14)(react@19.2.4)(use-sync-external-store@1.6.0(react@19.2.4)):
optionalDependencies:
'@types/react': 19.2.14
+6
View File
@@ -62,4 +62,10 @@ export const queryKeys = {
all: ["memory"] as const,
list: (params: Record<string, unknown>) => ["memory", params] as const,
},
kg: {
all: ["kg"] as const,
list: (params: Record<string, unknown>) => ["kg", params] as const,
stats: (agentId: string, userId?: string) => ["kg", "stats", agentId, userId] as const,
graph: (agentId: string, userId?: string) => ["kg", "graph", agentId, userId] as const,
},
};
@@ -12,6 +12,9 @@ import { Textarea } from "@/components/ui/textarea";
import type { BuiltinToolData } from "./hooks/use-builtin-tools";
import { MEDIA_TOOLS } from "./media-provider-params-schema";
import { MediaProviderChainForm } from "./media-provider-chain-form";
import { KGSettingsForm } from "./kg-settings-form";
const KG_TOOL = "knowledge_graph_search";
interface Props {
tool: BuiltinToolData | null;
@@ -22,10 +25,12 @@ interface Props {
export function BuiltinToolSettingsDialog({ tool, open, onOpenChange, onSave }: Props) {
const isMedia = tool ? MEDIA_TOOLS.has(tool.name) : false;
const isKG = tool?.name === KG_TOOL;
const wide = isMedia || isKG;
return (
<Dialog open={open} onOpenChange={onOpenChange}>
<DialogContent className={isMedia ? "sm:max-w-2xl" : "sm:max-w-md"}>
<DialogContent className={wide ? "sm:max-w-2xl" : "sm:max-w-md"}>
{isMedia && tool ? (
<MediaProviderChainForm
toolName={tool.name}
@@ -33,6 +38,12 @@ export function BuiltinToolSettingsDialog({ tool, open, onOpenChange, onSave }:
onSave={(settings) => onSave(tool.name, settings).then(() => onOpenChange(false))}
onCancel={() => onOpenChange(false)}
/>
) : isKG && tool ? (
<KGSettingsForm
initialSettings={tool.settings ?? {}}
onSave={(settings) => onSave(tool.name, settings).then(() => onOpenChange(false))}
onCancel={() => onOpenChange(false)}
/>
) : (
<JsonSettingsForm tool={tool} onOpenChange={onOpenChange} onSave={onSave} />
)}
@@ -0,0 +1,116 @@
import { useState, useEffect } from "react";
import {
DialogHeader,
DialogTitle,
DialogDescription,
DialogFooter,
} from "@/components/ui/dialog";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Label } from "@/components/ui/label";
import { Switch } from "@/components/ui/switch";
import { ProviderModelSelect } from "@/components/shared/provider-model-select";
interface KGSettings {
extract_on_memory_write: boolean;
extraction_provider: string;
extraction_model: string;
min_confidence: number;
}
const defaultSettings: KGSettings = {
extract_on_memory_write: false,
extraction_provider: "",
extraction_model: "",
min_confidence: 0.75,
};
interface Props {
initialSettings: Record<string, unknown>;
onSave: (settings: Record<string, unknown>) => Promise<void>;
onCancel: () => void;
}
export function KGSettingsForm({ initialSettings, onSave, onCancel }: Props) {
const [settings, setSettings] = useState<KGSettings>(defaultSettings);
const [saving, setSaving] = useState(false);
useEffect(() => {
setSettings({
...defaultSettings,
...initialSettings,
min_confidence: Number(initialSettings.min_confidence) || defaultSettings.min_confidence,
});
}, [initialSettings]);
const handleSave = async () => {
setSaving(true);
try {
await onSave(settings as unknown as Record<string, unknown>);
} finally {
setSaving(false);
}
};
return (
<>
<DialogHeader>
<DialogTitle>Knowledge Graph Settings</DialogTitle>
<DialogDescription>
Configure entity extraction from memory writes. Requires a provider/model capable of structured JSON output.
</DialogDescription>
</DialogHeader>
<div className="space-y-4 py-2">
<ProviderModelSelect
provider={settings.extraction_provider}
onProviderChange={(v) => setSettings((s) => ({ ...s, extraction_provider: v }))}
model={settings.extraction_model}
onModelChange={(v) => setSettings((s) => ({ ...s, extraction_model: v }))}
providerLabel="Extraction Provider"
modelLabel="Extraction Model"
providerTip="LLM provider used to extract entities and relations from text."
modelTip="Model ID for extraction. Should support structured JSON output."
/>
<div className="grid gap-1.5">
<Label htmlFor="kg-min-conf" className="text-sm">Min Confidence</Label>
<Input
id="kg-min-conf"
type="number"
min={0}
max={1}
step={0.05}
value={settings.min_confidence}
onChange={(e) => setSettings((s) => ({ ...s, min_confidence: Number(e.target.value) || 0.75 }))}
className="max-w-[120px]"
/>
<p className="text-xs text-muted-foreground">
Entities below this confidence score are discarded (0.0–1.0).
</p>
</div>
<div className="flex items-center justify-between rounded-md border p-3">
<div>
<Label htmlFor="kg-auto-extract" className="text-sm font-medium">Auto-extract on memory write</Label>
<p className="text-xs text-muted-foreground mt-0.5">
Automatically extract entities when agents write to memory files.
</p>
</div>
<Switch
id="kg-auto-extract"
checked={settings.extract_on_memory_write}
onCheckedChange={(v) => setSettings((s) => ({ ...s, extract_on_memory_write: v }))}
/>
</div>
</div>
<DialogFooter>
<Button variant="outline" onClick={onCancel}>Cancel</Button>
<Button onClick={handleSave} disabled={saving}>
{saving ? "Saving..." : "Save"}
</Button>
</DialogFooter>
</>
);
}
@@ -0,0 +1,187 @@
import { useCallback, useState } from "react";
import { useQuery, useQueryClient } from "@tanstack/react-query";
import { useHttp } from "@/hooks/use-ws";
import { queryKeys } from "@/lib/query-keys";
import { toast } from "@/stores/use-toast-store";
import type { KGEntity, KGRelation, KGStats, KGTraversalResult } from "@/types/knowledge-graph";
export interface KGFilters {
agentId: string;
userId?: string;
entityType?: string;
query?: string;
}
export function useKnowledgeGraph(filters: KGFilters) {
const http = useHttp();
const queryClient = useQueryClient();
const queryKey = queryKeys.kg.list({ ...filters });
const { data, isLoading, isFetching } = useQuery({
queryKey,
queryFn: async () => {
if (!filters.agentId) return [];
const params: Record<string, string> = {};
if (filters.userId) params.user_id = filters.userId;
if (filters.entityType) params.type = filters.entityType;
if (filters.query) params.q = filters.query;
params.limit = "100";
return (await http.get<KGEntity[]>(`/v1/agents/${filters.agentId}/kg/entities`, params)) ?? [];
},
enabled: !!filters.agentId,
placeholderData: (prev) => prev,
});
const entities = data ?? [];
const invalidate = useCallback(
() => queryClient.invalidateQueries({ queryKey: queryKeys.kg.all }),
[queryClient],
);
const getEntityWithRelations = useCallback(
async (entityId: string, userId?: string) => {
const params: Record<string, string> = {};
if (userId) params.user_id = userId;
return http.get<{ entity: KGEntity; relations: KGRelation[] }>(
`/v1/agents/${filters.agentId}/kg/entities/${entityId}`,
params,
);
},
[http, filters.agentId],
);
const upsertEntity = useCallback(
async (entity: Partial<KGEntity>) => {
try {
await http.post(`/v1/agents/${filters.agentId}/kg/entities`, entity);
await invalidate();
toast.success("Entity saved", entity.name || "");
} catch (err) {
toast.error("Failed to save entity", err instanceof Error ? err.message : "Unknown error");
throw err;
}
},
[http, filters.agentId, invalidate],
);
const deleteEntity = useCallback(
async (entityId: string, userId?: string) => {
try {
const qs = userId ? `?user_id=${encodeURIComponent(userId)}` : "";
await http.delete(`/v1/agents/${filters.agentId}/kg/entities/${entityId}${qs}`);
await invalidate();
toast.success("Entity deleted");
} catch (err) {
toast.error("Failed to delete entity", err instanceof Error ? err.message : "Unknown error");
throw err;
}
},
[http, filters.agentId, invalidate],
);
const extractFromText = useCallback(
async (text: string, provider: string, model: string, userId?: string) => {
try {
const res = await http.post<{ entities: number; relations: number }>(
`/v1/agents/${filters.agentId}/kg/extract`,
{ text, provider, model, user_id: userId || "" },
);
await invalidate();
toast.success("Extraction complete", `${res.entities} entities, ${res.relations} relations`);
return res;
} catch (err) {
toast.error("Extraction failed", err instanceof Error ? err.message : "Unknown error");
throw err;
}
},
[http, filters.agentId, invalidate],
);
return {
entities,
loading: isLoading,
fetching: isFetching,
refresh: invalidate,
getEntityWithRelations,
upsertEntity,
deleteEntity,
extractFromText,
};
}
export function useKGStats(agentId: string, userId?: string) {
const http = useHttp();
const { data, isLoading } = useQuery({
queryKey: queryKeys.kg.stats(agentId, userId),
queryFn: async () => {
const params: Record<string, string> = {};
if (userId) params.user_id = userId;
return http.get<KGStats>(`/v1/agents/${agentId}/kg/stats`, params);
},
enabled: !!agentId,
placeholderData: (prev) => prev,
});
return { stats: data, loading: isLoading };
}
export interface KGGraphData {
entities: KGEntity[];
relations: KGRelation[];
}
export function useKGGraph(agentId: string, userId?: string) {
const http = useHttp();
const { data, isLoading, isFetching } = useQuery({
queryKey: queryKeys.kg.graph(agentId, userId),
queryFn: async () => {
if (!agentId) return { entities: [], relations: [] };
const params: Record<string, string> = { limit: "200" };
if (userId) params.user_id = userId;
return http.get<KGGraphData>(`/v1/agents/${agentId}/kg/graph`, params);
},
enabled: !!agentId,
placeholderData: (prev) => prev,
});
return {
entities: data?.entities ?? [],
relations: data?.relations ?? [],
loading: isLoading,
fetching: isFetching,
};
}
export function useKGTraversal(agentId: string) {
const http = useHttp();
const [results, setResults] = useState<KGTraversalResult[]>([]);
const [traversing, setTraversing] = useState(false);
const traverse = useCallback(
async (entityId: string, userId?: string, maxDepth?: number) => {
setTraversing(true);
try {
const res = await http.post<KGTraversalResult[]>(`/v1/agents/${agentId}/kg/traverse`, {
entity_id: entityId,
user_id: userId || "",
max_depth: maxDepth || 2,
});
setResults(res ?? []);
return res ?? [];
} catch (err) {
toast.error("Traversal failed", err instanceof Error ? err.message : "Unknown error");
setResults([]);
return [];
} finally {
setTraversing(false);
}
},
[http, agentId],
);
return { results, traversing, traverse };
}
+231
View File
@@ -0,0 +1,231 @@
import { useState } from "react";
import { Network, Trash2, Search, GitFork, Sparkles, RefreshCw, LayoutGrid, Share2 } from "lucide-react";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Badge } from "@/components/ui/badge";
import { EmptyState } from "@/components/shared/empty-state";
import { TableSkeleton } from "@/components/shared/loading-skeleton";
import { ConfirmDialog } from "@/components/shared/confirm-dialog";
import { useDeferredLoading } from "@/hooks/use-deferred-loading";
import { useKnowledgeGraph, useKGStats, useKGGraph } from "./hooks/use-knowledge-graph";
import { KGEntityDetailDialog } from "./kg-entity-detail-dialog";
import { KGExtractDialog } from "./kg-extract-dialog";
import { KGGraphView } from "./kg-graph-view";
import type { KGEntity } from "@/types/knowledge-graph";
interface KGEntitiesTabProps {
agentId: string;
userId?: string;
}
type ViewMode = "table" | "graph";
export function KGEntitiesTab({ agentId, userId }: KGEntitiesTabProps) {
const [searchQuery, setSearchQuery] = useState("");
const [appliedQuery, setAppliedQuery] = useState("");
const [viewEntity, setViewEntity] = useState<KGEntity | null>(null);
const [deleteTarget, setDeleteTarget] = useState<KGEntity | null>(null);
const [deleteLoading, setDeleteLoading] = useState(false);
const [extractOpen, setExtractOpen] = useState(false);
const [viewMode, setViewMode] = useState<ViewMode>("table");
const { entities, loading, fetching, refresh, deleteEntity, getEntityWithRelations, extractFromText } = useKnowledgeGraph({
agentId,
userId,
query: appliedQuery || undefined,
});
const { stats } = useKGStats(agentId, userId);
const graphData = useKGGraph(agentId, userId);
const showSkeleton = useDeferredLoading(loading && entities.length === 0);
const handleSearch = () => setAppliedQuery(searchQuery.trim());
const handleKeyDown = (e: React.KeyboardEvent) => {
if (e.key === "Enter") handleSearch();
};
const handleDelete = async () => {
if (!deleteTarget) return;
setDeleteLoading(true);
try {
await deleteEntity(deleteTarget.id, deleteTarget.user_id);
setDeleteTarget(null);
} finally {
setDeleteLoading(false);
}
};
const handleExtract = (text: string, provider: string, model: string) =>
extractFromText(text, provider, model, userId);
return (
<div>
{/* Stats bar */}
{stats && (
<div className="flex gap-4 text-xs text-muted-foreground mb-3">
<span>Entities: {stats.entity_count}</span>
<span>Relations: {stats.relation_count}</span>
{Object.entries(stats.entity_types).map(([type, count]) => (
<span key={type}>{type}: {count}</span>
))}
</div>
)}
{/* Search + actions */}
<div className="flex gap-2 mb-4">
<Input
value={searchQuery}
onChange={(e) => setSearchQuery(e.target.value)}
onKeyDown={handleKeyDown}
placeholder="Search entities..."
className="max-w-sm"
/>
<Button variant="outline" size="sm" onClick={handleSearch} disabled={fetching} className="gap-1 h-9">
<Search className="h-3.5 w-3.5" /> Search
</Button>
{appliedQuery && (
<Button variant="ghost" size="sm" onClick={() => { setAppliedQuery(""); setSearchQuery(""); }} className="h-9">
Clear
</Button>
)}
<div className="flex-1" />
{/* View mode toggle */}
<div className="flex rounded-md border">
<Button
variant={viewMode === "table" ? "secondary" : "ghost"}
size="sm"
onClick={() => setViewMode("table")}
className="h-9 rounded-r-none gap-1 px-2.5"
>
<LayoutGrid className="h-3.5 w-3.5" />
</Button>
<Button
variant={viewMode === "graph" ? "secondary" : "ghost"}
size="sm"
onClick={() => setViewMode("graph")}
className="h-9 rounded-l-none gap-1 px-2.5"
>
<Share2 className="h-3.5 w-3.5" />
</Button>
</div>
<Button variant="outline" size="sm" onClick={() => refresh()} disabled={fetching} className="gap-1 h-9">
<RefreshCw className={"h-3.5 w-3.5" + (fetching ? " animate-spin" : "")} /> Refresh
</Button>
<Button variant="outline" size="sm" onClick={() => setExtractOpen(true)} className="gap-1 h-9">
<Sparkles className="h-3.5 w-3.5" /> Extract
</Button>
</div>
{/* Content area */}
{viewMode === "graph" ? (
<KGGraphView
entities={graphData.entities}
relations={graphData.relations}
onEntityClick={setViewEntity}
/>
) : showSkeleton ? (
<TableSkeleton rows={5} />
) : entities.length === 0 ? (
<EmptyState
icon={Network}
title="No entities"
description={appliedQuery ? "No entities match your search." : "No knowledge graph entities for this agent yet."}
/>
) : (
<div className="rounded-md border">
<table className="w-full text-sm">
<thead>
<tr className="border-b bg-muted/50">
<th className="px-4 py-3 text-left font-medium">Name</th>
<th className="px-4 py-3 text-left font-medium">Type</th>
<th className="px-4 py-3 text-left font-medium">Description</th>
<th className="px-4 py-3 text-left font-medium">Confidence</th>
<th className="px-4 py-3 text-right font-medium">Actions</th>
</tr>
</thead>
<tbody>
{entities.map((entity) => (
<tr key={entity.id} className="border-b last:border-0 hover:bg-muted/30">
<td className="px-4 py-3">
<button
className="text-left hover:underline cursor-pointer font-medium"
onClick={() => setViewEntity(entity)}
>
{entity.name}
</button>
<p className="font-mono text-[10px] text-muted-foreground">{entity.external_id}</p>
</td>
<td className="px-4 py-3">
<Badge variant="secondary">{entity.entity_type}</Badge>
</td>
<td className="px-4 py-3 text-xs text-muted-foreground max-w-[300px] truncate">
{entity.description || "-"}
</td>
<td className="px-4 py-3">
<ConfidenceBar value={entity.confidence} />
</td>
<td className="px-4 py-3 text-right">
<div className="flex items-center justify-end gap-1">
<Button variant="ghost" size="sm" onClick={() => setViewEntity(entity)} className="gap-1">
<GitFork className="h-3.5 w-3.5" />
</Button>
<Button
variant="ghost"
size="sm"
onClick={() => setDeleteTarget(entity)}
className="gap-1 text-destructive hover:text-destructive"
>
<Trash2 className="h-3.5 w-3.5" />
</Button>
</div>
</td>
</tr>
))}
</tbody>
</table>
</div>
)}
{/* Entity detail dialog */}
<KGEntityDetailDialog
open={!!viewEntity}
onOpenChange={(open) => !open && setViewEntity(null)}
agentId={agentId}
entity={viewEntity}
getEntityWithRelations={getEntityWithRelations}
/>
{/* Extract dialog */}
<KGExtractDialog
open={extractOpen}
onOpenChange={setExtractOpen}
onExtract={handleExtract}
/>
{/* Delete confirmation */}
<ConfirmDialog
open={!!deleteTarget}
onOpenChange={(open) => !open && setDeleteTarget(null)}
title="Delete Entity"
description={`Delete "${deleteTarget?.name}"? This will also delete all associated relations.`}
confirmLabel="Delete"
variant="destructive"
onConfirm={handleDelete}
loading={deleteLoading}
/>
</div>
);
}
function ConfidenceBar({ value }: { value: number }) {
const pct = Math.round(value * 100);
return (
<div className="flex items-center gap-1">
<div className="h-1.5 w-10 rounded-full bg-muted overflow-hidden">
<div className="h-full rounded-full bg-primary" style={{ width: `${pct}%` }} />
</div>
<span className="text-[10px] text-muted-foreground">{pct}%</span>
</div>
);
}
@@ -0,0 +1,152 @@
import { useState, useEffect, useCallback } from "react";
import {
Dialog,
DialogContent,
DialogHeader,
DialogTitle,
} from "@/components/ui/dialog";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { GitFork } from "lucide-react";
import { useKGTraversal } from "./hooks/use-knowledge-graph";
import type { KGEntity, KGRelation } from "@/types/knowledge-graph";
interface KGEntityDetailDialogProps {
open: boolean;
onOpenChange: (open: boolean) => void;
agentId: string;
entity: KGEntity | null;
getEntityWithRelations: (entityId: string, userId?: string) => Promise<{ entity: KGEntity; relations: KGRelation[] }>;
}
export function KGEntityDetailDialog({ open, onOpenChange, agentId, entity, getEntityWithRelations }: KGEntityDetailDialogProps) {
const [relations, setRelations] = useState<KGRelation[]>([]);
const [loadingRels, setLoadingRels] = useState(false);
const { results: traversalResults, traversing, traverse } = useKGTraversal(agentId);
const loadRelations = useCallback(async () => {
if (!entity) return;
setLoadingRels(true);
try {
const res = await getEntityWithRelations(entity.id, entity.user_id);
setRelations(res.relations ?? []);
} catch {
setRelations([]);
} finally {
setLoadingRels(false);
}
}, [entity, getEntityWithRelations]);
useEffect(() => {
if (open && entity) {
loadRelations();
}
}, [open, entity, loadRelations]);
const handleTraverse = () => {
if (!entity) return;
traverse(entity.id, entity.user_id, 2);
};
return (
<Dialog open={open} onOpenChange={(v) => !traversing && onOpenChange(v)}>
<DialogContent className="max-w-3xl max-h-[85vh] flex flex-col">
<DialogHeader>
<DialogTitle className="flex items-center gap-2">
<span>{entity?.name}</span>
{entity && (
<Badge variant="secondary" className="text-xs">
{entity.entity_type}
</Badge>
)}
</DialogTitle>
</DialogHeader>
<div className="flex-1 min-h-0 overflow-y-auto py-2 px-1 -mx-1 space-y-4">
{/* Entity info */}
{entity && (
<div className="grid grid-cols-2 gap-2 text-xs">
<div><span className="text-muted-foreground">External ID:</span> <span className="font-mono">{entity.external_id}</span></div>
<div><span className="text-muted-foreground">Confidence:</span> {Math.round(entity.confidence * 100)}%</div>
{entity.description && (
<div className="col-span-2"><span className="text-muted-foreground">Description:</span> {entity.description}</div>
)}
{entity.source_id && (
<div className="col-span-2"><span className="text-muted-foreground">Source:</span> <span className="font-mono">{entity.source_id}</span></div>
)}
{entity.properties && Object.keys(entity.properties).length > 0 && (
<div className="col-span-2">
<span className="text-muted-foreground">Properties:</span>
<pre className="mt-1 text-xs bg-muted/50 rounded p-2 whitespace-pre-wrap">{JSON.stringify(entity.properties, null, 2)}</pre>
</div>
)}
</div>
)}
{/* Relations */}
<div>
<div className="flex items-center justify-between mb-2">
<h4 className="text-sm font-medium">Relations</h4>
<Button variant="outline" size="sm" onClick={handleTraverse} disabled={traversing} className="gap-1">
<GitFork className="h-3.5 w-3.5" />
{traversing ? "Traversing..." : "Traverse (2-hop)"}
</Button>
</div>
{loadingRels ? (
<p className="text-xs text-muted-foreground">Loading...</p>
) : relations.length === 0 ? (
<p className="text-xs text-muted-foreground">No relations found.</p>
) : (
<div className="rounded-md border">
<table className="w-full text-xs">
<thead>
<tr className="border-b bg-muted/50">
<th className="px-3 py-2 text-left font-medium">Direction</th>
<th className="px-3 py-2 text-left font-medium">Relation</th>
<th className="px-3 py-2 text-left font-medium">Target</th>
<th className="px-3 py-2 text-left font-medium">Confidence</th>
</tr>
</thead>
<tbody>
{relations.map((rel) => (
<tr key={rel.id} className="border-b last:border-0 hover:bg-muted/30">
<td className="px-3 py-2">
{rel.source_entity_id === entity?.id ? "→ outgoing" : "← incoming"}
</td>
<td className="px-3 py-2 font-mono">{rel.relation_type}</td>
<td className="px-3 py-2 font-mono text-muted-foreground">
{rel.source_entity_id === entity?.id ? rel.target_entity_id.slice(0, 8) : rel.source_entity_id.slice(0, 8)}
</td>
<td className="px-3 py-2">{Math.round(rel.confidence * 100)}%</td>
</tr>
))}
</tbody>
</table>
</div>
)}
</div>
{/* Traversal results */}
{traversalResults.length > 0 && (
<div>
<h4 className="text-sm font-medium mb-2">Traversal Results ({traversalResults.length})</h4>
<div className="space-y-1">
{traversalResults.map((tr, i) => (
<div key={i} className="flex items-center gap-2 text-xs rounded border p-2">
<Badge variant="outline" className="text-[10px]">depth {tr.depth}</Badge>
{tr.via && <span className="font-mono text-muted-foreground">—[{tr.via}]→</span>}
<span className="font-medium">{tr.entity.name}</span>
<Badge variant="secondary" className="text-[10px]">{tr.entity.entity_type}</Badge>
{tr.entity.description && (
<span className="text-muted-foreground truncate max-w-[200px]">{tr.entity.description}</span>
)}
</div>
))}
</div>
</div>
)}
</div>
</DialogContent>
</Dialog>
);
}
@@ -0,0 +1,76 @@
import { useState } from "react";
import {
Dialog,
DialogContent,
DialogFooter,
DialogHeader,
DialogTitle,
} from "@/components/ui/dialog";
import { Button } from "@/components/ui/button";
import { Textarea } from "@/components/ui/textarea";
import { ProviderModelSelect } from "@/components/shared/provider-model-select";
interface KGExtractDialogProps {
open: boolean;
onOpenChange: (open: boolean) => void;
onExtract: (text: string, provider: string, model: string) => Promise<unknown>;
}
export function KGExtractDialog({ open, onOpenChange, onExtract }: KGExtractDialogProps) {
const [text, setText] = useState("");
const [provider, setProvider] = useState("");
const [model, setModel] = useState("");
const [loading, setLoading] = useState(false);
const handleSubmit = async () => {
if (!text.trim() || !provider || !model) return;
setLoading(true);
try {
await onExtract(text.trim(), provider, model);
setText("");
onOpenChange(false);
} finally {
setLoading(false);
}
};
return (
<Dialog open={open} onOpenChange={(v) => !loading && onOpenChange(v)}>
<DialogContent className="max-w-2xl max-h-[85vh] flex flex-col">
<DialogHeader>
<DialogTitle>Extract Entities from Text</DialogTitle>
</DialogHeader>
<div className="flex-1 min-h-0 overflow-y-auto py-2 px-1 -mx-1 space-y-4">
<ProviderModelSelect
provider={provider}
onProviderChange={setProvider}
model={model}
onModelChange={setModel}
providerLabel="Extraction Provider"
modelLabel="Extraction Model"
/>
<div className="grid gap-1.5">
<label className="text-xs font-medium">Text to extract from</label>
<Textarea
value={text}
onChange={(e) => setText(e.target.value)}
placeholder="Paste conversation text, notes, or any content to extract entities and relations from..."
rows={10}
/>
</div>
</div>
<DialogFooter>
<Button variant="outline" onClick={() => onOpenChange(false)} disabled={loading}>
Cancel
</Button>
<Button onClick={handleSubmit} disabled={loading || !text.trim() || !provider || !model}>
{loading ? "Extracting..." : "Extract"}
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
);
}
+198
View File
@@ -0,0 +1,198 @@
import { useMemo, useEffect, useCallback, useRef } from "react";
import {
ReactFlow,
useNodesState,
useEdgesState,
Background,
Controls,
MiniMap,
type Node,
type Edge,
Handle,
Position,
} from "@xyflow/react";
import { forceSimulation, forceLink, forceManyBody, forceCenter, forceCollide, forceX, forceY, type SimulationNodeDatum } from "d3-force";
import "@xyflow/react/dist/style.css";
import type { KGEntity, KGRelation } from "@/types/knowledge-graph";
// Color mapping for entity types
const TYPE_COLORS: Record<string, { bg: string; border: string; text: string }> = {
person: { bg: "#dbeafe", border: "#3b82f6", text: "#1e40af" },
project: { bg: "#dcfce7", border: "#22c55e", text: "#166534" },
task: { bg: "#fef3c7", border: "#f59e0b", text: "#92400e" },
event: { bg: "#fce7f3", border: "#ec4899", text: "#9d174d" },
concept: { bg: "#e0e7ff", border: "#6366f1", text: "#3730a3" },
location: { bg: "#ccfbf1", border: "#14b8a6", text: "#115e59" },
organization: { bg: "#fee2e2", border: "#ef4444", text: "#991b1b" },
};
const DEFAULT_COLOR = { bg: "#f3f4f6", border: "#9ca3af", text: "#374151" };
function EntityNode({ data }: { data: { label: string; type: string; description?: string } }) {
const colors = TYPE_COLORS[data.type] || DEFAULT_COLOR;
return (
<>
<Handle type="target" position={Position.Top} className="!bg-transparent !border-0 !w-3 !h-3" />
<div
className="px-3 py-2 rounded-lg shadow-sm border-2 min-w-[80px] max-w-[180px] cursor-grab"
style={{ background: colors.bg, borderColor: colors.border }}
>
<div className="text-xs font-semibold truncate" style={{ color: colors.text }}>
{data.label}
</div>
<div className="text-[10px] opacity-60" style={{ color: colors.text }}>
{data.type}
</div>
</div>
<Handle type="source" position={Position.Bottom} className="!bg-transparent !border-0 !w-3 !h-3" />
</>
);
}
const nodeTypes = { entity: EntityNode };
interface SimNode extends SimulationNodeDatum {
id: string;
}
function buildGraph(entities: KGEntity[], relations: KGRelation[]) {
const entityIds = new Set(entities.map((e) => e.id));
const nodes: Node[] = entities.map((e) => ({
id: e.id,
type: "entity",
position: { x: 0, y: 0 },
data: { label: e.name, type: e.entity_type, description: e.description },
}));
const edges: Edge[] = relations
.filter((r) => entityIds.has(r.source_entity_id) && entityIds.has(r.target_entity_id))
.map((r) => ({
id: r.id,
source: r.source_entity_id,
target: r.target_entity_id,
label: r.relation_type.replace(/_/g, " "),
animated: false,
style: { stroke: "#94a3b8", strokeWidth: 1.5 },
labelStyle: { fontSize: 10, fill: "#64748b" },
labelBgStyle: { fill: "#f8fafc", stroke: "#e2e8f0" },
labelBgPadding: [4, 2] as [number, number],
labelShowBg: true,
}));
return { nodes, edges };
}
function applyForceLayout(
nodes: Node[],
edges: Edge[],
onUpdate: (positioned: Node[]) => void,
) {
if (nodes.length === 0) return () => {};
const simNodes: SimNode[] = nodes.map((n) => ({ id: n.id, x: n.position.x, y: n.position.y }));
const simLinks = edges.map((e) => ({ source: e.source, target: e.target }));
const w = 600;
const h = 400;
const simulation = forceSimulation(simNodes)
.force("link", forceLink(simLinks).id((d: any) => d.id).distance(140))
.force("charge", forceManyBody().strength(-350))
.force("center", forceCenter(w / 2, h / 2))
.force("x", forceX(w / 2).strength(0.05))
.force("y", forceY(h / 2).strength(0.05))
.force("collide", forceCollide(55));
simulation.on("tick", () => {
const positioned = nodes.map((n, i) => ({
...n,
position: { x: simNodes[i]!.x ?? 0, y: simNodes[i]!.y ?? 0 },
}));
onUpdate(positioned);
});
// Run fast to settle
simulation.alpha(1).restart();
return () => simulation.stop();
}
interface KGGraphViewProps {
entities: KGEntity[];
relations: KGRelation[];
onEntityClick?: (entity: KGEntity) => void;
}
export function KGGraphView({ entities, relations, onEntityClick }: KGGraphViewProps) {
const { nodes: initialNodes, edges: initialEdges } = useMemo(
() => buildGraph(entities, relations),
[entities, relations],
);
const [nodes, setNodes, onNodesChange] = useNodesState(initialNodes);
const [edges, setEdges, onEdgesChange] = useEdgesState(initialEdges);
const simRef = useRef<(() => void) | null>(null);
// Run force layout when data changes
useEffect(() => {
setEdges(initialEdges);
simRef.current?.();
simRef.current = applyForceLayout(initialNodes, initialEdges, setNodes);
return () => simRef.current?.();
}, [initialNodes, initialEdges, setNodes, setEdges]);
const handleNodeClick = useCallback(
(_: React.MouseEvent, node: Node) => {
if (!onEntityClick) return;
const entity = entities.find((e) => e.id === node.id);
if (entity) onEntityClick(entity);
},
[entities, onEntityClick],
);
if (entities.length === 0) {
return (
<div className="flex items-center justify-center h-[400px] text-sm text-muted-foreground">
No entities to visualize
</div>
);
}
return (
<div className="h-[500px] rounded-md border bg-background">
<ReactFlow
nodes={nodes}
edges={edges}
onNodesChange={onNodesChange}
onEdgesChange={onEdgesChange}
onNodeClick={handleNodeClick}
nodeTypes={nodeTypes}
fitView
minZoom={0.2}
maxZoom={2}
proOptions={{ hideAttribution: true }}
>
<Background gap={20} size={1} />
<Controls showInteractive={false} />
<MiniMap
nodeColor={(n) => {
const type = (n.data as any)?.type as string;
return (TYPE_COLORS[type] || DEFAULT_COLOR).border;
}}
maskColor="rgba(0,0,0,0.1)"
/>
</ReactFlow>
{/* Legend */}
<div className="flex flex-wrap gap-2 px-3 py-2 border-t text-[10px]">
{Object.entries(TYPE_COLORS).map(([type, colors]) => (
<div key={type} className="flex items-center gap-1">
<div className="w-2.5 h-2.5 rounded-full" style={{ background: colors.border }} />
<span className="text-muted-foreground">{type}</span>
</div>
))}
</div>
</div>
);
}
+29 -2
View File
@@ -1,5 +1,5 @@
import { useState, useMemo } from "react";
import { Brain, Plus, RefreshCw, Search, Database, Trash2, RotateCw } from "lucide-react";
import { Brain, Plus, RefreshCw, Search, Database, Trash2, RotateCw, Network } from "lucide-react";
import { Button } from "@/components/ui/button";
import { Badge } from "@/components/ui/badge";
import { Label } from "@/components/ui/label";
@@ -13,6 +13,7 @@ import { useMemoryDocuments } from "./hooks/use-memory";
import { MemoryDocumentDialog } from "./memory-document-dialog";
import { MemoryCreateDialog } from "./memory-create-dialog";
import { MemorySearchDialog } from "./memory-search-dialog";
import { KGEntitiesTab } from "./kg-entities-tab";
import { useMinLoading } from "@/hooks/use-min-loading";
import { useDeferredLoading } from "@/hooks/use-deferred-loading";
import type { MemoryDocument } from "@/types/memory";
@@ -27,6 +28,7 @@ export function MemoryPage() {
const [deleteLoading, setDeleteLoading] = useState(false);
const [searchOpen, setSearchOpen] = useState(false);
const [indexAllLoading, setIndexAllLoading] = useState(false);
const [activeTab, setActiveTab] = useState<"documents" | "kg">("documents");
const [page, setPage] = useState(1);
const [pageSize, setPageSize] = useState(20);
@@ -167,8 +169,33 @@ export function MemoryPage() {
)}
</div>
{/* Tab bar (only when agent selected) */}
{agentId && (
<div className="mt-4 flex border-b">
<button
className={`px-4 py-2 text-sm font-medium border-b-2 -mb-px ${activeTab === "documents" ? "border-primary text-primary" : "border-transparent text-muted-foreground hover:text-foreground"}`}
onClick={() => setActiveTab("documents")}
>
<Brain className="inline h-3.5 w-3.5 mr-1.5" />Documents
</button>
<button
className={`px-4 py-2 text-sm font-medium border-b-2 -mb-px ${activeTab === "kg" ? "border-primary text-primary" : "border-transparent text-muted-foreground hover:text-foreground"}`}
onClick={() => setActiveTab("kg")}
>
<Network className="inline h-3.5 w-3.5 mr-1.5" />Knowledge Graph
</button>
</div>
)}
{/* KG tab */}
{agentId && activeTab === "kg" && (
<div className="mt-4">
<KGEntitiesTab agentId={agentId} userId={userIdFilter || undefined} />
</div>
)}
{/* Document table */}
<div className="mt-4">
<div className="mt-4" style={{ display: activeTab === "documents" || !agentId ? undefined : "none" }}>
{showSkeleton ? (
<TableSkeleton rows={5} />
) : documents.length === 0 ? (
+39
View File
@@ -0,0 +1,39 @@
export interface KGEntity {
id: string;
agent_id: string;
user_id?: string;
external_id: string;
name: string;
entity_type: string;
description?: string;
properties?: Record<string, string>;
source_id?: string;
confidence: number;
created_at: number;
updated_at: number;
}
export interface KGRelation {
id: string;
agent_id: string;
user_id?: string;
source_entity_id: string;
relation_type: string;
target_entity_id: string;
confidence: number;
properties?: Record<string, string>;
created_at: number;
}
export interface KGTraversalResult {
entity: KGEntity;
depth: number;
path: string[];
via: string;
}
export interface KGStats {
entity_count: number;
relation_count: number;
entity_types: Record<string, number>;
}