mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-13 00:14:33 +00:00
feat(kg): add tsvector FTS, entity deduplication, and embedding fixes
- Replace ILIKE search with PostgreSQL tsvector/GIN full-text search for KG entities (migration 000031) - Add entity deduplication system with dual-threshold strategy: auto-merge at 0.98+ similarity with Jaro-Winkler name match, flag 0.90-0.98 as candidates for manual review - Add ScanDuplicates for bulk on-demand duplicate detection - Add MergeEntities with advisory lock and tenant-scoped relation re-pointing (delete-then-update to avoid ON CONFLICT on UPDATE) - Wire dedup inline after KG extraction pipeline - Fix BackfillKGEmbeddings: was failing silently due to context.Background() missing tenant_id; now runs cross-tenant - Fix BackfillKGEmbeddings: break on error → continue with failed ID tracking and max consecutive error cap - Add EmbedEntity helper; UpsertEntity now generates embeddings in background goroutine - Add HTTP endpoints: POST /kg/dedup/scan, GET /kg/dedup, POST /kg/merge, POST /kg/dedup/dismiss - Add web UI: Dedup dialog with Scan All button, side-by-side entity comparison, merge/dismiss actions - Add Jaro-Winkler similarity algorithm with 34 unit tests - Update IngestExtraction to return upserted entity IDs - Bump RequiredSchemaVersion to 31
This commit is contained in:
viettranx committed
1 parent
97f4942943
commit
2e869f4ece
22 files changed
+1693
-57
No files matched your search
+12
-1
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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.
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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) {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
@@ -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;
|
||||
@@ -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);
|
||||
@@ -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.",
|
||||
|
||||
@@ -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.",
|
||||
|
||||
@@ -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}}」？这将同时移除所有关联关系。",
|
||||
|
||||
@@ -87,5 +87,6 @@ export const queryKeys = {
|
||||
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,
|
||||
dedup: (agentId: string, userId?: string) => ["kg", "dedup", agentId, userId] as const,
|
||||
},
|
||||
};
|
||||
@@ -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<string, string> = { limit: "50" };
|
||||
if (userId) params.user_id = userId;
|
||||
return (await http.get<KGDedupCandidate[]>(`/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<KGTraversalResult[]>([]);
|
||||
|
||||
@@ -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<string | null>(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 (
|
||||
<Dialog open={open} onOpenChange={onOpenChange}>
|
||||
<DialogContent className="max-sm:inset-0 sm:max-w-2xl max-h-[80dvh] flex flex-col">
|
||||
<DialogHeader>
|
||||
<DialogTitle className="flex items-center gap-2">
|
||||
<Merge className="h-4 w-4" />
|
||||
{t("kg.dedup.title")}
|
||||
</DialogTitle>
|
||||
<DialogDescription>{t("kg.dedup.description")}</DialogDescription>
|
||||
</DialogHeader>
|
||||
|
||||
<div className="flex items-center gap-2">
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
className="gap-1.5"
|
||||
disabled={scanning}
|
||||
onClick={async () => {
|
||||
setScanning(true);
|
||||
try { await scan(); } finally { setScanning(false); }
|
||||
}}
|
||||
>
|
||||
<ScanSearch className={"h-3.5 w-3.5" + (scanning ? " animate-spin" : "")} />
|
||||
{scanning ? t("kg.dedup.scanning") : t("kg.dedup.scan")}
|
||||
</Button>
|
||||
</div>
|
||||
|
||||
<div className="flex-1 min-h-0 overflow-y-auto space-y-3">
|
||||
{loading ? (
|
||||
<p className="text-sm text-muted-foreground py-8 text-center">{t("kg.entity.loading")}</p>
|
||||
) : candidates.length === 0 ? (
|
||||
<p className="text-sm text-muted-foreground py-8 text-center">{t("kg.dedup.noCandidates")}</p>
|
||||
) : (
|
||||
candidates.map((c) => (
|
||||
<CandidateCard
|
||||
key={c.id}
|
||||
candidate={c}
|
||||
loading={actionLoading === c.id}
|
||||
disabled={!!actionLoading}
|
||||
onMerge={() => handleMerge(c)}
|
||||
onDismiss={() => handleDismiss(c.id)}
|
||||
t={t}
|
||||
/>
|
||||
))
|
||||
)}
|
||||
</div>
|
||||
|
||||
{candidates.length > 0 && (
|
||||
<p className="text-xs text-muted-foreground text-center">
|
||||
{t("kg.dedup.count", { count: candidates.length })}
|
||||
</p>
|
||||
)}
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
);
|
||||
}
|
||||
|
||||
function CandidateCard({
|
||||
candidate,
|
||||
loading,
|
||||
disabled,
|
||||
onMerge,
|
||||
onDismiss,
|
||||
t,
|
||||
}: {
|
||||
candidate: KGDedupCandidate;
|
||||
loading: boolean;
|
||||
disabled: boolean;
|
||||
onMerge: () => void;
|
||||
onDismiss: () => void;
|
||||
t: (key: string, opts?: Record<string, unknown>) => string;
|
||||
}) {
|
||||
const similarity = Math.round(candidate.similarity * 100);
|
||||
const a = candidate.entity_a;
|
||||
const b = candidate.entity_b;
|
||||
|
||||
return (
|
||||
<div className="rounded-lg border p-3 space-y-2">
|
||||
{/* Similarity badge */}
|
||||
<div className="flex items-center justify-between">
|
||||
<Badge variant={similarity >= 98 ? "default" : "secondary"}>
|
||||
{similarity}% {t("kg.dedup.similar")}
|
||||
</Badge>
|
||||
<div className="flex gap-1">
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
className="h-7 gap-1 text-xs"
|
||||
disabled={disabled}
|
||||
onClick={onDismiss}
|
||||
>
|
||||
<X className="h-3 w-3" /> {t("kg.dedup.dismiss")}
|
||||
</Button>
|
||||
<Button
|
||||
size="sm"
|
||||
className="h-7 gap-1 text-xs"
|
||||
disabled={disabled}
|
||||
onClick={onMerge}
|
||||
>
|
||||
{loading ? (
|
||||
<span className="animate-spin">...</span>
|
||||
) : (
|
||||
<Check className="h-3 w-3" />
|
||||
)}
|
||||
{t("kg.dedup.merge")}
|
||||
</Button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Side-by-side entities */}
|
||||
<div className="grid grid-cols-2 gap-2">
|
||||
<EntityCard entity={a} label="A" />
|
||||
<EntityCard entity={b} label="B" />
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
function EntityCard({ entity, label }: { entity: KGDedupCandidate["entity_a"]; label: string }) {
|
||||
return (
|
||||
<div className="rounded-md bg-muted/50 p-2 space-y-1">
|
||||
<div className="flex items-center gap-1.5">
|
||||
<span className="text-[10px] font-mono text-muted-foreground">{label}</span>
|
||||
<span className="text-sm font-medium truncate">{entity.name}</span>
|
||||
</div>
|
||||
<div className="flex items-center gap-1">
|
||||
<Badge variant="outline" className="text-[10px] h-4">{entity.entity_type}</Badge>
|
||||
<span className="text-[10px] text-muted-foreground">{Math.round(entity.confidence * 100)}%</span>
|
||||
</div>
|
||||
{entity.description && (
|
||||
<p className="text-xs text-muted-foreground line-clamp-2">{entity.description}</p>
|
||||
)}
|
||||
<p className="font-mono text-[10px] text-muted-foreground/60 truncate">{entity.external_id}</p>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -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<KGEntity | null>(null);
|
||||
const [deleteLoading, setDeleteLoading] = useState(false);
|
||||
const [extractOpen, setExtractOpen] = useState(false);
|
||||
const [dedupOpen, setDedupOpen] = useState(false);
|
||||
const [viewMode, setViewMode] = useState<ViewMode>("table");
|
||||
|
||||
const { entities, loading, fetching, refresh, deleteEntity, getEntityWithRelations, extractFromText } = useKnowledgeGraph({
|
||||
@@ -115,6 +117,9 @@ export function KGEntitiesTab({ agentId, userId }: KGEntitiesTabProps) {
|
||||
<Button variant="outline" size="sm" onClick={() => refresh()} disabled={fetching} className="gap-1 h-8 px-2.5">
|
||||
<RefreshCw className={"h-3.5 w-3.5" + (fetching ? " animate-spin" : "")} />
|
||||
</Button>
|
||||
<Button variant="outline" size="sm" onClick={() => setDedupOpen(true)} className="gap-1 h-8 px-2.5">
|
||||
<Merge className="h-3.5 w-3.5" /> {t("kg.dedup.button")}
|
||||
</Button>
|
||||
<Button variant="outline" size="sm" onClick={() => setExtractOpen(true)} className="gap-1 h-8 px-2.5">
|
||||
<Sparkles className="h-3.5 w-3.5" /> {t("kg.extract")}
|
||||
</Button>
|
||||
@@ -201,6 +206,14 @@ export function KGEntitiesTab({ agentId, userId }: KGEntitiesTabProps) {
|
||||
getEntityWithRelations={getEntityWithRelations}
|
||||
/>
|
||||
|
||||
{/* Dedup dialog */}
|
||||
<KGDedupDialog
|
||||
open={dedupOpen}
|
||||
onOpenChange={setDedupOpen}
|
||||
agentId={agentId}
|
||||
userId={userId}
|
||||
/>
|
||||
|
||||
{/* Extract dialog */}
|
||||
<KGExtractDialog
|
||||
open={extractOpen}
|
||||
|
||||
@@ -37,3 +37,12 @@ export interface KGStats {
|
||||
relation_count: number;
|
||||
entity_types: Record<string, number>;
|
||||
}
|
||||
|
||||
export interface KGDedupCandidate {
|
||||
id: string;
|
||||
entity_a: KGEntity;
|
||||
entity_b: KGEntity;
|
||||
similarity: number;
|
||||
status: string;
|
||||
created_at: number;
|
||||
}
|
||||
Reference in new issue
Block a user