src/graphtrace/methods.py

"""The four compared methods behind one interface.

============  ==================================================================
text_rag      top-ranked source passages (corpus-wide BM25 or dense)
graph_rag     facts around entities linked from the question, serialised as triples
hybrid        graph facts + the source passages those facts came from (+ BM25 fill)
graphtrace    hybrid context + required cited path + verifier and decision policy
============  ==================================================================

All methods use the same base model, the same answer prompt and output schema, and the
same context token budget measured on the exact text the model sees. Text RAG also has
to state its reasoning path as triples, so path validity is comparable across methods.

Handles: passages are ``[C01]..``, graph facts ``{E01}..``, result records ``{R01}..``.
They are resolved deterministically to chunk IDs / edge IDs / result IDs.
"""

from __future__ import annotations

import re
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Literal

from pydantic import BaseModel
from rank_bm25 import BM25Okapi

from .graph import KG, value_groups, values_agree
from .linking import EntityLinker
from .llm import LLMClient, LLMError, UsageMeter, call_structured
from .paths import check_path
from .prompts import ANSWER_SYSTEM, RETRY_SUFFIX, AnswerOut, answer_user
from .retrieval import tokenize
from .schemas import BenchItem, Chunk, Edge, MethodOutput, ResultRecord
from .support import score_edge_list
from .tokens import count_tokens
from .verification import NLIVerifier, strip_headings


class MethodConfig(BaseModel):
    name: str
    kind: Literal["text_rag", "graph_rag", "hybrid", "graphtrace"]
    context_budget: int = 2048
    max_hops: int = 2
    graph_share: float = 0.5
    min_edge_support: float | None = 0.5
    retry_on_invalid_path: bool = True
    override_confidence: float = 0.5


@dataclass
class Resources:
    base: LLMClient
    kg: KG
    linker: EntityLinker
    chunks: dict[str, Chunk]
    titles: dict[str, str]  # doc_id -> title
    nli: NLIVerifier | None = None
    _bm25: BM25Okapi | None = field(default=None, repr=False)
    _order: list[str] = field(default_factory=list, repr=False)

    def bm25_rank(self, question: str) -> list[str]:
        if self._bm25 is None:
            self._order = sorted(self.chunks)
            self._bm25 = BM25Okapi([tokenize(self.titles.get(self.chunks[c].doc_id, "") + " " + self.chunks[c].text) or ["_"] for c in self._order])
        scores = self._bm25.get_scores(tokenize(question))
        return [self._order[i] for i in sorted(range(len(self._order)), key=lambda i: (-scores[i], i))]


# ---- context building --------------------------------------------------------------


@dataclass
class Context:
    text: str = ""
    chunk_handles: dict[str, str] = field(default_factory=dict)  # C01 -> chunk_id
    fact_handles: dict[str, list[str]] = field(default_factory=dict)  # E01 -> edge ids
    result_handles: dict[str, str] = field(default_factory=dict)  # R01 -> result id
    seeds: list[str] = field(default_factory=list)

    @property
    def tokens(self) -> int:
        return count_tokens(self.text)


def render_chunk(handle: str, c: Chunk, res: Resources) -> str:
    return f"[{handle}] (paper: {res.titles.get(c.doc_id, c.doc_id)}; section: {c.section_heading})\n{strip_headings(c.text)}"


def add_passages(ctx: Context, chunk_ids: list[str], res: Resources, budget: int, parts: list[str]) -> int:
    used = 0
    for cid in chunk_ids:
        if cid in ctx.chunk_handles.values():
            continue
        h = f"C{len(ctx.chunk_handles) + 1:02d}"
        block = render_chunk(h, res.chunks[cid], res)
        t = count_tokens(block) + 2  # + separator
        if used + t > budget:
            continue
        ctx.chunk_handles[h] = cid
        parts.append(block)
        used += t
    return used


def graph_facts(question: str, res: Resources, hops: int) -> tuple[list[str], list[list[Edge]], list[ResultRecord]]:
    links = res.linker.link_question(question)
    seeds = [lk.node_id for lk in links]
    facts: dict[tuple[str, str, str], list[Edge]] = {}
    for e in res.kg.neighbourhood(seeds, hops):
        facts.setdefault((e.subject, e.relation, e.object), []).append(e)
    methods = [s for s in seeds if res.kg.nodes[s].type == "Method"]
    datasets = [s for s in seeds if res.kg.nodes[s].type == "Dataset"]
    results: list[ResultRecord] = []
    for m in methods or [None]:
        for d in datasets or [None]:
            if m is None and d is None:
                continue
            results.extend(res.kg.results_for(method=m, dataset=d))
    uniq = {r.result_id: r for r in results}
    return seeds, list(facts.values()), sorted(uniq.values(), key=lambda r: (r.method, r.dataset, r.metric, r.result_id))


