tambo 0.3.4

"""Ontology-guided, full-suite Text2KGBench evaluation primitives."""

from __future__ import annotations

import hashlib
import json
import re
from collections import defaultdict
from pathlib import Path
from typing import Iterable


PIPELINE = "caboodle-text2kg-v3-ontology-guided"
SMOKE_SEED = "aegis-tydvlg.4-l1-smoke-v1"
PROMPT_TEMPLATE = "ontology-guided-json-evidence-v1"
PARSERS = ("strict-v1", "fenced-json-v2")


def canonical_json(value: object) -> str:
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True)


def sha256_bytes(value: bytes) -> str:
    return hashlib.sha256(value).hexdigest()


def sha256_file(path: Path) -> str:
    return sha256_bytes(path.read_bytes())


def read_jsonl(path: Path) -> list[dict]:
    return [json.loads(line) for line in path.read_text().splitlines() if line.strip()]


def normalize(value: object) -> str:
    return re.sub(r"[_\s]+", "", str(value)).lower()


def normalized_triple(triple: Iterable[object]) -> tuple[str, str, str]:
    return tuple(normalize(value) for value in triple)  # type: ignore[return-value]


def metric(tp: int, fp: int, fn: int) -> dict:
    precision = tp / (tp + fp) if tp + fp else 0.0
    recall = tp / (tp + fn) if tp + fn else 0.0
    f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0
    return {"tp": tp, "fp": fp, "fn": fn, "precision": precision, "recall": recall, "f1": f1}


def compile_ontology(ontology: dict, sentence: str) -> dict:
    """Compile ontology data into one canonical, domain-independent request."""
    concepts = sorted(
        ({"id": item["qid"], "label": item["label"]} for item in ontology["concepts"]),
        key=lambda item: (item["label"], item["id"]),
    )
    grouped: dict[str, set[tuple[str, str]]] = defaultdict(set)
    for item in ontology["relations"]:
        grouped[item["label"]].add((item["domain"], item["range"]))
    relations = [{"relation": label,
                  "signatures": [{"domain": domain, "range": range_}
                                 for domain, range_ in sorted(signatures)]}
                 for label, signatures in sorted(grouped.items())]
    schema = {
        "type": "object",
        "required": ["triples"],
        "properties": {"triples": {"type": "array", "items": {
            "type": "object", "required": ["subject", "relation", "object", "evidence_span"],
            "properties": {
                "subject": {"type": "string"},
                "relation": {"type": "string", "enum": [item["relation"] for item in relations]},
                "object": {"type": "string"},
                "evidence_span": {"type": "string"},
            },
        }}},
    }
    return {
        "instruction": "Extract only facts stated or unambiguously entailed by the isolated sentence.",
        "concepts": concepts,
        "relations": relations,
        "output_schema": schema,
        "sentence": sentence,
    }


def request_hash(request: dict, model: str, decoding: dict) -> str:
    return sha256_bytes(canonical_json({"request": request, "model": model, "decoding": decoding}).encode())


def smoke_ids(case_ids: Iterable[str], count: int = 5) -> list[str]:
    ranked = sorted(case_ids, key=lambda case_id: (
        hashlib.sha256(f"{SMOKE_SEED}\0{case_id}".encode()).hexdigest(), case_id))
    return ranked[:count]


