src/graphtrace/graph.py

"""Provenance-linked knowledge graph: build from extraction records, persist as JSONL,
load into a NetworkX ``MultiDiGraph``.

One edge per (subject, relation, object, source chunk, quote span): the same fact
stated in two places gives two edges, each with its own evidence. ``KG.has_fact`` asks
whether a fact holds regardless of source.
"""

from __future__ import annotations

import hashlib
import json
import re
from collections import Counter, defaultdict
from pathlib import Path

import networkx as nx

from .canonical import Canonicalizer
from .prompts import THIS_PAPER
from .schemas import Document, Edge, Node, ResultRecord


def _hid(prefix: str, *parts: object) -> str:
    return prefix + hashlib.sha1("|".join(map(str, parts)).encode("utf-8")).hexdigest()[:12]


def paper_node_id(doc_id: str) -> str:
    return f"paper:{doc_id}"


def parse_number(value: str) -> float | None:
    m = re.search(r"-?\d+(?:\.\d+)?", value.replace(",", ""))
    return float(m.group(0)) if m else None


def _decimals(value: str) -> int:
    m = re.search(r"-?\d+(?:\.(\d+))?", value.replace(",", ""))
    return len(m.group(1)) if m and m.group(1) else 0


def values_agree(a: str, b: str) -> bool:
    """Same number up to scale (x, 100x, x/100: fraction vs percentage) after rounding both
    to the coarser reported precision. "91.2" == "0.912" == "91.20" == "91.23"; "91.2" != "90.4"."""
    na, nb = parse_number(a), parse_number(b)
    if na is None or nb is None:
        return a.strip().lower() == b.strip().lower()
    da, db = _decimals(a), _decimals(b)
    for x, dx in ((na, da), (na * 100, max(da - 2, 0)), (na / 100, da + 2)):
        p = min(dx, db)
        if round(x, p) == round(nb, p):
            return True
    return False


def value_groups(values: list[str]) -> list[list[str]]:
    """Greedy clustering of reported values into groups that agree under ``values_agree``."""
    groups: list[list[str]] = []
    for v in values:
        for g in groups:
            if values_agree(g[0], v):
                g.append(v)
                break
        else:
            groups.append([v])
    return groups