def add_graph(ctx: Context, fact_groups: list[list[Edge]], results: list[ResultRecord], res: Resources, budget: int, parts: list[str]) -> int:
    used = 0
    for r in results:
        h = f"R{len(ctx.result_handles) + 1:02d}"
        line = f"{{{h}}} Result: {res.kg.name(r.method)} on {res.kg.name(r.dataset)}, {res.kg.name(r.metric)} = {r.value}"
        t = count_tokens(line) + 1
        if used + t > budget:
            continue
        ctx.result_handles[h] = r.result_id
        parts.append(line)
        used += t
    for edges in fact_groups:
        e = edges[0]
        h = f"E{len(ctx.fact_handles) + 1:02d}"
        line = f"{{{h}}} ({res.kg.name(e.subject)}) -[{e.relation}]-> ({res.kg.name(e.object)})"
        t = count_tokens(line) + 1
        if used + t > budget:
            continue
        ctx.fact_handles[h] = [x.edge_id for x in edges]
        parts.append(line)
        used += t
    return used


def build_context(cfg: MethodConfig, item: BenchItem, res: Resources) -> Context:
    ctx = Context()
    parts: list[str] = []
    if cfg.kind == "text_rag":
        add_passages(ctx, res.bm25_rank(item.question), res, cfg.context_budget, parts)
    else:
        seeds, groups, results = graph_facts(item.question, res, cfg.max_hops)
        ctx.seeds = seeds
        graph_budget = cfg.context_budget if cfg.kind == "graph_rag" else int(cfg.context_budget * cfg.graph_share)
        graph_parts: list[str] = []
        used = add_graph(ctx, groups, results, res, graph_budget - count_tokens("Graph facts:\n"), graph_parts)
        if graph_parts:
            parts.append("Graph facts:\n" + "\n".join(graph_parts))
        if cfg.kind in ("hybrid", "graphtrace"):
            # Source passages of the facts shown, most-cited first, then BM25 to fill the budget.
            counts: dict[str, int] = defaultdict(int)
            for eids in ctx.fact_handles.values():
                for eid in eids:
                    counts[res.kg.edges[eid].source_chunk_id] += 1
            for rid in ctx.result_handles.values():
                counts[res.kg.results[rid].source_chunk_id] += 1
            ordered = sorted(counts, key=lambda c: (-counts[c], c)) + res.bm25_rank(item.question)
            passage_parts: list[str] = []
            # Header and separators are counted against the budget too.
            add_passages(ctx, ordered, res, cfg.context_budget - used - count_tokens("Graph facts:\nSource passages:\n") - 8, passage_parts)
            if passage_parts:
                parts.append("Source passages:\n" + "\n\n".join(passage_parts))
    ctx.text = "\n\n".join(parts)
    return ctx


# ---- answering + verification ---------------------------------------------------------

_H = re.compile(r"([CER])\s*0*(\d+)")


def norm_handle(raw: str, kind: str) -> str | None:
    m = _H.search(raw.upper())
    return f"{m.group(1)}{int(m.group(2)):02d}" if m and m.group(1) == kind else None


def resolve(out: AnswerOut, ctx: Context, res: Resources) -> tuple[list[str], list[str], list[str], list[str]]:
    """-> (edge ids, chunk ids, result ids, invalid handles)."""
    edge_ids: list[str] = []
    result_ids: list[str] = []
    invalid: list[str] = []
    for raw in out.edge_ids:
        h = norm_handle(raw, "E") or norm_handle(raw, "R")
        if h and h in ctx.fact_handles:
            edge_ids.extend(ctx.fact_handles[h])
        elif h and h in ctx.result_handles:
            result_ids.append(ctx.result_handles[h])
        else:
            invalid.append(raw)
    chunk_ids: list[str] = []
    for raw in out.chunk_ids:
        h = norm_handle(raw, "C")
        if h and h in ctx.chunk_handles:
            chunk_ids.append(ctx.chunk_handles[h])
        else:
            invalid.append(raw)
    # Cited facts and results ground in the passages they were extracted from.
    sources = [res.kg.edges[e].source_chunk_id for e in edge_ids] + [res.kg.results[r].source_chunk_id for r in result_ids]
    return list(dict.fromkeys(edge_ids)), list(dict.fromkeys(chunk_ids + sources)), list(dict.fromkeys(result_ids)), invalid


def detect_conflict(ctx: Context, res: Resources) -> list[ResultRecord]:
    """Result records *shown in this method's budgeted context* for the linked
    (method, dataset[, metric]) whose values disagree across passages.

    Only records that were placed in the context are considered, so GraphTrace cannot use
    graph knowledge the model was not given. The agreement rule (``value_groups``) is the
    same one that defines C1 gold items; this overlap is a documented threat to validity
    and C1/C2 are reported separately from the headline.
    """
    kg = res.kg
    seeds = ctx.seeds
    methods = {s for s in seeds if kg.nodes[s].type == "Method"}
    datasets = {s for s in seeds if kg.nodes[s].type == "Dataset"}
    metrics = {s for s in seeds if kg.nodes[s].type == "Metric"}
    shown = [kg.results[rid] for rid in ctx.result_handles.values()]
    groups: dict[tuple[str, str, str], list[ResultRecord]] = defaultdict(list)
    for r in shown:
        if r.method in methods and r.dataset in datasets and (not metrics or r.metric in metrics) and r.value_num is not None:
            groups[(r.method, r.dataset, r.metric)].append(r)
    if not metrics and len({k[2] for k in groups}) > 1:
        return []  # several metrics and none named in the question: not a conflict
    for key in sorted(groups):
        rs = groups[key]
        if len(value_groups([r.value for r in rs])) >= 2 and len({r.source_chunk_id for r in rs}) >= 2:
            return rs
    return []


