federate.py

"""Cross-graph federation.

Merge N graphs (e.g. one per-source vault, one per-codebase) into a
universe graph. Same-label entities across graphs are linked or merged.

Three matching tiers:

  1. **exact id**: same id in two graphs → merged (deterministic).
  2. **canonical label** (case/whitespace-folded) + same `type` → merged.
  3. **fuzzy label** (token overlap) + same `type` + LLM tie-break →
     linked via a `same_as` edge with confidence='INFERRED'.

The output preserves provenance: when two nodes are merged, the merged
node lists `_merged_from = [id1, id2, ...]` so users can trace back.
Edges are de-duped by `(source, target, relation)` after id remapping.

`bridge.py`'s paper↔code shim now forwards to this module — see the
back-compat shim added at the bottom.
"""
from __future__ import annotations

import json
from collections import defaultdict
from pathlib import Path
from typing import Any, Callable

from .normalize import canonicalise_label


def federate(
    graphs: list[dict[str, Any]],
    *,
    fuzzy: bool = False,
    fuzzy_threshold: float = 0.7,
    llm_tiebreak: Any | None = None,
) -> dict[str, Any]:
    """Combine `graphs` into one universe.

    Parameters
    ----------
    graphs
        List of `{nodes, edges}` dicts (each typically read from a
        ``graph.json`` produced by ``Session.run``).
    fuzzy
        If True, also propose `same_as` edges between nodes whose labels
        share ≥``fuzzy_threshold`` token overlap (Jaccard) and have the
        same type. Disabled by default — fuzzy matches are advisory and
        easy to over-fire on small graphs.
    llm_tiebreak
        Optional LLM client. When ``fuzzy=True`` and a candidate pair
        scores below 0.9, the LLM is asked to confirm. ``None`` means
        accept everything ≥ ``fuzzy_threshold`` without LLM.

    Returns
    -------
    {"nodes": [...], "edges": [...], "linking_report": {...}}
    """
    # ------------------------------------------------------------------
    # Tier 1 + 2: exact id / canonical label remap
    # ------------------------------------------------------------------
    by_canon: dict[tuple[str, str], list[tuple[int, dict]]] = defaultdict(list)
    for gi, g in enumerate(graphs):
        for n in g.get("nodes", []):
            ckey = (canonicalise_label(n.get("label", "")), str(n.get("type") or ""))
            by_canon[ckey].append((gi, n))

    id_remap: dict[tuple[int, str], str] = {}   # (graph_index, original_id) → canonical_id
    canonical_nodes: list[dict] = []
    merge_groups: list[dict] = []

    for ckey, hits in by_canon.items():
        if len(hits) == 1:
            gi, n = hits[0]
            id_remap[(gi, n["id"])] = n["id"]
            canonical_nodes.append(dict(n))
        else:
            # Pick the first node's id as canonical; record the merge.
            target_gi, target_node = hits[0]
            target_id = target_node["id"]
            merged_node = dict(target_node)
            merged_from: list[str] = []
            for gi, n in hits:
                id_remap[(gi, n["id"])] = target_id
                merged_from.append(n["id"])
                # Carry over distinct attrs that the canonical didn't have.
                for k, v in n.items():
                    if k not in merged_node and v not in (None, ""):
                        merged_node[k] = v
            merged_node["_merged_from"] = merged_from
            canonical_nodes.append(merged_node)
            merge_groups.append({
                "into": target_id, "from": merged_from,
                "reason": "canonical label + type match",
            })

    # Edge consolidation: id-remap + dedup by (src, tgt, rel).
    seen_edges: set[tuple[str, str, str]] = set()
    canonical_edges: list[dict] = []
    for gi, g in enumerate(graphs):
        for e in g.get("edges", []) or g.get("links", []):
            src = id_remap.get((gi, e.get("source")), e.get("source"))
            tgt = id_remap.get((gi, e.get("target")), e.get("target"))
            rel = e.get("relation") or ""
            key = (src, tgt, rel)
            if key in seen_edges:
                continue
            seen_edges.add(key)
            canonical_edges.append(dict(e, source=src, target=tgt))

    # ------------------------------------------------------------------
    # Tier 3: fuzzy same_as suggestions (read-only — never merges)
    # ------------------------------------------------------------------
    fuzzy_links: list[dict] = []
    if fuzzy:
        from itertools import combinations

        # Group by type so we only compare within type.
        nodes_by_type: dict[str, list[dict]] = defaultdict(list)
        for n in canonical_nodes:
            nodes_by_type[str(n.get("type") or "")].append(n)

        for type_, nodes in nodes_by_type.items():
            for a, b in combinations(nodes, 2):
                if a["id"] == b["id"]:
                    continue
                score = _jaccard_tokens(a.get("label", ""), b.get("label", ""))
                if score < fuzzy_threshold:
                    continue
                if score < 0.9 and llm_tiebreak is not None:
                    if not _llm_says_same(llm_tiebreak, a, b):
                        continue
                fuzzy_links.append({
                    "source": a["id"],
                    "target": b["id"],
                    "relation": "same_as",
                    "confidence": "INFERRED",
                    "score": round(score, 3),
                    "source_file": "<federation>",
                })

    canonical_edges.extend(fuzzy_links)

    return {
        "nodes": canonical_nodes,
        "edges": canonical_edges,
        "linking_report": {
            "n_input_graphs": len(graphs),
            "n_input_nodes": sum(len(g.get("nodes", [])) for g in graphs),
            "n_output_nodes": len(canonical_nodes),
            "n_output_edges": len(canonical_edges),
            "n_merge_groups": len(merge_groups),
            "merge_groups": merge_groups[:50],
            "n_fuzzy_links": len(fuzzy_links),
            "fuzzy_links": fuzzy_links[:50],
        },
    }