def prefix_take(case_ids, remaining: int | None,
                selected: set[str] | None = None) -> list[str]:
    """The ids ONE group contributes to a `--max-cases` prefix, given `remaining`.

    `--max-cases` is a contiguous prefix of the run order: ontologies sorted by
    `(corpus, id)`, rows in file order. This is the per-group step of that walk,
    so the runner keeps its outer break and never opens an ontology past the
    budget — precomputing the whole prefix would read all 29 test files to run 5
    cases, which is the cost `--max-cases` exists to avoid.

    `remaining` is None for "no budget" (take everything this group offers). A
    `remaining` of 0 or less returns nothing: the runner subtracts what earlier
    ontologies consumed, so a later group can legitimately be handed 0. There is
    deliberately NO early-return fast path for that — the loop's own budget check
    already covers it, and a mutation proved the fast path was unreachable. Dead
    code in a guard reads as coverage it is not providing.

    WHY THIS EXISTS (aegis-687e2d). `--max-cases` shipped with no test, because
    reaching it through the CLI means standing up the 29-ontology universe and a
    dataset at a pinned git commit — `validate_manifest` checks both. The
    selection rule is arithmetic over ids and needs none of that, so it is
    extracted here where a plain list can exercise it. The runner CALLS this
    function rather than keeping a parallel copy: a helper the production path
    does not use is a test of nothing.

    TWO ORDERING PROPERTIES THAT ARE THE WHOLE CONTRACT, both easy to break
    without noticing:

      - the budget is counted AFTER the `selected` filter, so `--max-cases` with
        `--smoke` means "the first N of the PANEL", not "the first N rows, some
        of which the panel drops";
      - it is counted BEFORE the caller's already-done check, so re-running the
        same `--max-cases N` re-considers the same N cases rather than advancing
        N further into the corpus (aegis-xid7v6). A prefix that moves is not a
        prefix, and comparing two runs of it compares different work.
    """
    taken: list[str] = []
    for case_id in case_ids:
        if selected is not None and case_id not in selected:
            continue
        if remaining is not None and len(taken) >= remaining:
            break
        taken.append(case_id)
    return taken


def stratified_ids(groups, total: int) -> list[str]:
    """Select `total` case ids spread across groups, proportional to group size.

    `groups` is a sequence of `(group_id, case_ids)`. The result is deterministic
    and order-independent: within a group the ids are chosen by the same seeded
    ranking `smoke_ids` uses, and groups are ranked the same way, so the same
    `(groups, total)` always yields the same set however the input is ordered.

    WHY THIS EXISTS (aegis-mgw4x1). The runner's `--max-cases` takes a contiguous
    PREFIX, which is reproducible and re-runnable and is the right shape for "run
    a bounded block twice and compare". It is the wrong shape for estimating
    anything about the corpus: a 1,000-case prefix of Text2KGBench covered 10 of
    29 ontologies and ONE corpus, at 1,388 billed input tokens per case against
    the stratified panel's 3,117 — a 2.2x difference. Scaling either to the full
    universe is unsound, and the two projections agreeing ($27.86 vs $25.89) was
    two biased estimators coinciding rather than corroboration.

    EVERY NON-EMPTY GROUP GETS AT LEAST ONE CASE. That floor is the point of
    stratifying rather than sampling uniformly: without it a small ontology can
    round to zero and vanish from a result that claims to cover the corpus, which
    is the prefix's failure in a subtler form.
    """
    ranked_groups = sorted(
        ((gid, list(ids)) for gid, ids in groups if ids),
        key=lambda item: (hashlib.sha256(
            f"{SMOKE_SEED}\0{item[0]}".encode()).hexdigest(), item[0]),
    )
    if not ranked_groups or total <= 0:
        return []

    sizes = {gid: len(ids) for gid, ids in ranked_groups}
    available = sum(sizes.values())
    if total >= available:
        return sorted(cid for _, ids in ranked_groups for cid in ids)

    # Floor of one per group, for as many groups as the budget allows.
    if total <= len(ranked_groups):
        chosen = ranked_groups[:total]
        return sorted(smoke_ids(ids, 1)[0] for _, ids in chosen)

    allocation = {gid: 1 for gid, _ in ranked_groups}
    remaining = total - len(ranked_groups)

    # Largest-remainder over the proportional share of what is left, so the
    # totals add up exactly rather than drifting with rounding.
    surplus = {gid: sizes[gid] - 1 for gid, _ in ranked_groups}
    pool = sum(surplus.values())
    if pool:
        exact = {gid: remaining * surplus[gid] / pool for gid in surplus}
        for gid in surplus:
            allocation[gid] += min(int(exact[gid]), surplus[gid])
        short = total - sum(allocation.values())
        # Deterministic tiebreak: largest fractional part, then the seeded order.
        order = sorted(
            surplus,
            key=lambda gid: (-(exact[gid] - int(exact[gid])),
                             [g for g, _ in ranked_groups].index(gid)),
        )
        for gid in order:
            if short <= 0:
                break
            if allocation[gid] < sizes[gid]:
                allocation[gid] += 1
                short -= 1

    selected = []
    for gid, ids in ranked_groups:
        selected.extend(smoke_ids(ids, allocation[gid]))
    return sorted(selected)


