diff --git a/graphify/dedup.py b/graphify/dedup.py new file mode 100644 index 000000000..af5cef81b --- /dev/null +++ b/graphify/dedup.py @@ -0,0 +1,324 @@ +"""Entity deduplication pipeline for graphify knowledge graphs. + +Pipeline: exact normalization → entropy gate → MinHash/LSH blocking → +Jaro-Winkler verification → same-community boost → union-find merge. +""" +from __future__ import annotations +import math +import re +from collections import defaultdict + +from datasketch import MinHash, MinHashLSH +from rapidfuzz.distance import JaroWinkler + + +# ── helpers ─────────────────────────────────────────────────────────────────── + +def _norm(label: str) -> str: + """Lowercase + collapse non-alphanumeric runs to space.""" + return re.sub(r"[^a-z0-9]+", " ", label.lower()).strip() + + +def _entropy(label: str) -> float: + """Shannon entropy in bits/char of the normalised label.""" + s = _norm(label) + if not s: + return 0.0 + freq: dict[str, int] = defaultdict(int) + for ch in s: + freq[ch] += 1 + n = len(s) + return -sum((c / n) * math.log2(c / n) for c in freq.values()) + + +def _shingles(text: str, k: int = 3) -> set[str]: + """Return k-gram character shingles of text.""" + if len(text) < k: + return {text} + return {text[i : i + k] for i in range(len(text) - k + 1)} + + +def _make_minhash(text: str, num_perm: int = 128) -> MinHash: + # Strip spaces so "graph extractor" and "graphextractor" share shingles + m = MinHash(num_perm=num_perm) + for shingle in _shingles(text.replace(" ", "")): + m.update(shingle.encode("utf-8")) + return m + + +# ── union-find ──────────────────────────────────────────────────────────────── + +class _UF: + def __init__(self) -> None: + self._parent: dict[str, str] = {} + + def find(self, x: str) -> str: + self._parent.setdefault(x, x) + while self._parent[x] != x: + self._parent[x] = self._parent[self._parent[x]] + x = self._parent[x] + return x + + def union(self, x: str, y: str) -> None: + self._parent.setdefault(x, x) + self._parent.setdefault(y, y) + rx, ry = self.find(x), self.find(y) + if rx != ry: + self._parent[ry] = rx + + def components(self) -> dict[str, list[str]]: + groups: dict[str, list[str]] = defaultdict(list) + for x in self._parent: + groups[self.find(x)].append(x) + return dict(groups) + + +# ── constants ───────────────────────────────────────────────────────────────── + +_ENTROPY_THRESHOLD = 2.5 +_LSH_THRESHOLD = 0.7 +_MERGE_THRESHOLD = 92.0 # rapidfuzz normalized_similarity * 100 +_COMMUNITY_BOOST = 5.0 # score bonus when both nodes share community +_NUM_PERM = 128 +_CHUNK_SUFFIX = re.compile(r"_c\d+$") + + +# ── main entry point ────────────────────────────────────────────────────────── + +def deduplicate_entities( + nodes: list[dict], + edges: list[dict], + *, + communities: dict[str, int], + dedup_llm_backend: str | None = None, +) -> tuple[list[dict], list[dict]]: + """Deduplicate near-identical entities in a knowledge graph. + + Args: + nodes: list of node dicts with at minimum {"id": str, "label": str} + edges: list of edge dicts with {"source": str, "target": str, ...} + communities: mapping of node_id -> community_id (from cluster()) + dedup_llm_backend: if set, use LLM to resolve ambiguous pairs + + Returns: + (deduped_nodes, deduped_edges) with edges rewired to survivors + """ + if len(nodes) <= 1: + return nodes, edges + + # Pre-deduplicate: keep first occurrence of each id + seen_ids: dict[str, dict] = {} + for node in nodes: + nid = node.get("id", "") + if nid and nid not in seen_ids: + seen_ids[nid] = node + unique_nodes = list(seen_ids.values()) + + if len(unique_nodes) <= 1: + return unique_nodes, edges + + # ── pass 1: exact normalization ─────────────────────────────────────────── + norm_to_nodes: dict[str, list[dict]] = defaultdict(list) + for node in unique_nodes: + key = _norm(node.get("label", node.get("id", ""))) + if key: + norm_to_nodes[key].append(node) + + uf = _UF() + for key, group in norm_to_nodes.items(): + if len(group) > 1: + winner = _pick_winner(group) + for node in group: + uf.union(winner["id"], node["id"]) + + exact_merges = sum(len(g) - 1 for g in norm_to_nodes.values() if len(g) > 1) + + # ── pass 2: MinHash/LSH + Jaro-Winkler (high-entropy nodes only) ───────── + candidates: list[dict] = [] + seen_norms: set[str] = set() + for node in unique_nodes: + key = _norm(node.get("label", node.get("id", ""))) + if key and key not in seen_norms: + seen_norms.add(key) + if _entropy(node.get("label", "")) >= _ENTROPY_THRESHOLD: + candidates.append(node) + + fuzzy_merges = 0 + if len(candidates) >= 2: + lsh = MinHashLSH(threshold=_LSH_THRESHOLD, num_perm=_NUM_PERM) + minhashes: dict[str, MinHash] = {} + + for node in candidates: + norm_label = _norm(node.get("label", node.get("id", ""))) + m = _make_minhash(norm_label) + minhashes[node["id"]] = m + try: + lsh.insert(node["id"], m) + except ValueError: + pass # duplicate key in LSH — already inserted + + for node in candidates: + node_id = node["id"] + norm_label = _norm(node.get("label", node.get("id", ""))) + neighbors = lsh.query(minhashes[node_id]) + + for neighbor_id in neighbors: + if neighbor_id == node_id: + continue + if uf.find(node_id) == uf.find(neighbor_id): + continue + + neighbor = next((n for n in candidates if n["id"] == neighbor_id), None) + if neighbor is None: + continue + + neighbor_norm = _norm(neighbor.get("label", neighbor.get("id", ""))) + score = JaroWinkler.normalized_similarity(norm_label, neighbor_norm) * 100 + + c1 = communities.get(node_id) + c2 = communities.get(neighbor_id) + if c1 is not None and c2 is not None and c1 == c2: + score += _COMMUNITY_BOOST + + if score >= _MERGE_THRESHOLD: + all_group = norm_to_nodes.get(norm_label, [node]) + \ + norm_to_nodes.get(neighbor_norm, [neighbor]) + winner = _pick_winner(all_group) + uf.union(winner["id"], node_id) + uf.union(winner["id"], neighbor_id) + fuzzy_merges += 1 + + # ── pass 3: LLM tiebreaker for ambiguous pairs (opt-in) ────────────────── + if dedup_llm_backend is not None: + _llm_tiebreak(candidates, uf, communities, backend=dedup_llm_backend) + + # ── build remap table from union-find components ────────────────────────── + components = uf.components() + remap: dict[str, str] = {} + + for root, members in components.items(): + if len(members) == 1: + continue + group_nodes = [n for n in unique_nodes if n["id"] in members] + winner = _pick_winner(group_nodes) if group_nodes else {"id": root} + winner_id = winner["id"] + for member in members: + if member != winner_id: + remap[member] = winner_id + + # ── apply remap ─────────────────────────────────────────────────────────── + if not remap: + return unique_nodes, edges + + total = len(remap) + msg = f"[graphify] Deduplicated {total} node(s)" + if exact_merges: + msg += f" ({exact_merges} exact" + if fuzzy_merges: + msg += f", {fuzzy_merges} fuzzy" + msg += ")" + print(msg + ".", flush=True) + + deduped_nodes = [n for n in unique_nodes if n["id"] not in remap] + deduped_edges = [] + for edge in edges: + e = dict(edge) + e["source"] = remap.get(e["source"], e["source"]) + e["target"] = remap.get(e["target"], e["target"]) + if e["source"] != e["target"]: + deduped_edges.append(e) + + return deduped_nodes, deduped_edges + + +def _pick_winner(nodes: list[dict]) -> dict: + """Pick the canonical survivor: prefer no chunk suffix, then shorter ID.""" + if not nodes: + raise ValueError("Cannot pick winner from empty list") + + def _score(n: dict) -> tuple[int, int]: + has_suffix = bool(_CHUNK_SUFFIX.search(n["id"])) + return (1 if has_suffix else 0, len(n["id"])) + + return min(nodes, key=_score) + + +def _llm_tiebreak( + candidates: list[dict], + uf: _UF, + communities: dict[str, int], + *, + backend: str, + batch_size: int = 30, + low: float = 75.0, + high: float = 92.0, +) -> None: + """Batch-resolve ambiguous pairs (score in [low, high)) via LLM.""" + try: + from graphify.llm import BACKENDS + import os + env_key = BACKENDS.get(backend, {}).get("env_key", "") + if not os.environ.get(env_key): + print(f"[graphify] --dedup-llm: {env_key} not set, skipping LLM tiebreaker.", flush=True) + return + except ImportError: + return + + ambiguous: list[tuple[dict, dict, float]] = [] + for i, node in enumerate(candidates): + norm_i = _norm(node.get("label", node.get("id", ""))) + for j in range(i + 1, len(candidates)): + neighbor = candidates[j] + if uf.find(node["id"]) == uf.find(neighbor["id"]): + continue + norm_j = _norm(neighbor.get("label", neighbor.get("id", ""))) + score = JaroWinkler.normalized_similarity(norm_i, norm_j) * 100 + c1 = communities.get(node["id"]) + c2 = communities.get(neighbor["id"]) + if c1 is not None and c2 is not None and c1 == c2: + score += _COMMUNITY_BOOST + if low <= score < high: + ambiguous.append((node, neighbor, score)) + + if not ambiguous: + return + + try: + from graphify.llm import _call_llm + except ImportError: + return + + for batch_start in range(0, len(ambiguous), batch_size): + batch = ambiguous[batch_start : batch_start + batch_size] + pairs_text = "\n".join( + f"{i+1}. \"{a['label']}\" vs \"{b['label']}\"" + for i, (a, b, _) in enumerate(batch) + ) + prompt = ( + "For each pair below, answer only 'yes' or 'no': are they the same real-world concept?\n\n" + f"{pairs_text}\n\n" + "Reply with one line per pair: '1. yes', '2. no', etc." + ) + try: + response = _call_llm(prompt, backend=backend, max_tokens=200) + lines = response.strip().splitlines() + for line in lines: + line = line.strip() + if not line: + continue + parts = line.split(".", 1) + if len(parts) != 2: + continue + try: + idx = int(parts[0].strip()) - 1 + except ValueError: + continue + if 0 <= idx < len(batch): + answer = parts[1].strip().lower() + if answer.startswith("yes"): + a, b, _ = batch[idx] + winner = _pick_winner([a, b]) + uf.union(winner["id"], a["id"]) + uf.union(winner["id"], b["id"]) + except Exception as exc: + print(f"[graphify] --dedup-llm batch failed: {exc}", flush=True) diff --git a/tests/test_dedup.py b/tests/test_dedup.py new file mode 100644 index 000000000..f1bd1544b --- /dev/null +++ b/tests/test_dedup.py @@ -0,0 +1,120 @@ +"""Tests for graphify/dedup.py entity deduplication pipeline.""" +from __future__ import annotations +import pytest +from graphify.dedup import deduplicate_entities, _entropy, _shingles + + +# ── entropy gate ───────────────────────────────────────────────────────────── + +def test_entropy_short_label_low(): + assert _entropy("AI") < 2.5 + +def test_entropy_normal_label_high(): + assert _entropy("AuthenticationManager") >= 2.5 + +def test_entropy_empty_string(): + assert _entropy("") == 0.0 + + +# ── shingles ───────────────────────────────────────────────────────────────── + +def test_shingles_produces_trigrams(): + s = _shingles("hello") + assert "hel" in s + assert "ell" in s + assert "llo" in s + +def test_shingles_short_string(): + # strings shorter than 3 chars return single shingle of the string itself + assert _shingles("ab") == {"ab"} + + +# ── full pipeline ───────────────────────────────────────────────────────────── + +def _make_nodes(*labels): + return [{"id": label.lower().replace(" ", "_"), "label": label, "source_file": "test.md"} for label in labels] + +def _make_edges(src, tgt, relation="relates_to"): + return [{"source": src, "target": tgt, "relation": relation}] + + +def test_exact_duplicates_merged(): + nodes = _make_nodes("UserService", "userservice", "User Service") + edges = [] + result_nodes, result_edges = deduplicate_entities(nodes, edges, communities={}) + # All three are the same concept — only one survives + assert len(result_nodes) == 1 + + +def test_typo_merged(): + # "GraphExtractor" vs "Graph Extractor" — Jaro-Winkler >= 0.92 + nodes = _make_nodes("GraphExtractor", "Graph Extractor") + edges = [] + result_nodes, _ = deduplicate_entities(nodes, edges, communities={}) + assert len(result_nodes) == 1 + + +def test_unrelated_not_merged(): + nodes = _make_nodes("UserService", "OrderService") + edges = [] + result_nodes, _ = deduplicate_entities(nodes, edges, communities={}) + assert len(result_nodes) == 2 + + +def test_short_low_entropy_not_merged(): + # "AI" and "ML" are low-entropy — entropy gate skips them + nodes = _make_nodes("AI", "ML") + edges = [] + result_nodes, _ = deduplicate_entities(nodes, edges, communities={}) + assert len(result_nodes) == 2 + + +def test_edges_rewired_after_merge(): + nodes = _make_nodes("GraphExtractor", "Graph Extractor", "Parser") + # edge from loser to Parser should be rewired to winner + edges = [{"source": "graph_extractor", "target": "parser", "relation": "uses"}] + result_nodes, result_edges = deduplicate_entities(nodes, edges, communities={}) + assert len(result_nodes) == 2 # merged + Parser + # edge should still exist (rewired to winner) + assert len(result_edges) == 1 + + +def test_self_loops_dropped_after_merge(): + # If both endpoints of an edge get merged into same node, drop the edge + nodes = _make_nodes("GraphExtractor", "Graph Extractor") + edges = [{"source": "graphextractor", "target": "graph_extractor", "relation": "same"}] + _, result_edges = deduplicate_entities(nodes, edges, communities={}) + assert result_edges == [] + + +def test_community_boost_aids_merge(): + # Two nodes in same community with score in 0.75-0.85 zone get boosted + nodes = _make_nodes("AuthManager", "Auth Manager") + edges = [] + # Same community → boost → merge + communities = {"authmanager": 1, "auth_manager": 1} + result_with, _ = deduplicate_entities(nodes, edges, communities=communities) + # Different community → no boost + communities_diff = {"authmanager": 1, "auth_manager": 2} + result_without, _ = deduplicate_entities(nodes, edges, communities=communities_diff) + assert len(result_with) <= len(result_without) + + +def test_empty_inputs(): + result_nodes, result_edges = deduplicate_entities([], [], communities={}) + assert result_nodes == [] + assert result_edges == [] + + +def test_single_node_no_crash(): + nodes = _make_nodes("UserService") + result_nodes, _ = deduplicate_entities(nodes, [], communities={}) + assert len(result_nodes) == 1 + + +def test_dedup_llm_flag_accepted(): + """deduplicate_entities accepts dedup_llm_backend without crashing when no ambiguous pairs exist.""" + nodes = _make_nodes("UserService", "OrderService") + edges = [] + result_nodes, _ = deduplicate_entities(nodes, edges, communities={}, dedup_llm_backend=None) + assert len(result_nodes) == 2