def federate_files(
    paths: list[str | Path],
    *,
    fuzzy: bool = False,
    fuzzy_threshold: float = 0.7,
    llm_tiebreak: Any | None = None,
    out_path: str | Path | None = None,
) -> dict[str, Any]:
    """Convenience: load graphs from disk, federate, optionally write `out_path`."""
    graphs = []
    for p in paths:
        data = json.loads(Path(p).read_text())
        # Accept both NetworkX node-link and flat extraction shapes.
        if "links" in data and "edges" not in data:
            data["edges"] = data.pop("links")
        graphs.append(data)
    out = federate(graphs, fuzzy=fuzzy, fuzzy_threshold=fuzzy_threshold,
                   llm_tiebreak=llm_tiebreak)
    if out_path:
        Path(out_path).parent.mkdir(parents=True, exist_ok=True)
        Path(out_path).write_text(json.dumps(out, indent=2, ensure_ascii=False))
    return out


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _tokens(text: str) -> set[str]:
    return {t for t in canonicalise_label(text).split() if t}


def _jaccard_tokens(a: str, b: str) -> float:
    ta = _tokens(a)
    tb = _tokens(b)
    if not ta or not tb:
        return 0.0
    return len(ta & tb) / len(ta | tb)


_TIE_PROMPT = """Are these two nodes referring to the SAME real-world entity?
Answer strictly 'yes' or 'no'.

A: {a_label}  ({a_type})
B: {b_label}  ({b_type})
"""


def _llm_says_same(llm: Any, a: dict, b: dict) -> bool:
    prompt = _TIE_PROMPT.format(
        a_label=a.get("label"), a_type=a.get("type"),
        b_label=b.get("label"), b_type=b.get("type"),
    )
    try:
        if hasattr(llm, "chat"):
            out = llm.chat(prompt)
            text = out.get("content") if isinstance(out, dict) else (out or "")
            return text.strip().lower().startswith("y")
    except Exception:
        return False
    return False