def build_graph(
    records: list[dict],
    docs: dict[str, Document],
    canon: Canonicalizer,
    confidence: dict[str, float] | None = None,
    stats: dict | None = None,
) -> tuple[list[Node], list[Edge], list[ResultRecord]]:
    """``confidence`` optionally maps edge_id -> NLI support score; ``stats`` receives counters."""
    surface: dict[str, Counter] = defaultdict(Counter)
    ntype: dict[str, str] = {}
    nodes: dict[str, Node] = {}
    for doc_id in sorted({r["doc_id"] for r in records}):
        pid = paper_node_id(doc_id)
        nodes[pid] = Node(node_id=pid, type="Paper", name=docs[doc_id].title, doc_id=doc_id)

    def ent(name: str, etype: str, doc_id: str) -> str:
        if etype == "Paper":
            return paper_node_id(doc_id)
        nid = canon.node_id(name, etype)
        surface[nid][name] += 1
        ntype[nid] = etype
        return nid

    # A result whose method is THIS_PAPER means the paper's own method: resolve it only when
    # the paper proposes exactly one method; otherwise drop the record and count it.
    proposed_by_doc: dict[str, set[str]] = defaultdict(set)
    for rec in records:
        if rec.get("error"):
            continue
        for t in rec["triples"]:
            if t["relation"] == "proposes" and t["subject_type"] == "Paper":
                proposed_by_doc[rec["doc_id"]].add(canon.node_id(t["object"], "Method"))
    dropped_this_paper = 0

    edges: dict[str, Edge] = {}
    results: dict[str, ResultRecord] = {}
    for rec in records:
        if rec.get("error"):
            continue
        for t in rec["triples"]:
            s = ent(t["subject"], t["subject_type"], rec["doc_id"])
            o = ent(t["object"], t["object_type"], rec["doc_id"])
            if s == o:
                continue
            eid = _hid("e_", s, t["relation"], o, rec["chunk_id"], t["start_char"], t["end_char"])
            edges[eid] = Edge(
                edge_id=eid,
                subject=s,
                relation=t["relation"],
                object=o,
                source_chunk_id=rec["chunk_id"],
                quote=t["quote"],
                start_char=t["start_char"],
                end_char=t["end_char"],
                confidence=(confidence or {}).get(eid),
                extractor_model=rec["extractor_model"],
            )
        for r in rec["results"]:
            if r["method"].strip() == THIS_PAPER:
                own = proposed_by_doc.get(rec["doc_id"], set())
                if len(own) != 1:
                    dropped_this_paper += 1
                    continue
                (m,) = own
            else:
                m = ent(r["method"], "Method", rec["doc_id"])
            d = ent(r["dataset"], "Dataset", rec["doc_id"])
            k = ent(r["metric"], "Metric", rec["doc_id"])
            rid = _hid("r_", m, d, k, r["value"], rec["chunk_id"], r["start_char"])
            results[rid] = ResultRecord(
                result_id=rid,
                method=m,
                dataset=d,
                metric=k,
                value=r["value"],
                value_num=parse_number(r["value"]),
                source_chunk_id=rec["chunk_id"],
                quote=r["quote"],
                start_char=r["start_char"],
                end_char=r["end_char"],
                extractor_model=rec["extractor_model"],
            )
    if stats is not None:
        stats["results_dropped_this_paper"] = dropped_this_paper
    for nid, forms in surface.items():
        name = min(forms.items(), key=lambda kv: (-kv[1], kv[0]))[0]
        nodes[nid] = Node(node_id=nid, type=ntype[nid], name=name, aliases=sorted(forms), mentions=sum(forms.values()))
    # Aliases from the canonicaliser make the linker find abbreviations and long forms.
    for m in canon.merges:
        for nid, node in nodes.items():
            if node.type != "Paper" and nid.split(":", 1)[1] == m["canonical"] and m["alias"] not in node.aliases:
                node.aliases.append(m["alias"])
    return sorted(nodes.values(), key=lambda n: n.node_id), sorted(edges.values(), key=lambda e: e.edge_id), sorted(results.values(), key=lambda r: r.result_id)


def save_graph(out_dir: str | Path, nodes: list[Node], edges: list[Edge], results: list[ResultRecord], merges: list[dict]) -> None:
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    for name, rows in (("nodes", nodes), ("edges", edges), ("results", results)):
        with (out / f"{name}.jsonl").open("w", encoding="utf-8") as f:
            for r in rows:
                f.write(r.model_dump_json() + "\n")
    with (out / "alias_merges.jsonl").open("w", encoding="utf-8") as f:
        for m in merges:
            f.write(json.dumps(m, ensure_ascii=False) + "\n")


def _read_jsonl(p: Path) -> list[dict]:
    if not p.exists():
        return []
    return [json.loads(line) for line in p.read_text(encoding="utf-8").splitlines() if line.strip()]


