40 lines
1.8 KiB
Python
40 lines
1.8 KiB
Python
# assemble node+edge dicts into a NetworkX graph, preserving edge direction
|
|
from __future__ import annotations
|
|
import sys
|
|
import networkx as nx
|
|
from .validate import validate_extraction
|
|
|
|
|
|
def build_from_json(extraction: dict) -> nx.Graph:
|
|
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()
|
|
for node in extraction.get("nodes", []):
|
|
G.add_node(node["id"], **{k: v for k, v in node.items() if k != "id"})
|
|
node_set = set(G.nodes())
|
|
for edge in extraction.get("edges", []):
|
|
src, tgt = edge["source"], edge["target"]
|
|
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")}
|
|
# 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)
|
|
return G
|
|
|
|
|
|
def build(extractions: list[dict]) -> nx.Graph:
|
|
"""Merge multiple extraction results into one graph."""
|
|
combined: dict = {"nodes": [], "edges": [], "input_tokens": 0, "output_tokens": 0}
|
|
for ext in extractions:
|
|
combined["nodes"].extend(ext.get("nodes", []))
|
|
combined["edges"].extend(ext.get("edges", []))
|
|
combined["input_tokens"] += ext.get("input_tokens", 0)
|
|
combined["output_tokens"] += ext.get("output_tokens", 0)
|
|
return build_from_json(combined)
|