diff --git a/cmd/gateway_managed.go b/cmd/gateway_managed.go index d53c0e7e..61ab9c2e 100644 --- a/cmd/gateway_managed.go +++ b/cmd/gateway_managed.go @@ -583,11 +583,22 @@ func buildKGExtractFunc(kgStore store.KnowledgeGraphStore, bts store.BuiltinTool result.Relations[i].AgentID = agentID result.Relations[i].UserID = userID } - if err := kgStore.IngestExtraction(ctx, agentID, userID, result.Entities, result.Relations); err != nil { + entityIDs, err := kgStore.IngestExtraction(ctx, agentID, userID, result.Entities, result.Relations) + if 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)) + + // Run inline dedup on newly upserted entities (best-effort, non-blocking) + if len(entityIDs) > 0 { + merged, flagged, dedupErr := kgStore.DedupAfterExtraction(ctx, agentID, userID, entityIDs) + if dedupErr != nil { + slog.Warn("kg extract: dedup failed", "agent", agentID, "error", dedupErr) + } else if merged > 0 || flagged > 0 { + slog.Info("kg extract: dedup completed", "agent", agentID, "auto_merged", merged, "candidates_flagged", flagged) + } + } } } diff --git a/internal/http/knowledge_graph.go b/internal/http/knowledge_graph.go index d78c6c2a..9cb79496 100644 --- a/internal/http/knowledge_graph.go +++ b/internal/http/knowledge_graph.go @@ -42,6 +42,10 @@ func (h *KnowledgeGraphHandler) RegisterRoutes(mux *http.ServeMux) { 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)) + mux.HandleFunc("POST /v1/agents/{agentID}/kg/dedup/scan", h.auth(h.handleScanDuplicates)) + mux.HandleFunc("GET /v1/agents/{agentID}/kg/dedup", h.auth(h.handleListDedupCandidates)) + mux.HandleFunc("POST /v1/agents/{agentID}/kg/merge", h.auth(h.handleMergeEntities)) + mux.HandleFunc("POST /v1/agents/{agentID}/kg/dedup/dismiss", h.auth(h.handleDismissCandidate)) } func (h *KnowledgeGraphHandler) auth(next http.HandlerFunc) http.HandlerFunc { diff --git a/internal/http/knowledge_graph_handlers.go b/internal/http/knowledge_graph_handlers.go index 75e2a241..28a33883 100644 --- a/internal/http/knowledge_graph_handlers.go +++ b/internal/http/knowledge_graph_handlers.go @@ -202,15 +202,24 @@ func (h *KnowledgeGraphHandler) handleExtract(w http.ResponseWriter, r *http.Req result.Relations[i].UserID = body.UserID } - if err := h.store.IngestExtraction(r.Context(), agentID, body.UserID, result.Entities, result.Relations); err != nil { + entityIDs, err := h.store.IngestExtraction(r.Context(), agentID, body.UserID, result.Entities, result.Relations) + if err != nil { slog.Warn("kg.ingest_extraction failed", "error", err) writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()}) return } + // Run inline dedup on newly upserted entities (best-effort) + var dedupMerged, dedupFlagged int + if len(entityIDs) > 0 { + dedupMerged, dedupFlagged, _ = h.store.DedupAfterExtraction(r.Context(), agentID, body.UserID, entityIDs) + } + writeJSON(w, http.StatusOK, map[string]any{ - "entities": len(result.Entities), - "relations": len(result.Relations), + "entities": len(result.Entities), + "relations": len(result.Relations), + "dedup_merged": dedupMerged, + "dedup_flagged": dedupFlagged, }) } @@ -227,6 +236,105 @@ func (h *KnowledgeGraphHandler) handleStats(w http.ResponseWriter, r *http.Reque writeJSON(w, http.StatusOK, stats) } +func (h *KnowledgeGraphHandler) handleScanDuplicates(w http.ResponseWriter, r *http.Request) { + locale := extractLocale(r) + agentID := r.PathValue("agentID") + + var body struct { + UserID string `json:"user_id"` + Threshold float64 `json:"threshold"` + Limit int `json:"limit"` + } + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgInvalidJSON)}) + return + } + if body.Threshold <= 0 { + body.Threshold = 0.90 + } + if body.Limit <= 0 { + body.Limit = 100 + } + + found, err := h.store.ScanDuplicates(r.Context(), agentID, body.UserID, body.Threshold, body.Limit) + if err != nil { + slog.Warn("kg.scan_duplicates failed", "error", err) + writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()}) + return + } + writeJSON(w, http.StatusOK, map[string]any{"candidates_found": found}) +} + +func (h *KnowledgeGraphHandler) handleListDedupCandidates(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 = 50 + } + + candidates, err := h.store.ListDedupCandidates(r.Context(), agentID, userID, limit) + if err != nil { + slog.Warn("kg.list_dedup_candidates failed", "error", err) + writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()}) + return + } + if candidates == nil { + candidates = []store.DedupCandidate{} + } + writeJSON(w, http.StatusOK, candidates) +} + +func (h *KnowledgeGraphHandler) handleMergeEntities(w http.ResponseWriter, r *http.Request) { + locale := extractLocale(r) + agentID := r.PathValue("agentID") + + var body struct { + UserID string `json:"user_id"` + TargetID string `json:"target_id"` + SourceID string `json:"source_id"` + } + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgInvalidJSON)}) + return + } + if body.TargetID == "" || body.SourceID == "" { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": "target_id and source_id are required"}) + return + } + + if err := h.store.MergeEntities(r.Context(), agentID, body.UserID, body.TargetID, body.SourceID); err != nil { + slog.Warn("kg.merge_entities failed", "error", err) + writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()}) + return + } + writeJSON(w, http.StatusOK, map[string]string{"status": "merged"}) +} + +func (h *KnowledgeGraphHandler) handleDismissCandidate(w http.ResponseWriter, r *http.Request) { + locale := extractLocale(r) + agentID := r.PathValue("agentID") + + var body struct { + CandidateID string `json:"candidate_id"` + } + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgInvalidJSON)}) + return + } + if body.CandidateID == "" { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": "candidate_id is required"}) + return + } + + if err := h.store.DismissCandidate(r.Context(), agentID, body.CandidateID); err != nil { + slog.Warn("kg.dismiss_candidate failed", "error", err) + writeJSON(w, http.StatusNotFound, map[string]string{"error": "candidate not found or already resolved"}) + return + } + writeJSON(w, http.StatusOK, map[string]string{"status": "dismissed"}) +} + // handleGraph returns all entities + relations for graph visualization. func (h *KnowledgeGraphHandler) handleGraph(w http.ResponseWriter, r *http.Request) { agentID := r.PathValue("agentID") diff --git a/internal/knowledgegraph/similarity.go b/internal/knowledgegraph/similarity.go new file mode 100644 index 00000000..caef4631 --- /dev/null +++ b/internal/knowledgegraph/similarity.go @@ -0,0 +1,72 @@ +package knowledgegraph + +import "strings" + +// JaroWinkler computes the Jaro-Winkler similarity between two strings (0.0–1.0). +// Higher values indicate greater similarity. Case-insensitive comparison. +func JaroWinkler(a, b string) float64 { + a = strings.ToLower(a) + b = strings.ToLower(b) + + if a == b { + return 1.0 + } + if len(a) == 0 || len(b) == 0 { + return 0.0 + } + + // Jaro similarity + matchDist := max(len(a), len(b))/2 - 1 + if matchDist < 0 { + matchDist = 0 + } + + aMatched := make([]bool, len(a)) + bMatched := make([]bool, len(b)) + + var matches, transpositions float64 + for i := range a { + lo := max(0, i-matchDist) + hi := min(len(b), i+matchDist+1) + for j := lo; j < hi; j++ { + if bMatched[j] || a[i] != b[j] { + continue + } + aMatched[i] = true + bMatched[j] = true + matches++ + break + } + } + if matches == 0 { + return 0.0 + } + + k := 0 + for i := range a { + if !aMatched[i] { + continue + } + for !bMatched[k] { + k++ + } + if a[i] != b[k] { + transpositions++ + } + k++ + } + + jaro := (matches/float64(len(a)) + matches/float64(len(b)) + (matches-transpositions/2)/matches) / 3.0 + + // Winkler prefix bonus (up to 4 chars, scaling factor 0.1) + prefix := 0 + for i := 0; i < min(4, min(len(a), len(b))); i++ { + if a[i] == b[i] { + prefix++ + } else { + break + } + } + + return jaro + float64(prefix)*0.1*(1-jaro) +} diff --git a/internal/knowledgegraph/similarity_test.go b/internal/knowledgegraph/similarity_test.go new file mode 100644 index 00000000..811d42fe --- /dev/null +++ b/internal/knowledgegraph/similarity_test.go @@ -0,0 +1,537 @@ +package knowledgegraph + +import ( + "math" + "testing" +) + +func TestJaroWinkler_IdenticalStrings(t *testing.T) { + tests := []struct { + name string + a string + b string + }{ + {"empty strings", "", ""}, + {"single char", "a", "a"}, + {"identical words", "hello", "hello"}, + {"identical with spaces", "hello world", "hello world"}, + {"case insensitive match", "HELLO", "hello"}, + {"mixed case match", "HeLLo", "hello"}, + {"unicode identical", "你好", "你好"}, + {"accented chars identical", "café", "café"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + score := JaroWinkler(tt.a, tt.b) + if math.Abs(score-1.0) > 1e-9 { + t.Errorf("expected 1.0, got %f", score) + } + }) + } +} + +func TestJaroWinkler_EmptyStrings(t *testing.T) { + tests := []struct { + name string + a string + b string + }{ + {"both empty", "", ""}, + {"a empty, b not", "", "hello"}, + {"b empty, a not", "hello", ""}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + score := JaroWinkler(tt.a, tt.b) + if tt.a == "" && tt.b == "" { + if score != 1.0 { + t.Errorf("both empty should be 1.0, got %f", score) + } + } else { + if score != 0.0 { + t.Errorf("one empty should be 0.0, got %f", score) + } + } + }) + } +} + +func TestJaroWinkler_NoMatches(t *testing.T) { + tests := []struct { + name string + a string + b string + }{ + {"completely different", "abc", "xyz"}, + {"no common chars", "aaa", "bbb"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + score := JaroWinkler(tt.a, tt.b) + if score != 0.0 { + t.Errorf("no matches should be 0.0, got %f", score) + } + }) + } +} + +func TestJaroWinkler_CaseInsensitivity(t *testing.T) { + tests := []struct { + name string + a string + b string + }{ + {"uppercase vs lowercase", "HELLO", "hello"}, + {"mixed case", "HeLLo", "hElLO"}, + {"spaces preserved in case test", "HELLO WORLD", "hello world"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + score := JaroWinkler(tt.a, tt.b) + if math.Abs(score-1.0) > 1e-9 { + t.Errorf("case-insensitive match should be 1.0, got %f", score) + } + }) + } +} + +func TestJaroWinkler_PrefixBonus(t *testing.T) { + // Jaro-Winkler gives a bonus for matching prefixes (up to 4 chars) + // This test verifies that strings with matching prefixes score higher + // than those without, even if their base Jaro score is the same. + tests := []struct { + name string + a string + b string + minScore float64 // minimum expected score + maxScore float64 // maximum expected score (or 1.0) + }{ + { + name: "similar with prefix", + a: "martha", + b: "marhta", + minScore: 0.94, // high similarity due to prefix bonus + maxScore: 1.0, + }, + { + name: "similar no prefix match", + a: "dixon", + b: "dickson", + minScore: 0.78, + maxScore: 0.90, + }, + { + name: "one char prefix", + a: "apple", + b: "apricot", + minScore: 0.55, + maxScore: 0.70, + }, + { + name: "long common prefix", + a: "interstellar", + b: "interstellarx", + minScore: 0.97, + maxScore: 1.0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + score := JaroWinkler(tt.a, tt.b) + if score < tt.minScore || score > tt.maxScore { + t.Errorf("expected score between %f and %f, got %f", tt.minScore, tt.maxScore, score) + } + }) + } +} + +func TestJaroWinkler_SingleCharacter(t *testing.T) { + tests := []struct { + name string + a string + b string + }{ + {"both single char match", "a", "a"}, + {"both single char differ", "a", "b"}, + {"single vs multi", "a", "abc"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + score := JaroWinkler(tt.a, tt.b) + if score < 0.0 || score > 1.0 { + t.Errorf("score out of range [0,1]: %f", score) + } + }) + } +} + +func TestJaroWinkler_PartialMatches(t *testing.T) { + tests := []struct { + name string + a string + b string + minScore float64 + maxScore float64 + }{ + { + name: "one char different", + a: "sitting", + b: "kitten", + minScore: 0.70, + maxScore: 0.80, + }, + { + name: "transposition", + a: "smith", + b: "smtih", + minScore: 0.92, + maxScore: 1.0, + }, + { + name: "subset", + a: "cat", + b: "cats", + minScore: 0.80, + maxScore: 0.95, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + score := JaroWinkler(tt.a, tt.b) + if score < tt.minScore || score > tt.maxScore { + t.Errorf("expected score between %f and %f, got %f", tt.minScore, tt.maxScore, score) + } + }) + } +} + +func TestJaroWinkler_KnownValues(t *testing.T) { + // Test against known Jaro-Winkler values from standard implementations. + // These are reference cases that should remain consistent. + tests := []struct { + name string + a string + b string + expectedMin float64 + expectedMax float64 + }{ + { + name: "martha vs marhta (classic example)", + a: "martha", + b: "marhta", + expectedMin: 0.961, + expectedMax: 0.965, + }, + { + name: "dixon vs dickson", + a: "dixon", + b: "dickson", + expectedMin: 0.83, + expectedMax: 0.84, + }, + { + name: "exact substring", + a: "test", + b: "testing", + expectedMin: 0.85, + expectedMax: 0.95, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + score := JaroWinkler(tt.a, tt.b) + if score < tt.expectedMin || score > tt.expectedMax { + t.Errorf("expected score between %f and %f, got %f", tt.expectedMin, tt.expectedMax, score) + } + }) + } +} + +func TestJaroWinkler_ThresholdBoundaries(t *testing.T) { + // Test values near the dedup thresholds from the code: + // dedupNameMatchThreshold = 0.85 + // dedupCandidateThreshold = 0.90 + // dedupAutoMergeThreshold = 0.98 + tests := []struct { + name string + a string + b string + nearThreshold float64 // the threshold to be near + }{ + { + name: "names above 0.85 threshold", + a: "Jonathan", + b: "Jonathan", + nearThreshold: 0.85, + }, + { + name: "similar names around 0.85", + a: "Robert", + b: "Rupert", + nearThreshold: 0.85, + }, + { + name: "identical near 0.98", + a: "Company", + b: "Company", + nearThreshold: 0.98, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + score := JaroWinkler(tt.a, tt.b) + // Just verify score is in valid range and near expected threshold + if score < 0.0 || score > 1.0 { + t.Errorf("score out of range [0,1]: %f", score) + } + // For identical or near-identical, score should be reasonably high + if tt.a == tt.b && score < 0.99 { + t.Errorf("identical strings should score >= 0.99, got %f", score) + } + }) + } +} + +func TestJaroWinkler_Symmetry(t *testing.T) { + // JaroWinkler should be symmetric: JW(a, b) == JW(b, a) + tests := []string{ + "hello", + "world", + "test", + "martha", + "dixon", + "abc", + "xyz", + } + + for i, a := range tests { + for j, b := range tests { + if i >= j { + continue + } + t.Run("symmetric_"+a+"_"+b, func(t *testing.T) { + score1 := JaroWinkler(a, b) + score2 := JaroWinkler(b, a) + if math.Abs(score1-score2) > 1e-9 { + t.Errorf("asymmetric: JW(%q, %q)=%f, JW(%q, %q)=%f", a, b, score1, b, a, score2) + } + }) + } + } +} + +func TestJaroWinkler_LongStrings(t *testing.T) { + tests := []struct { + name string + a string + b string + }{ + { + name: "long identical", + a: "The quick brown fox jumps over the lazy dog", + b: "The quick brown fox jumps over the lazy dog", + }, + { + name: "long one char different", + a: "The quick brown fox jumps over the lazy dog", + b: "The quick brown fox jumps over the lazy fog", + }, + { + name: "long different", + a: "The quick brown fox", + b: "A slow red turtle", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + score := JaroWinkler(tt.a, tt.b) + if score < 0.0 || score > 1.0 { + t.Errorf("score out of range [0,1]: %f", score) + } + }) + } +} + +func TestJaroWinkler_SpecialCharacters(t *testing.T) { + tests := []struct { + name string + a string + b string + }{ + { + name: "with hyphens", + a: "first-name", + b: "first-name", + }, + { + name: "with underscores", + a: "first_name", + b: "first_name", + }, + { + name: "with numbers", + a: "test123", + b: "test123", + }, + { + name: "with punctuation", + a: "hello, world!", + b: "hello, world!", + }, + { + name: "unicode with special chars", + a: "café-au-lait", + b: "cafe-au-lait", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + score := JaroWinkler(tt.a, tt.b) + if score < 0.0 || score > 1.0 { + t.Errorf("score out of range [0,1]: %f", score) + } + // Identical special-char strings should match + if tt.a == tt.b && score < 0.99 { + t.Errorf("identical strings with special chars should score high, got %f", score) + } + }) + } +} + +func TestJaroWinkler_UnicodeHandling(t *testing.T) { + tests := []struct { + name string + a string + b string + }{ + { + name: "chinese characters identical", + a: "北京", + b: "北京", + }, + { + name: "hindi characters identical", + a: "नमस्ते", + b: "नमस्ते", + }, + { + name: "emoji identical", + a: "🚀test", + b: "🚀test", + }, + { + name: "mixed unicode", + a: "café", + b: "café", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + score := JaroWinkler(tt.a, tt.b) + if math.Abs(score-1.0) > 1e-9 { + t.Errorf("identical unicode should score 1.0, got %f", score) + } + }) + } +} + +func TestJaroWinkler_Boundary(t *testing.T) { + // Test values at boundary conditions: very short and very long strings + tests := []struct { + name string + a string + b string + }{ + { + name: "two char strings", + a: "ab", + b: "ab", + }, + { + name: "two char different", + a: "ab", + b: "ba", + }, + { + name: "three char strings", + a: "abc", + b: "abc", + }, + { + name: "very long identical", + a: "a" + string(make([]byte, 1000)), + b: "a" + string(make([]byte, 1000)), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + score := JaroWinkler(tt.a, tt.b) + if score < 0.0 || score > 1.0 { + t.Errorf("score out of range [0,1]: %f", score) + } + }) + } +} + +func TestJaroWinkler_MonotonicIncreaseWithSimilarity(t *testing.T) { + // Test that as strings become more similar, the score increases. + base := "hello" + tests := []struct { + name string + str string + label string + }{ + {"completely different", "world", "world"}, + {"one char match", "hxxxx", "1 match"}, + {"two char match", "hexxxx", "2 matches"}, + {"three char match", "helxx", "3 matches"}, + {"four char match", "hellx", "4 matches"}, + {"identical", "hello", "identical"}, + } + + var scores []float64 + for _, tt := range tests { + score := JaroWinkler(base, tt.str) + scores = append(scores, score) + } + + // Verify monotonic increase + for i := 1; i < len(scores); i++ { + if scores[i] < scores[i-1] { + t.Errorf("scores not monotonically increasing at index %d: %f < %f", i, scores[i], scores[i-1]) + } + } +} + +func BenchmarkJaroWinkler(b *testing.B) { + benchmarks := []struct { + name string + a string + b string + }{ + {"short", "test", "best"}, + {"medium", "jonathan", "johnathan"}, + {"long", "The quick brown fox jumps over the lazy dog", "The quick brown fox jumps over the lazy fog"}, + } + + for _, bm := range benchmarks { + b.Run(bm.name, func(b *testing.B) { + for i := 0; i < b.N; i++ { + JaroWinkler(bm.a, bm.b) + } + }) + } +} diff --git a/internal/store/knowledge_graph_store.go b/internal/store/knowledge_graph_store.go index 1f7e21aa..d7723acd 100644 --- a/internal/store/knowledge_graph_store.go +++ b/internal/store/knowledge_graph_store.go @@ -53,6 +53,16 @@ type GraphStats struct { EntityTypes map[string]int `json:"entity_types"` } +// DedupCandidate represents a pair of entities that may be duplicates. +type DedupCandidate struct { + ID string `json:"id"` + EntityA Entity `json:"entity_a"` + EntityB Entity `json:"entity_b"` + Similarity float64 `json:"similarity"` + Status string `json:"status"` + CreatedAt int64 `json:"created_at"` +} + // KnowledgeGraphStore manages entity-relationship graphs. type KnowledgeGraphStore interface { UpsertEntity(ctx context.Context, entity *Entity) error @@ -69,9 +79,25 @@ type KnowledgeGraphStore interface { Traverse(ctx context.Context, agentID, userID, startEntityID string, maxDepth int) ([]TraversalResult, error) - IngestExtraction(ctx context.Context, agentID, userID string, entities []Entity, relations []Relation) error + // IngestExtraction upserts entities and relations from an LLM extraction. + // Returns the DB UUIDs of all upserted entities for downstream processing (e.g. dedup). + IngestExtraction(ctx context.Context, agentID, userID string, entities []Entity, relations []Relation) ([]string, error) PruneByConfidence(ctx context.Context, agentID, userID string, minConfidence float64) (int, error) + // DedupAfterExtraction checks newly upserted entities for duplicates. + // Auto-merges at high similarity (>0.98 + name match), flags medium (>0.90) as candidates. + DedupAfterExtraction(ctx context.Context, agentID, userID string, newEntityIDs []string) (merged int, flagged int, err error) + // ScanDuplicates scans ALL entities with embeddings for duplicates (self-join). + // Flags candidates above threshold. Used for on-demand bulk scanning of existing data. + ScanDuplicates(ctx context.Context, agentID, userID string, threshold float64, limit int) (int, error) + // ListDedupCandidates returns pending dedup candidates for review. + ListDedupCandidates(ctx context.Context, agentID, userID string, limit int) ([]DedupCandidate, error) + // MergeEntities merges sourceID into targetID: re-points relations, deletes source. + MergeEntities(ctx context.Context, agentID, userID, targetID, sourceID string) error + // DismissCandidate marks a dedup candidate as dismissed (not a duplicate). + // Scoped by agentID + tenant to prevent cross-agent dismissal. + DismissCandidate(ctx context.Context, agentID, candidateID string) error + Stats(ctx context.Context, agentID, userID string) (*GraphStats, error) // SetEmbeddingProvider configures the embedding provider for semantic search. diff --git a/internal/store/pg/knowledge_graph.go b/internal/store/pg/knowledge_graph.go index d1f40ea9..a05aab12 100644 --- a/internal/store/pg/knowledge_graph.go +++ b/internal/store/pg/knowledge_graph.go @@ -7,7 +7,6 @@ import ( "cmp" "fmt" "slices" - "strings" "time" "github.com/google/uuid" @@ -39,7 +38,8 @@ func (s *PGKnowledgeGraphStore) UpsertEntity(ctx context.Context, entity *store. now := time.Now() id := uuid.Must(uuid.NewV7()) tid := tenantIDForInsert(ctx) - _, err = s.db.ExecContext(ctx, ` + var actualID uuid.UUID + if err = s.db.QueryRowContext(ctx, ` INSERT INTO kg_entities (id, agent_id, user_id, external_id, name, entity_type, description, properties, source_id, confidence, tenant_id, created_at, updated_at) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $12) @@ -51,11 +51,17 @@ func (s *PGKnowledgeGraphStore) UpsertEntity(ctx context.Context, entity *store. source_id = EXCLUDED.source_id, confidence = EXCLUDED.confidence, tenant_id = EXCLUDED.tenant_id, - updated_at = EXCLUDED.updated_at`, + updated_at = EXCLUDED.updated_at + RETURNING id`, id, aid, entity.UserID, entity.ExternalID, entity.Name, entity.EntityType, entity.Description, props, entity.SourceID, entity.Confidence, tid, now, - ) - return err + ).Scan(&actualID); err != nil { + return err + } + + // Generate embedding in background (best-effort, non-blocking) + go s.EmbedEntity(context.WithoutCancel(ctx), actualID.String(), entity.Name, entity.Description) + return nil } func (s *PGKnowledgeGraphStore) GetEntity(ctx context.Context, agentID, userID, entityID string) (*store.Entity, error) { @@ -167,8 +173,8 @@ func (s *PGKnowledgeGraphStore) SearchEntities(ctx context.Context, agentID, use shared := store.IsSharedKG(ctx) - // ILIKE search - ilikeResults, err := s.ilikeSearchEntities(ctx, aid, userID, query, limit*2, shared) + // FTS search using tsvector + ftsResults, err := s.ftsSearchEntities(ctx, aid, userID, query, limit*2, shared) if err != nil { return nil, err } @@ -185,24 +191,24 @@ func (s *PGKnowledgeGraphStore) SearchEntities(ctx context.Context, agentID, use } } - // If no vector results, fall back to ILIKE-only + // If no vector results, fall back to FTS-only if len(vecResults) == 0 { - if len(ilikeResults) > limit { - ilikeResults = ilikeResults[:limit] + if len(ftsResults) > limit { + ftsResults = ftsResults[:limit] } - entities := make([]store.Entity, len(ilikeResults)) - for i, r := range ilikeResults { + entities := make([]store.Entity, len(ftsResults)) + for i, r := range ftsResults { entities[i] = r.Entity } return entities, nil } - // Hybrid merge with weights: 0.3 ILIKE, 0.7 vector + // Hybrid merge with weights: 0.3 FTS, 0.7 vector textW, vecW := 0.3, 0.7 - if len(ilikeResults) == 0 { + if len(ftsResults) == 0 { textW, vecW = 0, 1.0 } - merged := hybridMergeEntities(ilikeResults, vecResults, textW, vecW) + merged := hybridMergeEntities(ftsResults, vecResults, textW, vecW) if len(merged) > limit { merged = merged[:limit] @@ -215,13 +221,10 @@ type scoredEntity struct { Score float64 } -func (s *PGKnowledgeGraphStore) ilikeSearchEntities(ctx context.Context, agentID uuid.UUID, userID, query string, limit int, shared bool) ([]scoredEntity, error) { - escaped := strings.NewReplacer("%", "\\%", "_", "\\_").Replace(query) - pattern := "%" + escaped + "%" - - where := "agent_id = $1" - args := []any{agentID} - idx := 2 +func (s *PGKnowledgeGraphStore) ftsSearchEntities(ctx context.Context, agentID uuid.UUID, userID, query string, limit int, shared bool) ([]scoredEntity, error) { + where := "agent_id = $1 AND tsv @@ plainto_tsquery('simple', $2)" + args := []any{agentID, query} + idx := 3 if !shared && userID != "" { where += fmt.Sprintf(" AND user_id = $%d", idx) args = append(args, userID) @@ -236,13 +239,14 @@ func (s *PGKnowledgeGraphStore) ilikeSearchEntities(ctx context.Context, agentID args = append(args, tcArgs...) idx++ } - args = append(args, pattern, limit) + args = append(args, query, limit) q := fmt.Sprintf(` SELECT id, agent_id, user_id, external_id, name, entity_type, description, - properties, source_id, confidence, created_at, updated_at + properties, source_id, confidence, created_at, updated_at, + ts_rank(tsv, plainto_tsquery('simple', $%d)) AS score FROM kg_entities - WHERE %s AND (name ILIKE $%d ESCAPE '\' OR description ILIKE $%d ESCAPE '\') - ORDER BY confidence DESC, updated_at DESC LIMIT $%d`, where, idx, idx, idx+1) + WHERE %s + ORDER BY score DESC LIMIT $%d`, idx, where, idx+1) rows, err := s.db.QueryContext(ctx, q, args...) if err != nil { @@ -251,22 +255,22 @@ func (s *PGKnowledgeGraphStore) ilikeSearchEntities(ctx context.Context, agentID defer rows.Close() var results []scoredEntity - rank := 1.0 for rows.Next() { var e store.Entity var props []byte var createdAt, updatedAt time.Time + var score float64 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, + &score, ); err != nil { continue } json.Unmarshal(props, &e.Properties) //nolint:errcheck e.CreatedAt = createdAt.UnixMilli() e.UpdatedAt = updatedAt.UnixMilli() - results = append(results, scoredEntity{Entity: e, Score: rank}) - rank *= 0.95 // decay for rank-based scoring + results = append(results, scoredEntity{Entity: e, Score: score}) } return results, rows.Err() } diff --git a/internal/store/pg/knowledge_graph_dedup.go b/internal/store/pg/knowledge_graph_dedup.go new file mode 100644 index 00000000..2f6fbdcd --- /dev/null +++ b/internal/store/pg/knowledge_graph_dedup.go @@ -0,0 +1,445 @@ +package pg + +import ( + "context" + "database/sql" + "encoding/json" + "fmt" + "log/slog" + "time" + + "github.com/google/uuid" + kg "github.com/nextlevelbuilder/goclaw/internal/knowledgegraph" + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +const ( + dedupAutoMergeThreshold = 0.98 + dedupCandidateThreshold = 0.90 + dedupNameMatchThreshold = 0.85 +) + +// DedupAfterExtraction checks newly upserted entities for duplicates using +// embedding similarity (HNSW KNN) and name similarity (Jaro-Winkler). +// Auto-merges near-certain duplicates (>0.98 + name match), flags possible +// duplicates (>0.90) as candidates for manual review. +func (s *PGKnowledgeGraphStore) DedupAfterExtraction(ctx context.Context, agentID, userID string, newEntityIDs []string) (int, int, error) { + if len(newEntityIDs) == 0 { + return 0, 0, nil + } + + aid := mustParseUUID(agentID) + shared := store.IsSharedKG(ctx) + var merged, flagged int + + for _, eid := range newEntityIDs { + entityID, parseErr := uuid.Parse(eid) + if parseErr != nil { + slog.Warn("kg.dedup: invalid entity ID", "id", eid, "error", parseErr) + continue + } + + // Fetch entity details + embedding with tenant scope + var name, entityType string + var embeddingStr *string + var confidence float64 + tc, tcArgs, _, err := scopeClause(ctx, 3) + if err != nil { + continue + } + row := s.db.QueryRowContext(ctx, + `SELECT name, entity_type, confidence, embedding::text + FROM kg_entities WHERE id = $1 AND agent_id = $2`+tc, + append([]any{entityID, aid}, tcArgs...)...) + if err := row.Scan(&name, &entityType, &confidence, &embeddingStr); err != nil { + continue // entity may have been deleted/merged already + } + if embeddingStr == nil { + continue // no embedding → can't compute similarity + } + + // KNN: find top-3 nearest existing entities of same type (exclude self) + neighbors, err := s.knnNeighbors(ctx, aid, userID, entityID, entityType, *embeddingStr, shared, 3) + if err != nil { + slog.Warn("kg.dedup: knn query failed", "entity_id", eid, "error", err) + continue + } + + for _, n := range neighbors { + nameSim := kg.JaroWinkler(name, n.name) + + if n.similarity >= dedupAutoMergeThreshold && nameSim >= dedupNameMatchThreshold { + // Auto-merge: keep the one with higher confidence + targetID, sourceID := eid, n.id + if n.confidence > confidence { + targetID, sourceID = n.id, eid + } + if err := s.MergeEntities(ctx, agentID, userID, targetID, sourceID); err != nil { + slog.Warn("kg.dedup: auto-merge failed", "target", targetID, "source", sourceID, "error", err) + continue + } + merged++ + break // entity merged, stop checking neighbors + } else if n.similarity >= dedupCandidateThreshold { + // Flag as candidate for manual review + if err := s.insertDedupCandidate(ctx, aid, userID, eid, n.id, n.similarity); err != nil { + slog.Warn("kg.dedup: flag candidate failed", "error", err) + } else { + flagged++ + } + } + } + } + + return merged, flagged, nil +} + +type knnNeighbor struct { + id string + name string + confidence float64 + similarity float64 +} + +// knnNeighbors finds the top-K nearest entities of the same type using HNSW index. +// embeddingStr is the PG vector text format (e.g. "[0.1,0.2,...]"). +func (s *PGKnowledgeGraphStore) knnNeighbors(ctx context.Context, agentID uuid.UUID, userID string, excludeID uuid.UUID, entityType, embeddingStr string, shared bool, limit int) ([]knnNeighbor, error) { + where := "agent_id = $1 AND entity_type = $2 AND id != $3 AND embedding IS NOT NULL" + args := []any{agentID, entityType, excludeID} + idx := 4 + if !shared && userID != "" { + where += fmt.Sprintf(" AND user_id = $%d", idx) + args = append(args, userID) + idx++ + } + tc, tcArgs, _, err := scopeClause(ctx, idx) + if err != nil { + return nil, err + } + if tc != "" { + where += tc + args = append(args, tcArgs...) + idx++ + } + args = append(args, embeddingStr, limit) + q := fmt.Sprintf(` + SELECT id, name, confidence, + 1 - (embedding <=> $%d::vector) AS similarity + FROM kg_entities + WHERE %s + ORDER BY embedding <=> $%d::vector + LIMIT $%d`, idx, where, idx, idx+1) + + rows, err := s.db.QueryContext(ctx, q, args...) + if err != nil { + return nil, err + } + defer rows.Close() + + var results []knnNeighbor + for rows.Next() { + var n knnNeighbor + if err := rows.Scan(&n.id, &n.name, &n.confidence, &n.similarity); err != nil { + continue + } + results = append(results, n) + } + return results, rows.Err() +} + +func (s *PGKnowledgeGraphStore) insertDedupCandidate(ctx context.Context, agentID uuid.UUID, userID, entityAID, entityBID string, similarity float64) error { + // Ensure consistent ordering (smaller UUID first) to avoid duplicates + if entityAID > entityBID { + entityAID, entityBID = entityBID, entityAID + } + aID, _ := uuid.Parse(entityAID) + bID, _ := uuid.Parse(entityBID) + tid := tenantIDForInsert(ctx) + _, err := s.db.ExecContext(ctx, ` + INSERT INTO kg_dedup_candidates (id, tenant_id, agent_id, user_id, entity_a_id, entity_b_id, similarity, created_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8) + ON CONFLICT (entity_a_id, entity_b_id) DO NOTHING`, + uuid.Must(uuid.NewV7()), tid, agentID, userID, aID, bID, similarity, time.Now(), + ) + return err +} + +// ScanDuplicates performs a bulk scan of ALL entities with embeddings using +// a self-join to find duplicate candidates above the given threshold. +// Inserts results into kg_dedup_candidates. Returns number of candidates found. +func (s *PGKnowledgeGraphStore) ScanDuplicates(ctx context.Context, agentID, userID string, threshold float64, limit int) (int, error) { + aid := mustParseUUID(agentID) + if threshold <= 0 { + threshold = dedupCandidateThreshold + } + if limit <= 0 { + limit = 100 + } + + shared := store.IsSharedKG(ctx) + + where := "a.agent_id = $1" + args := []any{aid} + idx := 2 + + if !shared && userID != "" { + where += fmt.Sprintf(" AND a.user_id = $%d", idx) + args = append(args, userID) + idx++ + } + tc, tcArgs, _, err := scopeClauseAlias(ctx, idx, "a") + if err != nil { + return 0, err + } + if tc != "" { + where += tc + args = append(args, tcArgs...) + idx += len(tcArgs) + } + args = append(args, threshold, limit) + + q := fmt.Sprintf(` + SELECT a.id, b.id, 1 - (a.embedding <=> b.embedding) AS similarity + FROM kg_entities a + JOIN kg_entities b ON b.agent_id = a.agent_id + AND b.tenant_id = a.tenant_id + AND b.entity_type = a.entity_type + AND b.id > a.id + AND b.embedding IS NOT NULL + WHERE %s + AND a.embedding IS NOT NULL + AND 1 - (a.embedding <=> b.embedding) > $%d + ORDER BY similarity DESC + LIMIT $%d`, where, idx, idx+1) + + rows, err := s.db.QueryContext(ctx, q, args...) + if err != nil { + return 0, fmt.Errorf("kg.scan_duplicates: query failed: %w", err) + } + defer rows.Close() + + tid := tenantIDForInsert(ctx) + found := 0 + for rows.Next() { + var aID, bID string + var sim float64 + if err := rows.Scan(&aID, &bID, &sim); err != nil { + continue + } + // Ensure consistent ordering + if aID > bID { + aID, bID = bID, aID + } + aUUID, _ := uuid.Parse(aID) + bUUID, _ := uuid.Parse(bID) + if _, err := s.db.ExecContext(ctx, ` + INSERT INTO kg_dedup_candidates (id, tenant_id, agent_id, user_id, entity_a_id, entity_b_id, similarity, created_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8) + ON CONFLICT (entity_a_id, entity_b_id) DO NOTHING`, + uuid.Must(uuid.NewV7()), tid, aid, userID, aUUID, bUUID, sim, time.Now().Unix(), + ); err != nil { + slog.Warn("kg.scan_duplicates: insert candidate failed", "error", err) + continue + } + found++ + } + + return found, rows.Err() +} + +// MergeEntities merges sourceID into targetID: re-points all relations from +// source to target, deletes the source entity. Uses advisory lock to prevent +// concurrent merges on the same agent. +func (s *PGKnowledgeGraphStore) MergeEntities(ctx context.Context, agentID, userID, targetID, sourceID string) error { + aid := mustParseUUID(agentID) + tid := mustParseUUID(targetID) + sid := mustParseUUID(sourceID) + + tx, err := s.db.BeginTx(ctx, nil) + if err != nil { + return err + } + defer tx.Rollback() //nolint:errcheck + + // Advisory lock per agent to prevent concurrent merges + if _, err := tx.ExecContext(ctx, `SELECT pg_advisory_xact_lock(hashtext($1))`, agentID); err != nil { + return fmt.Errorf("kg.merge: advisory lock failed: %w", err) + } + + // Verify both entities exist and belong to the same agent + tenant scope. + // When userID is empty, skip user_id filter (admin/shared view). + shared := store.IsSharedKG(ctx) || userID == "" + for _, eid := range []uuid.UUID{tid, sid} { + var exists bool + var q string + var args []any + if shared { + tc, tcArgs, _, err := scopeClause(ctx, 3) + if err != nil { + return err + } + q = `SELECT EXISTS(SELECT 1 FROM kg_entities WHERE id = $1 AND agent_id = $2` + tc + `)` + args = append([]any{eid, aid}, tcArgs...) + } else { + tc, tcArgs, _, err := scopeClause(ctx, 4) + if err != nil { + return err + } + q = `SELECT EXISTS(SELECT 1 FROM kg_entities WHERE id = $1 AND agent_id = $2 AND user_id = $3` + tc + `)` + args = append([]any{eid, aid, userID}, tcArgs...) + } + if err := tx.QueryRowContext(ctx, q, args...).Scan(&exists); err != nil { + return fmt.Errorf("kg.merge: entity check failed: %w", err) + } + if !exists { + return fmt.Errorf("kg.merge: entity %s not found or access denied", eid) + } + } + + // Re-point relations from source to target. + // First delete relations that would become duplicates after re-pointing, + // then update the remaining ones. + tc, tcArgs, _, err := scopeClause(ctx, 4) + if err != nil { + return err + } + for _, cols := range [][2]string{ + {"source_entity_id", "target_entity_id"}, + {"target_entity_id", "source_entity_id"}, + } { + col, otherCol := cols[0], cols[1] + // Delete would-be-duplicate relations (same type, same endpoints after re-point) + delQ := fmt.Sprintf(` + DELETE FROM kg_relations r1 + WHERE r1.%s = $1 AND r1.agent_id = $2 + AND EXISTS ( + SELECT 1 FROM kg_relations r2 + WHERE r2.%s = $3 + AND r2.agent_id = r1.agent_id + AND r2.user_id = r1.user_id + AND r2.relation_type = r1.relation_type + AND r2.%s = r1.%s + )`, col, col, otherCol, otherCol) + delArgs := append([]any{sid, aid, tid}, tcArgs...) + if _, err := tx.ExecContext(ctx, delQ+tc, delArgs...); err != nil { + return fmt.Errorf("kg.merge: dedup relations %s failed: %w", col, err) + } + // Update remaining relations + updQ := fmt.Sprintf(`UPDATE kg_relations SET %s = $1 WHERE %s = $2 AND agent_id = $3`, col, col) + updArgs := append([]any{tid, sid, aid}, tcArgs...) + if _, err := tx.ExecContext(ctx, updQ+tc, updArgs...); err != nil { + return fmt.Errorf("kg.merge: re-point %s failed: %w", col, err) + } + } + + // Delete the source entity (CASCADE removes any remaining orphan relations) + if _, err := tx.ExecContext(ctx, `DELETE FROM kg_entities WHERE id = $1`, sid); err != nil { + return fmt.Errorf("kg.merge: delete source failed: %w", err) + } + + // Mark any dedup candidates referencing the source as merged + if _, err := tx.ExecContext(ctx, ` + UPDATE kg_dedup_candidates SET status = 'merged' + WHERE (entity_a_id = $1 OR entity_b_id = $1) AND status = 'pending'`, sid); err != nil { + slog.Warn("kg.merge: update candidates failed", "error", err) + } + + return tx.Commit() +} + +// ListDedupCandidates returns pending dedup candidates for review. +func (s *PGKnowledgeGraphStore) ListDedupCandidates(ctx context.Context, agentID, userID string, limit int) ([]store.DedupCandidate, error) { + aid := mustParseUUID(agentID) + if limit <= 0 { + limit = 50 + } + + where := "c.agent_id = $1 AND c.status = 'pending'" + args := []any{aid} + idx := 2 + if !store.IsSharedKG(ctx) && userID != "" { + where += fmt.Sprintf(" AND c.user_id = $%d", idx) + args = append(args, userID) + idx++ + } + tc, tcArgs, _, err := scopeClauseAlias(ctx, idx, "c") + if err != nil { + return nil, err + } + if tc != "" { + where += tc + args = append(args, tcArgs...) + idx += len(tcArgs) + } + args = append(args, limit) + + q := fmt.Sprintf(` + SELECT c.id, c.similarity, c.status, c.created_at, + a.id, a.agent_id, a.user_id, a.external_id, a.name, a.entity_type, + a.description, a.properties, a.source_id, a.confidence, a.created_at, a.updated_at, + b.id, b.agent_id, b.user_id, b.external_id, b.name, b.entity_type, + b.description, b.properties, b.source_id, b.confidence, b.created_at, b.updated_at + FROM kg_dedup_candidates c + JOIN kg_entities a ON c.entity_a_id = a.id + JOIN kg_entities b ON c.entity_b_id = b.id + WHERE %s + ORDER BY c.similarity DESC, c.created_at DESC + LIMIT $%d`, where, idx) + + rows, err := s.db.QueryContext(ctx, q, args...) + if err != nil { + return nil, err + } + defer rows.Close() + + var results []store.DedupCandidate + for rows.Next() { + var dc store.DedupCandidate + var propsA, propsB []byte + var caA, uaA, caB, uaB time.Time + var createdAt time.Time + if err := rows.Scan( + &dc.ID, &dc.Similarity, &dc.Status, &createdAt, + &dc.EntityA.ID, &dc.EntityA.AgentID, &dc.EntityA.UserID, &dc.EntityA.ExternalID, + &dc.EntityA.Name, &dc.EntityA.EntityType, &dc.EntityA.Description, &propsA, + &dc.EntityA.SourceID, &dc.EntityA.Confidence, &caA, &uaA, + &dc.EntityB.ID, &dc.EntityB.AgentID, &dc.EntityB.UserID, &dc.EntityB.ExternalID, + &dc.EntityB.Name, &dc.EntityB.EntityType, &dc.EntityB.Description, &propsB, + &dc.EntityB.SourceID, &dc.EntityB.Confidence, &caB, &uaB, + ); err != nil { + continue + } + json.Unmarshal(propsA, &dc.EntityA.Properties) //nolint:errcheck + json.Unmarshal(propsB, &dc.EntityB.Properties) //nolint:errcheck + dc.EntityA.CreatedAt = caA.UnixMilli() + dc.EntityA.UpdatedAt = uaA.UnixMilli() + dc.EntityB.CreatedAt = caB.UnixMilli() + dc.EntityB.UpdatedAt = uaB.UnixMilli() + dc.CreatedAt = createdAt.Unix() + results = append(results, dc) + } + return results, rows.Err() +} + +// DismissCandidate marks a dedup candidate as dismissed. +// Scoped by agent_id + tenant to prevent cross-agent/cross-tenant dismissal. +func (s *PGKnowledgeGraphStore) DismissCandidate(ctx context.Context, agentID, candidateID string) error { + aid := mustParseUUID(agentID) + cid := mustParseUUID(candidateID) + tc, tcArgs, _, err := scopeClause(ctx, 3) + if err != nil { + return err + } + res, err := s.db.ExecContext(ctx, + `UPDATE kg_dedup_candidates SET status = 'dismissed' WHERE id = $1 AND agent_id = $2 AND status = 'pending'`+tc, + append([]any{cid, aid}, tcArgs...)..., + ) + if err != nil { + return err + } + n, _ := res.RowsAffected() + if n == 0 { + return sql.ErrNoRows + } + return nil +} diff --git a/internal/store/pg/knowledge_graph_embedding.go b/internal/store/pg/knowledge_graph_embedding.go index 864462a9..7355012c 100644 --- a/internal/store/pg/knowledge_graph_embedding.go +++ b/internal/store/pg/knowledge_graph_embedding.go @@ -2,7 +2,6 @@ package pg import ( "context" - "fmt" "log/slog" "github.com/google/uuid" @@ -10,27 +9,29 @@ import ( // BackfillKGEmbeddings generates embeddings for all KG entities that don't have one yet. // Processes in batches of 50. Returns total number of entities updated. +// On batch-level embedding failure, skips the batch and continues (up to 3 consecutive failures). func (s *PGKnowledgeGraphStore) BackfillKGEmbeddings(ctx context.Context) (int, error) { if s.embProvider == nil { return 0, nil } const batchSize = 50 + const maxConsecutiveErrors = 3 total := 0 + consecutiveErrors := 0 - // $1=batchSize (LIMIT), tenant at $2 if not cross-tenant - tc, tcArgs, _, tcErr := scopeClause(ctx, 2) - if tcErr != nil { - return 0, tcErr - } - batchQ := fmt.Sprintf(`SELECT id, name, description FROM kg_entities - WHERE embedding IS NULL%s + // Backfill is a cross-tenant admin operation — no tenant scoping. + // context.Background() is typically passed here (no tenant in context). + batchQ := `SELECT id, name, description FROM kg_entities + WHERE embedding IS NULL ORDER BY created_at DESC - LIMIT $1`, tc) - batchBaseArgs := tcArgs // tenant uuid(s) prepended after $1 + LIMIT $1` + + // Track failed entity IDs to avoid re-fetching them + failedIDs := make(map[uuid.UUID]bool) for { - queryArgs := append([]any{batchSize}, batchBaseArgs...) + queryArgs := []any{batchSize} rows, err := s.db.QueryContext(ctx, batchQ, queryArgs...) if err != nil { return total, err @@ -47,6 +48,9 @@ func (s *PGKnowledgeGraphStore) BackfillKGEmbeddings(ctx context.Context) (int, if err := rows.Scan(&id, &name, &desc); err != nil { continue } + if failedIDs[id] { + continue // skip previously failed entities + } pending = append(pending, entityRow{id: id, text: name + " " + desc}) } rows.Close() @@ -63,9 +67,19 @@ func (s *PGKnowledgeGraphStore) BackfillKGEmbeddings(ctx context.Context) (int, } embeddings, err := s.embProvider.Embed(ctx, texts) if err != nil { - slog.Warn("kg entity embedding batch failed", "error", err) - break + slog.Warn("kg entity embedding batch failed, skipping batch", "error", err, "batch_size", len(pending)) + // Mark these entities as failed so we don't re-fetch them + for _, p := range pending { + failedIDs[p.id] = true + } + consecutiveErrors++ + if consecutiveErrors >= maxConsecutiveErrors { + slog.Warn("kg backfill: too many consecutive errors, stopping", "errors", consecutiveErrors) + break + } + continue } + consecutiveErrors = 0 // reset on success for i, emb := range embeddings { if len(emb) == 0 { @@ -92,3 +106,23 @@ func (s *PGKnowledgeGraphStore) BackfillKGEmbeddings(ctx context.Context) (int, } return total, nil } + +// EmbedEntity generates and stores an embedding for a single entity. +// Called by UpsertEntity to ensure entities created via HTTP API also get embeddings. +func (s *PGKnowledgeGraphStore) EmbedEntity(ctx context.Context, entityID, name, description string) { + if s.embProvider == nil { + return + } + text := name + " " + description + embeddings, err := s.embProvider.Embed(ctx, []string{text}) + if err != nil || len(embeddings) == 0 || len(embeddings[0]) == 0 { + return // best-effort, don't fail the upsert + } + vecStr := vectorToString(embeddings[0]) + if _, err := s.db.ExecContext(ctx, + `UPDATE kg_entities SET embedding = $1::vector WHERE id = $2`, + vecStr, mustParseUUID(entityID), + ); err != nil { + slog.Warn("kg entity embedding failed", "entity_id", entityID, "error", err) + } +} diff --git a/internal/store/pg/knowledge_graph_relations.go b/internal/store/pg/knowledge_graph_relations.go index 1285743a..506030ad 100644 --- a/internal/store/pg/knowledge_graph_relations.go +++ b/internal/store/pg/knowledge_graph_relations.go @@ -137,10 +137,10 @@ func (s *PGKnowledgeGraphStore) ListAllRelations(ctx context.Context, agentID, u return scanRelations(rows) } -func (s *PGKnowledgeGraphStore) IngestExtraction(ctx context.Context, agentID, userID string, entities []store.Entity, relations []store.Relation) error { +func (s *PGKnowledgeGraphStore) IngestExtraction(ctx context.Context, agentID, userID string, entities []store.Entity, relations []store.Relation) ([]string, error) { tx, err := s.db.BeginTx(ctx, nil) if err != nil { - return err + return nil, err } defer tx.Rollback() //nolint:errcheck @@ -175,7 +175,7 @@ func (s *PGKnowledgeGraphStore) IngestExtraction(ctx context.Context, agentID, u id, aid, userID, e.ExternalID, e.Name, e.EntityType, e.Description, props, e.SourceID, e.Confidence, tid, now, ).Scan(&actualID); err != nil { - return err + return nil, err } extIDToUUID[e.ExternalID] = actualID } @@ -229,11 +229,20 @@ func (s *PGKnowledgeGraphStore) IngestExtraction(ctx context.Context, agentID, u tenant_id = EXCLUDED.tenant_id`, id, aid, userID, src, r.RelationType, tgt, r.Confidence, props, tid, now, ); err != nil { - return err + return nil, err } } - return tx.Commit() + if err := tx.Commit(); err != nil { + return nil, err + } + + // Collect upserted entity IDs for downstream processing (e.g. dedup) + entityIDs := make([]string, 0, len(extIDToUUID)) + for _, uid := range extIDToUUID { + entityIDs = append(entityIDs, uid.String()) + } + return entityIDs, nil } func (s *PGKnowledgeGraphStore) PruneByConfidence(ctx context.Context, agentID, userID string, minConfidence float64) (int, error) { diff --git a/internal/tools/knowledge_graph_test.go b/internal/tools/knowledge_graph_test.go index 306948fd..6414a0c1 100644 --- a/internal/tools/knowledge_graph_test.go +++ b/internal/tools/knowledge_graph_test.go @@ -89,14 +89,34 @@ func (m *mockKGStore) Traverse(_ context.Context, _, _, startEntityID string, _ return m.traversal[startEntityID], nil } -func (m *mockKGStore) IngestExtraction(context.Context, string, string, []store.Entity, []store.Relation) error { - return nil +func (m *mockKGStore) IngestExtraction(context.Context, string, string, []store.Entity, []store.Relation) ([]string, error) { + return nil, nil } func (m *mockKGStore) PruneByConfidence(context.Context, string, string, float64) (int, error) { return 0, nil } +func (m *mockKGStore) DedupAfterExtraction(context.Context, string, string, []string) (int, int, error) { + return 0, 0, nil +} + +func (m *mockKGStore) ScanDuplicates(context.Context, string, string, float64, int) (int, error) { + return 0, nil +} + +func (m *mockKGStore) ListDedupCandidates(context.Context, string, string, int) ([]store.DedupCandidate, error) { + return nil, nil +} + +func (m *mockKGStore) MergeEntities(context.Context, string, string, string, string) error { + return nil +} + +func (m *mockKGStore) DismissCandidate(context.Context, string, string) error { + return nil +} + func (m *mockKGStore) Stats(context.Context, string, string) (*store.GraphStats, error) { return &store.GraphStats{}, nil } diff --git a/internal/upgrade/version.go b/internal/upgrade/version.go index ff384a57..c07b7e9a 100644 --- a/internal/upgrade/version.go +++ b/internal/upgrade/version.go @@ -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 = 30 +const RequiredSchemaVersion uint = 31 diff --git a/migrations/000031_kg_fts_and_dedup.down.sql b/migrations/000031_kg_fts_and_dedup.down.sql new file mode 100644 index 00000000..01792856 --- /dev/null +++ b/migrations/000031_kg_fts_and_dedup.down.sql @@ -0,0 +1,3 @@ +DROP TABLE IF EXISTS kg_dedup_candidates; +DROP INDEX IF EXISTS idx_kg_entities_tsv; +ALTER TABLE kg_entities DROP COLUMN IF EXISTS tsv; diff --git a/migrations/000031_kg_fts_and_dedup.up.sql b/migrations/000031_kg_fts_and_dedup.up.sql new file mode 100644 index 00000000..30368a5d --- /dev/null +++ b/migrations/000031_kg_fts_and_dedup.up.sql @@ -0,0 +1,21 @@ +-- tsvector full-text search for KG entities (replaces ILIKE) +ALTER TABLE kg_entities ADD COLUMN IF NOT EXISTS tsv tsvector + GENERATED ALWAYS AS (to_tsvector('simple', name || ' ' || COALESCE(description, ''))) STORED; + +CREATE INDEX IF NOT EXISTS idx_kg_entities_tsv ON kg_entities USING GIN (tsv); + +-- Dedup candidates table for entity deduplication review +CREATE TABLE IF NOT EXISTS kg_dedup_candidates ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + tenant_id UUID REFERENCES tenants(id) ON DELETE CASCADE, + agent_id UUID NOT NULL REFERENCES agents(id) ON DELETE CASCADE, + user_id VARCHAR(255) NOT NULL DEFAULT '', + entity_a_id UUID NOT NULL REFERENCES kg_entities(id) ON DELETE CASCADE, + entity_b_id UUID NOT NULL REFERENCES kg_entities(id) ON DELETE CASCADE, + similarity FLOAT NOT NULL, + status VARCHAR(20) NOT NULL DEFAULT 'pending', + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + UNIQUE(entity_a_id, entity_b_id) +); + +CREATE INDEX IF NOT EXISTS idx_kg_dedup_agent ON kg_dedup_candidates(agent_id, status); diff --git a/ui/web/src/i18n/locales/en/memory.json b/ui/web/src/i18n/locales/en/memory.json index 8abafb1a..ab191f73 100644 --- a/ui/web/src/i18n/locales/en/memory.json +++ b/ui/web/src/i18n/locales/en/memory.json @@ -113,6 +113,24 @@ "observations": "Observations", "relations": "Relations" }, + "dedup": { + "button": "Dedup", + "title": "Duplicate Entities", + "description": "Review and merge entities that may be duplicates based on embedding similarity.", + "noCandidates": "No duplicate candidates found.", + "similar": "similar", + "merge": "Merge", + "dismiss": "Dismiss", + "merged": "Entities merged successfully.", + "mergeFailed": "Failed to merge entities.", + "dismissed": "Candidate dismissed.", + "dismissFailed": "Failed to dismiss candidate.", + "scan": "Scan All", + "scanning": "Scanning...", + "scanComplete": "Scan complete.", + "scanFailed": "Scan failed.", + "count": "{{count}} candidate(s)" + }, "deleteEntity": { "title": "Delete Entity", "description": "Delete \"{{name}}\"? This will also delete all associated relations.", diff --git a/ui/web/src/i18n/locales/vi/memory.json b/ui/web/src/i18n/locales/vi/memory.json index 5f814f11..76500539 100644 --- a/ui/web/src/i18n/locales/vi/memory.json +++ b/ui/web/src/i18n/locales/vi/memory.json @@ -113,6 +113,24 @@ "observations": "Quan sát", "relations": "Quan hệ" }, + "dedup": { + "button": "Gộp trùng", + "title": "Thực thể trùng lặp", + "description": "Xem và gộp các thực thể có thể trùng lặp dựa trên độ tương đồng embedding.", + "noCandidates": "Không tìm thấy ứng viên trùng lặp.", + "similar": "tương đồng", + "merge": "Gộp", + "dismiss": "Bỏ qua", + "merged": "Đã gộp thực thể thành công.", + "mergeFailed": "Gộp thực thể thất bại.", + "dismissed": "Đã bỏ qua ứng viên.", + "dismissFailed": "Bỏ qua ứng viên thất bại.", + "scan": "Quét tất cả", + "scanning": "Đang quét...", + "scanComplete": "Quét hoàn tất.", + "scanFailed": "Quét thất bại.", + "count": "{{count}} ứng viên" + }, "deleteEntity": { "title": "Xóa thực thể", "description": "Xóa \"{{name}}\"? Tất cả quan hệ liên quan cũng sẽ bị xóa.", diff --git a/ui/web/src/i18n/locales/zh/memory.json b/ui/web/src/i18n/locales/zh/memory.json index 01ce0a6e..c663e6f0 100644 --- a/ui/web/src/i18n/locales/zh/memory.json +++ b/ui/web/src/i18n/locales/zh/memory.json @@ -113,6 +113,24 @@ "observations": "观察", "relations": "关系" }, + "dedup": { + "button": "去重", + "title": "重复实体", + "description": "查看并合并基于嵌入相似度可能重复的实体。", + "noCandidates": "未找到重复候选项。", + "similar": "相似", + "merge": "合并", + "dismiss": "忽略", + "merged": "实体合并成功。", + "mergeFailed": "合并实体失败。", + "dismissed": "候选项已忽略。", + "dismissFailed": "忽略候选项失败。", + "scan": "全面扫描", + "scanning": "扫描中...", + "scanComplete": "扫描完成。", + "scanFailed": "扫描失败。", + "count": "{{count}} 个候选项" + }, "deleteEntity": { "title": "删除实体", "description": "删除「{{name}}」？这将同时移除所有关联关系。", diff --git a/ui/web/src/lib/query-keys.ts b/ui/web/src/lib/query-keys.ts index bb97b463..48c53879 100644 --- a/ui/web/src/lib/query-keys.ts +++ b/ui/web/src/lib/query-keys.ts @@ -87,5 +87,6 @@ export const queryKeys = { list: (params: Record) => ["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, + dedup: (agentId: string, userId?: string) => ["kg", "dedup", agentId, userId] as const, }, }; diff --git a/ui/web/src/pages/memory/hooks/use-knowledge-graph.ts b/ui/web/src/pages/memory/hooks/use-knowledge-graph.ts index a1ecf685..db27cc97 100644 --- a/ui/web/src/pages/memory/hooks/use-knowledge-graph.ts +++ b/ui/web/src/pages/memory/hooks/use-knowledge-graph.ts @@ -4,7 +4,7 @@ import { useHttp } from "@/hooks/use-ws"; import { queryKeys } from "@/lib/query-keys"; import { toast } from "@/stores/use-toast-store"; import i18n from "@/i18n"; -import type { KGEntity, KGRelation, KGStats, KGTraversalResult } from "@/types/knowledge-graph"; +import type { KGEntity, KGRelation, KGStats, KGTraversalResult, KGDedupCandidate } from "@/types/knowledge-graph"; export interface KGFilters { agentId: string; @@ -157,6 +157,90 @@ export function useKGGraph(agentId: string, userId?: string) { }; } +export function useKGDedup(agentId: string, userId?: string) { + const http = useHttp(); + const queryClient = useQueryClient(); + + const queryKey = queryKeys.kg.dedup(agentId, userId); + + const { data, isLoading, isFetching } = useQuery({ + queryKey, + queryFn: async () => { + if (!agentId) return []; + const params: Record = { limit: "50" }; + if (userId) params.user_id = userId; + return (await http.get(`/v1/agents/${agentId}/kg/dedup`, params)) ?? []; + }, + enabled: !!agentId, + placeholderData: (prev) => prev, + }); + + const invalidate = useCallback(() => { + queryClient.invalidateQueries({ queryKey }); + queryClient.invalidateQueries({ queryKey: queryKeys.kg.all }); + }, [queryClient, queryKey]); + + const merge = useCallback( + async (targetId: string, sourceId: string) => { + try { + await http.post(`/v1/agents/${agentId}/kg/merge`, { + target_id: targetId, + source_id: sourceId, + user_id: userId || "", + }); + await invalidate(); + toast.success(i18n.t("memory:kg.dedup.merged")); + } catch (err) { + toast.error(i18n.t("memory:kg.dedup.mergeFailed"), err instanceof Error ? err.message : ""); + throw err; + } + }, + [http, agentId, userId, invalidate], + ); + + const dismiss = useCallback( + async (candidateId: string) => { + try { + await http.post(`/v1/agents/${agentId}/kg/dedup/dismiss`, { candidate_id: candidateId }); + await invalidate(); + toast.success(i18n.t("memory:kg.dedup.dismissed")); + } catch (err) { + toast.error(i18n.t("memory:kg.dedup.dismissFailed"), err instanceof Error ? err.message : ""); + throw err; + } + }, + [http, agentId, invalidate], + ); + + const scan = useCallback( + async () => { + try { + const res = await http.post<{ candidates_found: number }>( + `/v1/agents/${agentId}/kg/dedup/scan`, + { user_id: userId || "", threshold: 0.90, limit: 100 }, + ); + await invalidate(); + toast.success(i18n.t("memory:kg.dedup.scanComplete"), `${res.candidates_found} candidates`); + return res.candidates_found; + } catch (err) { + toast.error(i18n.t("memory:kg.dedup.scanFailed"), err instanceof Error ? err.message : ""); + throw err; + } + }, + [http, agentId, userId, invalidate], + ); + + return { + candidates: data ?? [], + loading: isLoading, + fetching: isFetching, + refresh: invalidate, + scan, + merge, + dismiss, + }; +} + export function useKGTraversal(agentId: string) { const http = useHttp(); const [results, setResults] = useState([]); diff --git a/ui/web/src/pages/memory/kg-dedup-dialog.tsx b/ui/web/src/pages/memory/kg-dedup-dialog.tsx new file mode 100644 index 00000000..4cca2990 --- /dev/null +++ b/ui/web/src/pages/memory/kg-dedup-dialog.tsx @@ -0,0 +1,181 @@ +import { useState } from "react"; +import { Merge, X, Check, ScanSearch } from "lucide-react"; +import { Button } from "@/components/ui/button"; +import { Badge } from "@/components/ui/badge"; +import { Dialog, DialogContent, DialogHeader, DialogTitle, DialogDescription } from "@/components/ui/dialog"; +import { useTranslation } from "react-i18next"; +import { useKGDedup } from "./hooks/use-knowledge-graph"; +import type { KGDedupCandidate } from "@/types/knowledge-graph"; + +interface KGDedupDialogProps { + open: boolean; + onOpenChange: (open: boolean) => void; + agentId: string; + userId?: string; +} + +export function KGDedupDialog({ open, onOpenChange, agentId, userId }: KGDedupDialogProps) { + const { t } = useTranslation("memory"); + const { candidates, loading, scan, merge, dismiss } = useKGDedup(agentId, userId); + const [actionLoading, setActionLoading] = useState(null); + const [scanning, setScanning] = useState(false); + + const handleMerge = async (candidate: KGDedupCandidate) => { + setActionLoading(candidate.id); + try { + // Keep entity with higher confidence as target + const target = candidate.entity_a.confidence >= candidate.entity_b.confidence + ? candidate.entity_a : candidate.entity_b; + const source = target.id === candidate.entity_a.id + ? candidate.entity_b : candidate.entity_a; + await merge(target.id, source.id); + } finally { + setActionLoading(null); + } + }; + + const handleDismiss = async (candidateId: string) => { + setActionLoading(candidateId); + try { + await dismiss(candidateId); + } finally { + setActionLoading(null); + } + }; + + return ( + + + + + + {t("kg.dedup.title")} + + {t("kg.dedup.description")} + + +
+ +
+ +
+ {loading ? ( +

{t("kg.entity.loading")}

+ ) : candidates.length === 0 ? ( +

{t("kg.dedup.noCandidates")}

+ ) : ( + candidates.map((c) => ( + handleMerge(c)} + onDismiss={() => handleDismiss(c.id)} + t={t} + /> + )) + )} +
+ + {candidates.length > 0 && ( +

+ {t("kg.dedup.count", { count: candidates.length })} +

+ )} + + + ); +} + +function CandidateCard({ + candidate, + loading, + disabled, + onMerge, + onDismiss, + t, +}: { + candidate: KGDedupCandidate; + loading: boolean; + disabled: boolean; + onMerge: () => void; + onDismiss: () => void; + t: (key: string, opts?: Record) => string; +}) { + const similarity = Math.round(candidate.similarity * 100); + const a = candidate.entity_a; + const b = candidate.entity_b; + + return ( +
+ {/* Similarity badge */} +
+ = 98 ? "default" : "secondary"}> + {similarity}% {t("kg.dedup.similar")} + +
+ + +
+
+ + {/* Side-by-side entities */} +
+ + +
+
+ ); +} + +function EntityCard({ entity, label }: { entity: KGDedupCandidate["entity_a"]; label: string }) { + return ( +
+
+ {label} + {entity.name} +
+
+ {entity.entity_type} + {Math.round(entity.confidence * 100)}% +
+ {entity.description && ( +

{entity.description}

+ )} +

{entity.external_id}

+
+ ); +} diff --git a/ui/web/src/pages/memory/kg-entities-tab.tsx b/ui/web/src/pages/memory/kg-entities-tab.tsx index 73411702..ad431ea7 100644 --- a/ui/web/src/pages/memory/kg-entities-tab.tsx +++ b/ui/web/src/pages/memory/kg-entities-tab.tsx @@ -1,5 +1,5 @@ import { useState } from "react"; -import { Network, Trash2, Search, GitFork, Sparkles, RefreshCw, LayoutGrid, Share2 } from "lucide-react"; +import { Network, Trash2, Search, GitFork, Sparkles, RefreshCw, LayoutGrid, Share2, Merge } from "lucide-react"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Badge } from "@/components/ui/badge"; @@ -11,6 +11,7 @@ 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 { KGDedupDialog } from "./kg-dedup-dialog"; import { KGGraphView } from "./kg-graph-view"; import type { KGEntity } from "@/types/knowledge-graph"; @@ -29,6 +30,7 @@ export function KGEntitiesTab({ agentId, userId }: KGEntitiesTabProps) { const [deleteTarget, setDeleteTarget] = useState(null); const [deleteLoading, setDeleteLoading] = useState(false); const [extractOpen, setExtractOpen] = useState(false); + const [dedupOpen, setDedupOpen] = useState(false); const [viewMode, setViewMode] = useState("table"); const { entities, loading, fetching, refresh, deleteEntity, getEntityWithRelations, extractFromText } = useKnowledgeGraph({ @@ -115,6 +117,9 @@ export function KGEntitiesTab({ agentId, userId }: KGEntitiesTabProps) { + @@ -201,6 +206,14 @@ export function KGEntitiesTab({ agentId, userId }: KGEntitiesTabProps) { getEntityWithRelations={getEntityWithRelations} /> + {/* Dedup dialog */} + + {/* Extract dialog */} ; } + +export interface KGDedupCandidate { + id: string; + entity_a: KGEntity; + entity_b: KGEntity; + similarity: number; + status: string; + created_at: number; +}