class KG:
    def __init__(self, nodes: list[Node], edges: list[Edge], results: list[ResultRecord]):
        self.nodes = {n.node_id: n for n in nodes}
        self.edges = {e.edge_id: e for e in edges}
        self.results = {r.result_id: r for r in results}
        self.g = nx.MultiDiGraph()
        for n in nodes:
            self.g.add_node(n.node_id, type=n.type, name=n.name)
        for e in edges:
            self.g.add_edge(e.subject, e.object, key=e.edge_id, relation=e.relation)
        self._facts: dict[tuple[str, str, str], list[str]] = defaultdict(list)
        for e in edges:
            self._facts[(e.subject, e.relation, e.object)].append(e.edge_id)

    def filtered(self, min_confidence: float | None) -> KG:
        """The graph without edges whose NLI support is below ``min_confidence``.

        One filtered view is used for benchmark gold *and* for every method's retrieval and
        for path checks, so gold and methods always see the same edge set. Unscored edges
        (confidence None) are kept; low-support edges stay in edges.jsonl for audit.
        """
        if min_confidence is None:
            return self
        keep = [e for e in self.edges.values() if e.confidence is None or e.confidence >= min_confidence]
        return KG(list(self.nodes.values()), keep, list(self.results.values()))

    @classmethod
    def load(cls, graph_dir: str | Path, min_confidence: float | None = None) -> KG:
        d = Path(graph_dir)
        return cls(
            [Node(**r) for r in _read_jsonl(d / "nodes.jsonl")],
            [Edge(**r) for r in _read_jsonl(d / "edges.jsonl")],
            [ResultRecord(**r) for r in _read_jsonl(d / "results.jsonl")],
        ).filtered(min_confidence)

    # ---- queries ----

    def name(self, node_id: str) -> str:
        n = self.nodes.get(node_id)
        return n.name if n else node_id

    def has_fact(self, s: str, r: str, o: str) -> bool:
        return (s, r, o) in self._facts

    def fact_edges(self, s: str, r: str, o: str) -> list[Edge]:
        return [self.edges[i] for i in self._facts.get((s, r, o), [])]

    def out_edges(self, node_id: str, relation: str | None = None) -> list[Edge]:
        if node_id not in self.g:
            return []
        return [self.edges[k] for _, _, k, d in self.g.out_edges(node_id, keys=True, data=True) if relation is None or d["relation"] == relation]

    def in_edges(self, node_id: str, relation: str | None = None) -> list[Edge]:
        if node_id not in self.g:
            return []
        return [self.edges[k] for _, _, k, d in self.g.in_edges(node_id, keys=True, data=True) if relation is None or d["relation"] == relation]

    def objects(self, s: str, relation: str) -> set[str]:
        return {e.object for e in self.out_edges(s, relation)}

    def subjects(self, o: str, relation: str) -> set[str]:
        return {e.subject for e in self.in_edges(o, relation)}

    def neighbourhood(self, seeds: list[str], hops: int) -> list[Edge]:
        """Edges reachable within ``hops`` of the seeds (ignoring direction), BFS order."""
        seen_nodes = {s for s in seeds if s in self.g}
        frontier = list(seen_nodes)
        out: list[Edge] = []
        seen_edges: set[str] = set()
        for _ in range(hops):
            nxt = []
            for n in frontier:
                for e in self.out_edges(n) + self.in_edges(n):
                    if e.edge_id in seen_edges:
                        continue
                    seen_edges.add(e.edge_id)
                    out.append(e)
                    other = e.object if e.subject == n else e.subject
                    if other not in seen_nodes:
                        seen_nodes.add(other)
                        nxt.append(other)
            frontier = nxt
        return out

    def results_for(self, method: str | None = None, dataset: str | None = None, metric: str | None = None) -> list[ResultRecord]:
        return [
            r
            for r in self.results.values()
            if (method is None or r.method == method) and (dataset is None or r.dataset == dataset) and (metric is None or r.metric == metric)
        ]

    def stats(self) -> dict:
        deg = Counter()
        for e in self.edges.values():
            deg[e.subject] += 1
            deg[e.object] += 1
        und = nx.Graph(self.g)
        comps = sorted((len(c) for c in nx.connected_components(und)), reverse=True) if len(und) else []
        return {
            "nodes": len(self.nodes),
            "nodes_by_type": dict(Counter(n.type for n in self.nodes.values())),
            "edges": len(self.edges),
            "distinct_facts": len(self._facts),
            "edges_by_relation": dict(Counter(e.relation for e in self.edges.values())),
            "results": len(self.results),
            "components": len(comps),
            "largest_component_nodes": comps[0] if comps else 0,
            "top_degree": [(self.name(n), d) for n, d in deg.most_common(10)],
            "edges_with_confidence": sum(1 for e in self.edges.values() if e.confidence is not None),
        }