code/self_healing_finetuning/data_pipeline/repair_coverage.py

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

from n21.config import write_json
from observability.audit_log import utc_now


DEFAULT_MINIMUMS = {
    "numeric_reasoning": 300,
    "fact_inference_separation": 300,
    "neutral_language": 200,
    "risk_tradeoff": 200,
    "critical_reasoning": 300,
}


def _read_jsonl(path: Path) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    with path.open("r", encoding="utf-8-sig") as handle:
        for line in handle:
            if not line.strip():
                continue
            row = json.loads(line)
            if isinstance(row, dict):
                rows.append(row)
    return rows


def _text(row: dict[str, Any]) -> str:
    parts: list[str] = []
    for message in row.get("messages") or []:
        if isinstance(message, dict):
            parts.append(str(message.get("content") or ""))
    return "\n".join(parts).lower()


def _metadata(row: dict[str, Any]) -> dict[str, Any]:
    meta = row.get("metadata")
    return meta if isinstance(meta, dict) else {}


def _has_any(text: str, terms: list[str]) -> bool:
    return any(term in text for term in terms)


def classify_repair_row(row: dict[str, Any]) -> set[str]:
    meta = _metadata(row)
    text = _text(row)
    classes: set[str] = set()

    failure_classes = meta.get("failure_classes") if isinstance(meta.get("failure_classes"), list) else []
    missing_checks = meta.get("missing_checks") if isinstance(meta.get("missing_checks"), list) else []
    synthetic_method = str(meta.get("synthetic_method") or "")
    task = str(meta.get("task") or "")
    rubric_target = str(meta.get("rubric_target") or "")
    repair_task = task in {"paired_eval_failure_repair", "grounded_critical_reasoning_sft"}
    explicit_repair = bool(failure_classes or missing_checks or repair_task or synthetic_method == "grounded_template_reasoning_v1")

    if "missing_numeric_reasoning" in failure_classes or "numeric_reasoning_present" in missing_checks:
        classes.add("numeric_reasoning")
    if "missing_fact_inference_separation" in failure_classes or "fact_inference_separation" in missing_checks:
        classes.add("fact_inference_separation")
    if "non_neutral_or_underqualified_language" in failure_classes or "neutral_language" in missing_checks:
        classes.add("neutral_language")
    if "missing_risk_or_tradeoff" in failure_classes or "risk_or_tradeoff_identified" in missing_checks:
        classes.add("risk_tradeoff")

    if repair_task or (explicit_repair and "critical" in rubric_target):
        classes.add("critical_reasoning")

    if not explicit_repair:
        return classes

    if _has_any(text, ["numeric anchor", "calculate", "calculation", "%", "basis point", "ratio", "compare"]):
        classes.add("numeric_reasoning")
    if _has_any(text, ["reported facts:", "facts:", "inference:", "observed facts", "separate fact"]):
        classes.add("fact_inference_separation")
    if _has_any(text, ["may", "could", "suggests", "uncertainty", "investigate", "rather than"]):
        classes.add("neutral_language")
    if _has_any(text, ["risk", "tradeoff", "red flag", "downside", "failure mode"]):
        classes.add("risk_tradeoff")

    if synthetic_method == "grounded_template_reasoning_v1":
        # These examples are intentionally generated only after a quality stall.
        # Count them across the core repair dimensions when the text carries the
        # required rubric language.
        classes.update({"critical_reasoning", "fact_inference_separation", "neutral_language", "risk_tradeoff"})
        if _has_any(text, ["numeric anchor", "quantitative", "valuation", "cash flow", "margin", "roic"]):
            classes.add("numeric_reasoning")

    return classes


def evaluate_repair_coverage(
    *,
    selected_training: Path,
    output_path: Path | None = None,
    minimums: dict[str, int] | None = None,
) -> dict[str, Any]:
    if not selected_training.exists():
        raise FileNotFoundError(f"selected training JSONL not found: {selected_training}")

    required = dict(DEFAULT_MINIMUMS)
    if minimums:
        required.update({key: int(value) for key, value in minimums.items() if value is not None})

    rows = _read_jsonl(selected_training)
    counts = {key: 0 for key in required}
    repair_rows = 0
    latest_failure_ids: set[str] = set()
    source_run_ids: set[str] = set()

    for row in rows:
        classes = classify_repair_row(row)
        if classes:
            repair_rows += 1
        for key in counts:
            if key in classes:
                counts[key] += 1
        meta = _metadata(row)
        if meta.get("source_prediction_id"):
            latest_failure_ids.add(str(meta["source_prediction_id"]))
        if meta.get("source_run_id"):
            source_run_ids.add(str(meta["source_run_id"]))

    checks = {
        key: {
            "ok": counts.get(key, 0) >= threshold,
            "detail": f"{counts.get(key, 0)} >= {threshold}",
        }
        for key, threshold in required.items()
    }
    errors = [f"{key}: {value['detail']}" for key, value in checks.items() if not value["ok"]]
    result = {
        "schema_version": "shft_repair_coverage_gate_v1",
        "selected_training": str(selected_training),
        "record_count": len(rows),
        "repair_row_count": repair_rows,
        "counts": counts,
        "minimums": required,
        "checks": checks,
        "latest_failure_id_count": len(latest_failure_ids),
        "source_run_ids": sorted(source_run_ids),
        "ok": not errors,
        "errors": errors,
        "created_at": utc_now(),
    }
    if output_path:
        write_json(output_path, result)
    return result