def validate_candidate(candidate: object, sentence: str, ontology: dict) -> tuple[str, dict | None]:
    if not isinstance(candidate, dict) or any(k not in candidate for k in
                                               ("subject", "relation", "object", "evidence_span")):
        return "invalid_triple_shape", None
    values = {key: str(candidate[key]) for key in ("subject", "relation", "object", "evidence_span")}
    relation = next((item for item in ontology["relations"] if item["label"] == values["relation"]), None)
    if relation is None:
        return "relation_outside_ontology", None
    if values["evidence_span"] not in sentence:
        return "evidence_not_grounded", None
    if normalize(values["subject"]) not in normalize(sentence):
        return "subject_not_grounded", None
    if normalize(values["object"]) not in normalize(sentence):
        return "object_not_grounded", None
    return "accepted", values


def parse_response(raw: str, *, parser: str = "strict-v1") -> tuple[dict | None, bool]:
    if parser not in PARSERS:
        raise ValueError(f"unknown response parser: {parser}")
    if parser == "fenced-json-v2":
        return _parse_fenced_json(raw)
    # Preserve the published parser byte-for-byte for reproducible comparisons.
    try:
        return json.loads(raw), False
    except json.JSONDecodeError:
        match = re.fullmatch(r"\s*```(?:json)?\s*(\{.*\})\s*```\s*", raw, re.DOTALL)
        if not match:
            return None, False
        repaired = re.sub(r",\s*([}\]])", r"\1", match.group(1))
        try:
            return json.loads(repaired), True
        except json.JSONDecodeError:
            return None, False


def _remove_trailing_commas(text: str) -> str:
    """Repair punctuation outside strings; a regex corrupts literal ',}' values."""
    result = []
    in_string = escaped = False
    for index, char in enumerate(text):
        if in_string:
            result.append(char)
            if escaped:
                escaped = False
            elif char == "\\":
                escaped = True
            elif char == '"':
                in_string = False
        elif char == '"':
            in_string = True
            result.append(char)
        elif char == ',' and text[index + 1:].lstrip().startswith(('}', ']')):
            continue
        else:
            result.append(char)
    return ''.join(result)


def _parse_fenced_json(raw: str) -> tuple[dict | None, bool]:
    """Read one explicit JSON block, not arbitrary braces in model commentary.

    Multiple blocks are ambiguous and refused, even if one would score better.
    This parser receives no ontology, sentence, gold, or case identity. It repairs
    the response envelope only; the existing evidence validator still decides
    which triples are grounded.
    """
    def shape(value):
        return value if isinstance(value, dict) and isinstance(value.get('triples'), list) else None

    try:
        return shape(json.loads(raw)), False
    except json.JSONDecodeError:
        pass
    fences = list(re.finditer(r"(?m)^[ \t]*```[^\n]*$", raw))
    if len(fences) != 2:
        return None, False
    opening, closing = fences
    if not re.fullmatch(r"[ \t]*```(?:json)?[ \t]*", opening.group()):
        return None, False
    if not re.fullmatch(r"[ \t]*```[ \t]*", closing.group()):
        return None, False
    content = raw[opening.end():closing.start()]
    try:
        # Try valid JSON first: repair must not modify string values.
        payload = json.loads(content)
    except json.JSONDecodeError:
        try:
            payload = json.loads(_remove_trailing_commas(content))
        except json.JSONDecodeError:
            return None, False
    payload = shape(payload)
    return (payload, True) if payload is not None else (None, False)


