mirror of
https://github.com/safishamsi/graphify.git
synced 2026-07-15 03:47:05 +00:00
Safi
354 lines
15 KiB
Python
354 lines
15 KiB
Python
# assemble node+edge dicts into a NetworkX graph, preserving edge direction
|
||
#
|
||
# Node deduplication — three layers:
|
||
#
|
||
# 1. Within a file (AST): each extractor tracks a `seen_ids` set. A node ID is
|
||
# emitted at most once per file, so duplicate class/function definitions in
|
||
# the same source file are collapsed to the first occurrence.
|
||
#
|
||
# 2. Between files (build): NetworkX G.add_node() is idempotent — calling it
|
||
# twice with the same ID overwrites the attributes with the second call's
|
||
# values. Nodes are added in extraction order (AST first, then semantic),
|
||
# so if the same entity is extracted by both passes the semantic node
|
||
# silently overwrites the AST node. This is intentional: semantic nodes
|
||
# carry richer labels and cross-file context, while AST nodes have precise
|
||
# source_location. If you need to change the priority, reorder extractions
|
||
# passed to build().
|
||
#
|
||
# 3. Semantic merge (skill): before calling build(), the skill merges cached
|
||
# and new semantic results using an explicit `seen` set keyed on node["id"],
|
||
# so duplicates across cache hits and new extractions are resolved there
|
||
# before any graph construction happens.
|
||
#
|
||
from __future__ import annotations
|
||
import json
|
||
import re
|
||
import sys
|
||
import unicodedata
|
||
from pathlib import Path
|
||
import networkx as nx
|
||
from .validate import validate_extraction
|
||
|
||
|
||
# Synonym mapper for known invalid file_type values that LLM subagents commonly
|
||
# emit. Keeps semantic intent close (markdown→document, tool→code) and falls
|
||
# back to "concept" for any other invalid value (see #840).
|
||
_FILE_TYPE_SYNONYMS = {
|
||
"markdown": "document",
|
||
"text": "document",
|
||
"tool": "code",
|
||
"library": "code",
|
||
"pattern": "concept",
|
||
"principle": "concept",
|
||
"constraint": "concept",
|
||
"tech": "concept",
|
||
"technology": "concept",
|
||
"data-source": "concept",
|
||
"data_source": "concept",
|
||
"gotcha": "concept",
|
||
"framework": "concept",
|
||
}
|
||
|
||
|
||
def _normalize_id(s: str) -> str:
|
||
r"""Normalize an ID string the same way extract._make_id does.
|
||
|
||
Used to reconcile edge endpoints when the LLM generates IDs with slightly
|
||
different punctuation or casing than the AST extractor. Must stay in sync
|
||
with extract._make_id — NFKC normalization, \w with re.UNICODE, underscore
|
||
collapse, and casefold must all match (#811).
|
||
"""
|
||
s = unicodedata.normalize("NFKC", s)
|
||
cleaned = re.sub(r"[^\w]+", "_", s, flags=re.UNICODE)
|
||
cleaned = re.sub(r"_+", "_", cleaned)
|
||
return cleaned.strip("_").casefold()
|
||
|
||
|
||
def _norm_source_file(p: str | None) -> str | None:
|
||
"""Normalize path separators to forward slashes so Windows backslash paths
|
||
and POSIX paths from semantic subagents resolve to the same node identity."""
|
||
return p.replace("\\", "/") if p else p
|
||
|
||
|
||
def edge_data(G: nx.Graph, u: str, v: str) -> dict:
|
||
"""Return one edge attribute dict for (u, v), tolerating MultiGraph.
|
||
|
||
For MultiGraph/MultiDiGraph there can be multiple parallel edges;
|
||
this returns the first one (sufficient for callers that only need
|
||
relation/confidence for rendering). Fixes #796.
|
||
"""
|
||
raw = G[u][v]
|
||
if isinstance(G, (nx.MultiGraph, nx.MultiDiGraph)):
|
||
return next(iter(raw.values()), {})
|
||
return raw
|
||
|
||
|
||
def edge_datas(G: nx.Graph, u: str, v: str) -> list[dict]:
|
||
"""Return every edge attribute dict for (u, v); always a list."""
|
||
raw = G[u][v]
|
||
if isinstance(G, (nx.MultiGraph, nx.MultiDiGraph)):
|
||
return list(raw.values())
|
||
return [raw]
|
||
|
||
|
||
def build_from_json(extraction: dict, *, directed: bool = False) -> nx.Graph:
|
||
"""Build a NetworkX graph from an extraction dict.
|
||
|
||
directed=True produces a DiGraph that preserves edge direction (source→target).
|
||
directed=False (default) produces an undirected Graph for backward compatibility.
|
||
"""
|
||
# NetworkX <= 3.1 serialised edges as "links"; remap to "edges" for compatibility.
|
||
if "edges" not in extraction and "links" in extraction:
|
||
extraction = dict(extraction, edges=extraction["links"])
|
||
|
||
# Canonicalize legacy node/edge schema before validation.
|
||
for node in extraction.get("nodes", []):
|
||
if not isinstance(node, dict):
|
||
continue
|
||
if "source" in node and "source_file" not in node:
|
||
# Count edges that reference this node so the warning is actionable (#479)
|
||
node_id = node.get("id", "?")
|
||
affected_edges = sum(
|
||
1 for e in extraction.get("edges", [])
|
||
if e.get("source") == node_id or e.get("target") == node_id
|
||
)
|
||
print(
|
||
f"[graphify] WARNING: node '{node_id}' uses field 'source' instead of "
|
||
f"'source_file' — {affected_edges} edge(s) may be misrouted. "
|
||
f"Rename the field to 'source_file' to silence this warning.",
|
||
file=sys.stderr,
|
||
)
|
||
node["source_file"] = node.pop("source")
|
||
# Default missing/None file_type to "concept" so legacy graph.json
|
||
# entries (and stub nodes preserved by `_rebuild_code` from older
|
||
# graphify versions that didn't always populate file_type) don't
|
||
# trigger spurious "invalid file_type 'None'" validator warnings (#660).
|
||
if node.get("file_type") in (None, ""):
|
||
node["file_type"] = "concept"
|
||
ft = node.get("file_type", "")
|
||
if ft and ft not in {"code", "document", "paper", "image", "rationale", "concept"}:
|
||
node["file_type"] = _FILE_TYPE_SYNONYMS.get(ft, "concept")
|
||
|
||
errors = validate_extraction(extraction)
|
||
# Dangling edges (stdlib/external imports) are expected - only warn about real schema errors.
|
||
real_errors = [e for e in errors if "does not match any node id" not in e]
|
||
if real_errors:
|
||
print(f"[graphify] Extraction warning ({len(real_errors)} issues): {real_errors[0]}", file=sys.stderr)
|
||
G: nx.Graph = nx.DiGraph() if directed else nx.Graph()
|
||
for node in extraction.get("nodes", []):
|
||
if "source_file" in node:
|
||
node["source_file"] = _norm_source_file(node["source_file"])
|
||
G.add_node(node["id"], **{k: v for k, v in node.items() if k != "id"})
|
||
node_set = set(G.nodes())
|
||
# Normalized ID map: lets edges survive when the LLM generates IDs with
|
||
# slightly different casing or punctuation than the AST extractor.
|
||
# e.g. "Session_ValidateToken" maps to "session_validatetoken".
|
||
norm_to_id: dict[str, str] = {_normalize_id(nid): nid for nid in node_set}
|
||
for edge in extraction.get("edges", []):
|
||
if "source" not in edge and "from" in edge:
|
||
edge["source"] = edge["from"]
|
||
if "target" not in edge and "to" in edge:
|
||
edge["target"] = edge["to"]
|
||
if "source" not in edge or "target" not in edge:
|
||
continue
|
||
src, tgt = edge["source"], edge["target"]
|
||
# Remap mismatched IDs via normalization before dropping the edge.
|
||
if src not in node_set:
|
||
src = norm_to_id.get(_normalize_id(src), src)
|
||
if tgt not in node_set:
|
||
tgt = norm_to_id.get(_normalize_id(tgt), tgt)
|
||
if src not in node_set or tgt not in node_set:
|
||
continue # skip edges to external/stdlib nodes - expected, not an error
|
||
attrs = {k: v for k, v in edge.items() if k not in ("source", "target")}
|
||
if "source_file" in attrs:
|
||
attrs["source_file"] = _norm_source_file(attrs["source_file"])
|
||
# Preserve original edge direction - undirected graphs lose it otherwise,
|
||
# causing display functions to show edges backwards.
|
||
attrs["_src"] = src
|
||
attrs["_tgt"] = tgt
|
||
G.add_edge(src, tgt, **attrs)
|
||
hyperedges = extraction.get("hyperedges", [])
|
||
if hyperedges:
|
||
G.graph["hyperedges"] = hyperedges
|
||
return G
|
||
|
||
|
||
def build(
|
||
extractions: list[dict],
|
||
*,
|
||
directed: bool = False,
|
||
dedup: bool = True,
|
||
dedup_llm_backend: str | None = None,
|
||
) -> nx.Graph:
|
||
"""Merge multiple extraction results into one graph.
|
||
|
||
directed=True produces a DiGraph that preserves edge direction (source→target).
|
||
directed=False (default) produces an undirected Graph for backward compatibility.
|
||
dedup=True (default) runs entity deduplication before building the graph.
|
||
dedup_llm_backend: if set (e.g. "gemini", "claude", or "kimi"), uses LLM to resolve
|
||
ambiguous pairs in the 75–92 Jaro-Winkler score zone.
|
||
|
||
Extractions are merged in order. For nodes with the same ID, the last
|
||
extraction's attributes win (NetworkX add_node overwrites). Pass AST
|
||
results before semantic results so semantic labels take precedence, or
|
||
reverse the order if you prefer AST source_location precision to win.
|
||
"""
|
||
from graphify.dedup import deduplicate_entities
|
||
combined: dict = {"nodes": [], "edges": [], "hyperedges": [], "input_tokens": 0, "output_tokens": 0}
|
||
for ext in extractions:
|
||
combined["nodes"].extend(ext.get("nodes", []))
|
||
combined["edges"].extend(ext.get("edges", []))
|
||
combined["hyperedges"].extend(ext.get("hyperedges", []))
|
||
combined["input_tokens"] += ext.get("input_tokens", 0)
|
||
combined["output_tokens"] += ext.get("output_tokens", 0)
|
||
if dedup and combined["nodes"]:
|
||
combined["nodes"], combined["edges"] = deduplicate_entities(
|
||
combined["nodes"], combined["edges"], communities={},
|
||
dedup_llm_backend=dedup_llm_backend,
|
||
)
|
||
return build_from_json(combined, directed=directed)
|
||
|
||
|
||
def _norm_label(label: str) -> str:
|
||
"""Canonical dedup key — lowercase, alphanumeric only."""
|
||
return re.sub(r"[^a-z0-9 ]", "", label.lower()).strip()
|
||
|
||
|
||
def deduplicate_by_label(nodes: list[dict], edges: list[dict]) -> tuple[list[dict], list[dict]]:
|
||
"""Merge nodes that share a normalised label, rewriting edge references.
|
||
|
||
Prefers IDs without chunk suffixes (_c\\d+) and shorter IDs when tied.
|
||
Drops self-loops created by the merge. Called in build() automatically.
|
||
"""
|
||
_CHUNK_SUFFIX = re.compile(r"_c\d+$")
|
||
canonical: dict[str, dict] = {} # norm_label -> surviving node
|
||
remap: dict[str, str] = {} # old_id -> surviving_id
|
||
|
||
for node in nodes:
|
||
key = _norm_label(node.get("label", node.get("id", "")))
|
||
if not key:
|
||
continue
|
||
existing = canonical.get(key)
|
||
if existing is None:
|
||
canonical[key] = node
|
||
else:
|
||
has_suffix = bool(_CHUNK_SUFFIX.search(node["id"]))
|
||
existing_has_suffix = bool(_CHUNK_SUFFIX.search(existing["id"]))
|
||
if has_suffix and not existing_has_suffix:
|
||
remap[node["id"]] = existing["id"]
|
||
elif existing_has_suffix and not has_suffix:
|
||
remap[existing["id"]] = node["id"]
|
||
canonical[key] = node
|
||
elif len(node["id"]) < len(existing["id"]):
|
||
remap[existing["id"]] = node["id"]
|
||
canonical[key] = node
|
||
else:
|
||
remap[node["id"]] = existing["id"]
|
||
|
||
if not remap:
|
||
return nodes, edges
|
||
|
||
print(f"[graphify] Deduplicated {len(remap)} duplicate node(s) by label.", file=sys.stderr)
|
||
deduped_nodes = list(canonical.values())
|
||
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 build_merge(
|
||
new_chunks: list[dict],
|
||
graph_path: str | Path = "graphify-out/graph.json",
|
||
prune_sources: list[str] | None = None,
|
||
*,
|
||
directed: bool = False,
|
||
dedup: bool = True,
|
||
dedup_llm_backend: str | None = None,
|
||
) -> nx.Graph:
|
||
"""Load existing graph.json, merge new chunks into it, and save back.
|
||
|
||
Never replaces - only grows (or prunes deleted-file nodes via prune_sources).
|
||
Safe to call repeatedly: existing nodes and edges are preserved.
|
||
"""
|
||
graph_path = Path(graph_path)
|
||
if graph_path.exists():
|
||
# Read JSON directly instead of going through node_link_graph().
|
||
# The latter rebuilds an undirected nx.Graph and then enumerating
|
||
# edges() yields endpoints based on node insertion order, which
|
||
# silently flips directional edges (e.g. `calls`) when the callee
|
||
# was inserted before the caller. The _src/_tgt direction-preserving
|
||
# attrs are popped before saving in export.py, so going through the
|
||
# NetworkX round-trip loses direction permanently (#760).
|
||
data = json.loads(graph_path.read_text(encoding="utf-8"))
|
||
links_key = "links" if "links" in data else "edges"
|
||
existing_nodes = list(data.get("nodes", []))
|
||
existing_edges = list(data.get(links_key, []))
|
||
base = [{"nodes": existing_nodes, "edges": existing_edges}]
|
||
else:
|
||
existing_nodes = []
|
||
base = []
|
||
|
||
all_chunks = base + list(new_chunks)
|
||
G = build(all_chunks, directed=directed, dedup=dedup, dedup_llm_backend=dedup_llm_backend)
|
||
|
||
# Prune nodes from deleted source files
|
||
if prune_sources:
|
||
to_remove = [
|
||
n for n, d in G.nodes(data=True)
|
||
if d.get("source_file") in prune_sources
|
||
]
|
||
G.remove_nodes_from(to_remove)
|
||
n_files = len(prune_sources)
|
||
n_nodes = len(to_remove)
|
||
if n_nodes:
|
||
print(
|
||
f"[graphify] Pruned {n_nodes} node(s) from {n_files} deleted source file(s).",
|
||
file=sys.stderr,
|
||
)
|
||
else:
|
||
print(
|
||
f"[graphify] {n_files} source file(s) deleted since last run — "
|
||
f"no matching nodes in graph, already clean.",
|
||
file=sys.stderr,
|
||
)
|
||
|
||
# Safety check: refuse to shrink the graph silently (#479)
|
||
# Skip when dedup or prune_sources is active — shrinkage is intentional there.
|
||
if graph_path.exists() and not dedup and not prune_sources:
|
||
existing_n = len(existing_nodes)
|
||
new_n = G.number_of_nodes()
|
||
if new_n < existing_n:
|
||
raise ValueError(
|
||
f"graphify: build_merge would shrink graph from {existing_n} → {new_n} nodes. "
|
||
f"Pass prune_sources explicitly if you intend to remove nodes."
|
||
)
|
||
|
||
return G
|
||
|
||
|
||
def prefix_graph_for_global(G: nx.Graph, repo_tag: str) -> nx.Graph:
|
||
"""Return a copy of G with all node IDs prefixed with repo_tag::.
|
||
|
||
Labels are preserved unchanged (for display). A 'local_id' attribute
|
||
is added to each node so the original ID can be recovered. Edges are
|
||
rewritten to match the new prefixed IDs. The 'repo' attribute is set
|
||
on every node.
|
||
"""
|
||
relabel = {n: f"{repo_tag}::{n}" for n in G.nodes}
|
||
H = nx.relabel_nodes(G, relabel, copy=True)
|
||
for node, data in H.nodes(data=True):
|
||
data["repo"] = repo_tag
|
||
data.setdefault("local_id", node.split("::", 1)[1])
|
||
return H
|
||
|
||
|
||
def prune_repo_from_graph(G: nx.Graph, repo_tag: str) -> int:
|
||
"""Remove all nodes tagged with repo_tag from G in-place. Returns count removed."""
|
||
to_remove = [n for n, d in G.nodes(data=True) if d.get("repo") == repo_tag]
|
||
G.remove_nodes_from(to_remove)
|
||
return len(to_remove)
|