code/self_healing_finetuning/tests/test_paired_eval_diagnostics.py

from __future__ import annotations

import json
import tempfile
import unittest
from pathlib import Path
from unittest import mock

from data_pipeline import paired_eval_diagnostics
from data_pipeline.paired_eval_diagnostics import build_paired_eval_diagnostics


def write_jsonl(path: Path, rows: list[dict[str, object]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text("\n".join(json.dumps(row) for row in rows) + "\n", encoding="utf-8")


class PairedEvalDiagnosticsTests(unittest.TestCase):
    def test_builds_bucketed_diagnostics_and_repair_targets(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            workspace = Path(tmp)
            run_id = "run_test_paired_diagnostics"
            eval_dir = workspace / "runs" / run_id / "eval"
            predictions = eval_dir / "paired_predictions.jsonl"
            write_jsonl(
                predictions,
                [
                    {
                        "id": "margin_loss",
                        "task": "finance_qa",
                        "prompt": "Gross margin rose from 5% to 9%, but operating margin fell as opex rose.",
                        "baseline_answer": "<think>hidden</think> The margin picture is mixed: gross margin rose from 5% to 9%, but operating margin fell as opex rose, an offsetting cost pressure on profit.",
                        "candidate_answer": "<think>hidden</think> The company is clearly improving because gross margin rose.",
                        "baseline_score": {"score": 1.0, "critical_pass": True},
                        "candidate_score": {
                            "score": 0.5,
                            "critical_pass": False,
                            "checks": {
                                "neutral_language": False,
                                "risk_or_tradeoff_identified": False,
                                "numeric_reasoning_present": True,
                            },
                        },
                        "delta": -0.5,
                    },
                    {
                        "id": "tie",
                        "task": "finance_qa",
                        "prompt": "Summarize the known facts.",
                        "baseline_answer": "Known facts only.",
                        "candidate_answer": "Known facts only.",
                        "baseline_score": {"score": 1.0, "critical_pass": True},
                        "candidate_score": {"score": 1.0, "critical_pass": True, "checks": {}},
                        "delta": 0.0,
                    },
                ],
            )
            (eval_dir / "paired_eval_report.json").write_text(
                json.dumps({"improvement": {"losses": 1, "pairwise_loss_rate": 0.5}}),
                encoding="utf-8",
            )

            with mock.patch.object(paired_eval_diagnostics, "SHFT_WORKSPACE_ROOT", workspace):
                result = build_paired_eval_diagnostics(run_id=run_id, asset_class="equity", role="researcher")

            self.assertTrue(result["ok"], result)
            self.assertEqual(result["summary"]["prediction_count"], 2)
            self.assertEqual(result["summary"]["source_pairwise_loss_count"], 1)
            self.assertEqual(result["summary"]["pairwise_loss_count"], 1)
            self.assertEqual(result["summary"]["diagnosed_pairwise_loss_count"], 1)
            self.assertEqual(result["summary"]["latest_loss_coverage_ratio"], 1.0)
            self.assertTrue(result["summary"]["latest_loss_coverage_ok"])
            self.assertEqual(result["summary"]["critical_failure_count"], 1)
            self.assertEqual(result["summary"]["accepted_repair_target_count"], 1)

            diagnostics = [
                json.loads(line)
                for line in Path(result["diagnostics_jsonl_path"]).read_text(encoding="utf-8").splitlines()
            ]
            self.assertEqual(diagnostics[0]["prompt_id"], "margin_loss")
            self.assertEqual(diagnostics[0]["winner"], "baseline")
            self.assertEqual(diagnostics[0]["failure_bucket"], "margin_analysis")
            self.assertTrue(diagnostics[0]["judge_rationale"])
            self.assertTrue(diagnostics[0]["repair_target"]["admitted_to_training"])
            self.assertNotIn("<think>", diagnostics[0]["repair_target"]["answer"])
            self.assertTrue(Path(result["repair_targets_jsonl_path"]).exists())
            self.assertTrue(Path(result["markdown_path"]).exists())


if __name__ == "__main__":
    unittest.main()