"""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(), )