def validate_manifest(dataset_root: Path, manifest: dict) -> dict:
    errors: list[str] = []
    entries = manifest.get("ontologies", [])
    if len(entries) != 29:
        errors.append(f"expected 29 ontologies, found {len(entries)}")
    seen_ids: set[str] = set()
    totals: dict[str, dict[str, int]] = defaultdict(lambda: {"ontologies": 0, "cases": 0, "gold": 0})
    declared_gold: set[str] = set()
    for entry in entries:
        totals[entry["corpus"]]["ontologies"] += 1
        files = entry["files"]
        for role, record in files.items():
            path = dataset_root / record["path"]
            if not path.is_file():
                errors.append(f"missing {role}: {record['path']}")
            elif sha256_file(path) != record["sha256"]:
                errors.append(f"hash mismatch {role}: {record['path']}")
        declared_gold.add(files["gold"]["path"])
        ontology_path = dataset_root / files["ontology"]["path"]
        if ontology_path.is_file():
            ontology = json.loads(ontology_path.read_text())
            compiled = compile_ontology(ontology, "CANARY SENTENCE")
            if len(compiled["relations"]) != len({item["label"] for item in ontology.get("relations", [])}):
                errors.append(f"ontology compiler relation loss: {entry['id']}")
            if compiled != compile_ontology(ontology, "CANARY SENTENCE"):
                errors.append(f"nondeterministic ontology compiler: {entry['id']}")
        if not errors or (dataset_root / files["gold"]["path"]).is_file():
            gold_rows = read_jsonl(dataset_root / files["gold"]["path"])
            test_rows = read_jsonl(dataset_root / files["test"]["path"])
            gold_ids, test_ids = [row["id"] for row in gold_rows], [row["id"] for row in test_rows]
            duplicate = (set(gold_ids) & seen_ids) | {x for x in gold_ids if gold_ids.count(x) > 1}
            if duplicate:
                errors.append(f"duplicate case ids in {entry['id']}: {sorted(duplicate)[:3]}")
            seen_ids.update(gold_ids)
            if gold_ids != test_ids:
                errors.append(f"test/gold id mismatch: {entry['id']}")
            cases, gold = len(gold_rows), sum(len(row["triples"]) for row in gold_rows)
            if cases != entry["case_count"] or gold != entry["gold_triple_count"]:
                errors.append(f"count mismatch: {entry['id']}")
            totals[entry["corpus"]]["cases"] += cases
            totals[entry["corpus"]]["gold"] += gold
    actual_gold = {str(path.relative_to(dataset_root)) for path in
                   (dataset_root / "data").glob("*/ground_truth/*_ground_truth.jsonl")}
    extra = sorted(actual_gold - declared_gold)
    if extra:
        errors.append(f"unmanifested ground truth: {extra}")
    expected = manifest["totals"]
    observed = {"ontologies": len(entries), "cases": len(seen_ids),
                "gold_triples": sum(value["gold"] for value in totals.values())}
    if observed != expected:
        errors.append(f"suite totals mismatch: {observed} != {expected}")
    if errors:
        raise ValueError("; ".join(errors))
    return {"pipeline": PIPELINE, "totals": observed, "corpora": dict(totals)}


def load_predictions(path: Path) -> dict[str, dict]:
    return {row["id"]: row for row in read_jsonl(path)}