def answer_nodes(answer: list[str], res: Resources) -> set[str]:
    out = set()
    for a in answer:
        lk = res.linker.link_name(a)
        if lk:
            out.add(lk.node_id)
    return out


def edge_support(edge_ids: list[str], res: Resources, cfg: MethodConfig) -> dict[str, float]:
    """NLI support for cited edges that were not scored at graph-build time."""
    if res.nli is None:
        return {}
    todo = [res.kg.edges[e] for e in edge_ids if res.kg.edges[e].confidence is None]
    return score_edge_list(res.kg, todo, res.nli)


def run_method(cfg: MethodConfig, item: BenchItem, res: Resources) -> MethodOutput:
    meter = UsageMeter()
    trace: dict = {}
    try:
        ctx = build_context(cfg, item, res)
        trace["seeds"] = ctx.seeds
        trace["context"] = ctx.text
        trace["handles"] = {"chunks": ctx.chunk_handles, "facts": ctx.fact_handles, "results": ctx.result_handles}
        user = answer_user(item.question, ctx.text or "(empty)")
        out = call_structured(res.base, meter, "answer", ANSWER_SYSTEM, user, AnswerOut, max_tokens=2048)
        trace["raw_answer"] = out.model_dump()
        decision, answer, confidence = out.decision, list(out.answer), float(min(1.0, max(0.0, out.confidence)))
        path = [[t.subject, t.relation, t.object] for t in out.graph_path]
        edge_ids, chunk_ids, result_ids, invalid = resolve(out, ctx, res)
        trace["invalid_handles"] = invalid

        if cfg.kind == "graphtrace":
            report: dict = {"overrides": []}
            conflict = detect_conflict(ctx, res)
            if conflict:
                report["conflict_results"] = [r.result_id for r in conflict]
                if decision != "flag_conflict":
                    report["overrides"].append(f"{decision}->flag_conflict (conflicting result records)")
                    decision, confidence = "flag_conflict", cfg.override_confidence
                answer = sorted({r.value for r in conflict})
                chunk_ids = list(dict.fromkeys(chunk_ids + [r.source_chunk_id for r in conflict]))
            elif decision == "answer":
                for attempt in range(2):
                    support = edge_support(edge_ids, res, cfg)
                    low = {e for e, s in support.items() if cfg.min_edge_support is not None and s < cfg.min_edge_support}
                    # A value answer is backed only by a result record that was in the context and cited.
                    is_value = any(values_agree(res.kg.results[r].value, a) for r in result_ids for a in answer)
                    pr = check_path(res.kg, res.linker, path, set(ctx.seeds), set() if is_value else answer_nodes(answer, res), cfg.min_edge_support)
                    problems = pr.problems() + [f"cited fact {e} is not supported by its source" for e in sorted(low)]
                    if not path and not is_value:
                        problems.append("no reasoning path was given")
                    report[f"path_check_{attempt}"] = {**pr.to_dict(), "low_support_edges": sorted(low), "problems": problems}
                    if not problems:
                        break
                    if attempt == 0 and cfg.retry_on_invalid_path:
                        out = call_structured(
                            res.base,
                            meter,
                            "answer:retry",
                            ANSWER_SYSTEM,
                            user + RETRY_SUFFIX.format(problems="\n".join(f"- {p}" for p in problems)),
                            AnswerOut,
                            max_tokens=2048,
                        )
                        trace["raw_answer_retry"] = out.model_dump()
                        decision, answer, confidence = out.decision, list(out.answer), float(min(1.0, max(0.0, out.confidence)))
                        path = [[t.subject, t.relation, t.object] for t in out.graph_path]
                        edge_ids, chunk_ids, result_ids, invalid = resolve(out, ctx, res)
                        if decision != "answer":
                            break
                        continue
                    report["overrides"].append("answer->abstain (path failed verification)")
                    decision, answer, confidence = "abstain", [], cfg.override_confidence
                    break
            trace["verifier"] = report

        return MethodOutput(
            item_id=item.item_id,
            method=cfg.name,
            answer=answer if decision != "abstain" else [],
            confidence=confidence,
            decision=decision,
            graph_path=path,
            cited_edge_ids=edge_ids,
            source_chunk_ids=chunk_ids,
            context_tokens=ctx.tokens,
            trace=trace,
            usage=meter.summary(),
        )
    except Exception as e:  # noqa: BLE001 - one item must never abort a run; errors are scored as incorrect
        kind = "llm" if isinstance(e, LLMError) else "internal"
        return MethodOutput(
            item_id=item.item_id,
            method=cfg.name,
            answer=[],
            confidence=0.0,
            decision="abstain",
            error=f"{kind}: {type(e).__name__}: {e}",
            trace=trace,
            usage=meter.summary(),
        )