def validate_reconcilers(registry: dict, ontology_ids: set[tuple[str, str]]) -> None:
    required = {"id", "version", "implementation_sha256", "scope", "pipeline", "mode",
                "input", "forbidden", "claim"}
    for item in registry.get("reconcilers", []):
        missing = required - set(item)
        if missing:
            raise ValueError(f"reconciler {item.get('id', '<unknown>')} missing {sorted(missing)}")
        if set(item) != required:
            raise ValueError(f"reconciler {item['id']} has unknown fields {sorted(set(item) - required)}")
        if "gold" in item["input"] or "gold" not in item["forbidden"]:
            raise ValueError(f"reconciler {item['id']} permits gold input")
        scope = {(corpus, ontology) for corpus in item["scope"].get("corpora", [])
                 for ontology in item["scope"].get("ontologies", [])}
        if not scope or not scope <= ontology_ids:
            raise ValueError(f"reconciler {item['id']} has invalid scope")
        if item["mode"] == "general" and scope != ontology_ids:
            raise ValueError(f"general reconciler {item['id']} must cover all ontologies")
        if item["mode"] not in {"general", "domain_specific"}:
            raise ValueError(f"reconciler {item['id']} has invalid mode")


def score_suite(dataset_root: Path, manifest: dict, response_root: Path, lever: str = "L0",
                *, parser: str = "strict-v1") -> dict:
    if parser not in PARSERS:
        raise ValueError(f"unknown response parser: {parser}")
    per_ontology, case_rows, stage_rows = [], [], []
    aggregate: dict[str, dict[str, list[int]]] = defaultdict(
        lambda: {"strict": [0, 0, 0], "relation_filtered": [0, 0, 0]})
    for entry in sorted(manifest["ontologies"], key=lambda x: (x["corpus"], x["id"])):
        ontology = json.loads((dataset_root / entry["files"]["ontology"]["path"]).read_text())
        gold_rows = read_jsonl(dataset_root / entry["files"]["gold"]["path"])
        response_file = (dataset_root / entry["files"]["upstream_response"]["path"] if lever == "L0"
                         else response_root / entry["corpus"] / f"{entry['id']}.jsonl")
        responses = load_predictions(response_file) if response_file.is_file() else {}
        counts = {"strict": [0, 0, 0], "relation_filtered": [0, 0, 0]}
        for row in gold_rows:
            case_id, gold = row["id"], {normalized_triple((t["sub"], t["rel"], t["obj"])) for t in row["triples"]}
            response = responses.get(case_id)
            accepted: set[tuple[str, str, str]] = set()
            stage = defaultdict(int)
            if response is None:
                stage["missing_output"] = 1
            else:
                payload, repaired = (({"triples": response["triples"]}, False)
                                     if lever == "L0" and isinstance(response.get("triples"), list)
                                     else parse_response(str(response.get("raw_response", response.get("response", ""))), parser=parser))
                if repaired:
                    stage["syntax_repaired"] += 1
                if payload is None:
                    stage["malformed_output"] += 1
                else:
                    candidates = payload.get("triples", []) if isinstance(payload, dict) else []
                    stage["generated"] = len(candidates)
                    for candidate in candidates:
                        if lever == "L0" and isinstance(candidate, list) and len(candidate) == 3:
                            verdict, value = "accepted", {"subject": candidate[0], "relation": candidate[1],
                                                          "object": candidate[2]}
                        else:
                            verdict, value = validate_candidate(candidate, row["sent"], ontology)
                        stage[verdict] += 1
                        if value:
                            normalized = normalized_triple((value["subject"], value["relation"], value["object"]))
                            if normalized in accepted:
                                stage["accepted"] -= 1
                                stage["duplicate"] += 1
                            else:
                                accepted.add(normalized)
                    rejected = sum(value for key, value in stage.items() if key not in
                                   {"generated", "accepted", "syntax_repaired"})
                    if stage["generated"] != len(accepted) + rejected:
                        raise RuntimeError(f"stage conservation failed for {case_id}")
            relations = {triple[1] for triple in gold}
            filtered = {triple for triple in accepted if triple[1] in relations}
            strict_values = (len(gold & accepted), len(accepted - gold), len(gold - accepted))
            filtered_values = (len(gold & filtered), len(filtered - gold), len(gold - filtered))
            for label, values in (("strict", strict_values), ("relation_filtered", filtered_values)):
                for index in range(3):
                    counts[label][index] += values[index]
            case_rows.append({"corpus": entry["corpus"], "ontology": entry["id"], "id": case_id,
                              "status": "ok" if response else "missing_output",
                              "strict": metric(*strict_values), "relation_filtered": metric(*filtered_values)})
            stage_rows.append({"corpus": entry["corpus"], "ontology": entry["id"], "id": case_id, **stage})
        for label in counts:
            for index in range(3):
                aggregate[entry["corpus"]][label][index] += counts[label][index]
        per_ontology.append({"corpus": entry["corpus"], "ontology": entry["id"],
                             **{label: metric(*values) for label, values in counts.items()}})
    corpus_rows = [{"corpus": corpus, **{label: metric(*values) for label, values in counts.items()}}
                   for corpus, counts in sorted(aggregate.items())]
    selected_counts = {"strict": [0, 0, 0], "relation_filtered": [0, 0, 0]}
    selected_ids: set[str] = set()
    for entry in manifest["ontologies"]:
        selected_file = entry["files"].get("selected_ids")
        if selected_file:
            selected_ids.update(line.strip() for line in
                                (dataset_root / selected_file["path"]).read_text().splitlines() if line.strip())
    for row in case_rows:
        if row["id"] in selected_ids:
            for label in selected_counts:
                for index, field in enumerate(("tp", "fp", "fn")):
                    selected_counts[label][index] += row[label][field]
    strata = [{"name": "wikidata_selected", "cases": len(selected_ids),
               **{label: metric(*values) for label, values in selected_counts.items()}}]

    unseen_counts = {"strict": [0, 0, 0], "relation_filtered": [0, 0, 0]}
    unseen_cases = 0
    for entry in manifest["ontologies"]:
        files = entry["files"]
        if "unseen_gold" not in files:
            continue
        ontology = json.loads((dataset_root / files["ontology"]["path"]).read_text())
        gold_rows = read_jsonl(dataset_root / files["unseen_gold"]["path"])
        response_file = (dataset_root / files["unseen_response"]["path"] if lever == "L0" else
                         response_root / entry["corpus"] / "unseen" / f"{entry['id']}.jsonl")
        responses = load_predictions(response_file) if response_file.is_file() else {}
        unseen_cases += len(gold_rows)
        for row in gold_rows:
            gold = {normalized_triple((t["sub"], t["rel"], t["obj"])) for t in row["triples"]}
            response = responses.get(row["id"], {})
            predicted: set[tuple[str, str, str]] = set()
            if lever == "L0":
                predicted = {normalized_triple(item) for item in response.get("triples", [])
                             if isinstance(item, list) and len(item) == 3}
            else:
                payload, _ = parse_response(str(response.get("raw_response", "")), parser=parser)
                for candidate in payload.get("triples", []) if isinstance(payload, dict) else []:
                    verdict, value = validate_candidate(candidate, row["sent"], ontology)
                    if verdict == "accepted" and value:
                        predicted.add(normalized_triple((value["subject"], value["relation"], value["object"])))
            filtered = {item for item in predicted if item[1] in {gold_item[1] for gold_item in gold}}
            for label, values in (("strict", predicted), ("relation_filtered", filtered)):
                counts = (len(gold & values), len(values - gold), len(gold - values))
                for index in range(3):
                    unseen_counts[label][index] += counts[index]
    strata.append({"name": "wikidata_unseen", "cases": unseen_cases,
                   **{label: metric(*values) for label, values in unseen_counts.items()}})
    return {"pipeline": PIPELINE if parser == "strict-v1" else PIPELINE + "+" + parser,
            "parser": parser, "lever": lever, "corpora": corpus_rows, "ontologies": per_ontology,
            "strata": strata, "cases": case_rows, "stages": stage_rows}