code/self_healing_finetuning/tests/test_repair_answer_quality.py

from __future__ import annotations

import unittest

from data_pipeline.repair_answer_quality import answer_admitted_to_training, repair_quality_checks


class RepairAnswerQualityTests(unittest.TestCase):
    def test_rejects_think_tags_and_chatty_cot(self) -> None:
        checks = repair_quality_checks(
            prompt="Debt increased from $7 million to $11 million while EBITDA stayed flat at $175 million.",
            answer="<think>Okay, let's see. The user wants me to calculate leverage.</think> It rose.",
            task="quantitative_qa",
            require_sections=False,
        )
        self.assertFalse(checks["no_think_tags"])
        self.assertFalse(checks["no_banned_cot_or_chatty_text"])
        self.assertFalse(answer_admitted_to_training(checks))

    def test_rejects_numeric_answer_without_terminal_numeric_conclusion(self) -> None:
        checks = repair_quality_checks(
            prompt="Debt increased from $7 million to $11 million while EBITDA stayed flat at $175 million.",
            answer=(
                "Reported facts: Debt increased while EBITDA stayed flat. "
                "Inference: leverage risk may be higher. "
                "Risk/tradeoff: balance-sheet flexibility could be lower. "
                "Conclusion: flag the change and request more evidence."
            ),
            task="quantitative_qa",
        )
        self.assertFalse(checks["retains_required_prompt_numbers"])
        self.assertFalse(checks["terminal_numeric_conclusion_when_required"])
        self.assertFalse(answer_admitted_to_training(checks))

    def test_accepts_clean_sectioned_leverage_answer(self) -> None:
        answer = (
            "Reported facts: Debt increased from $7 million to $11 million while EBITDA stayed flat at $175 million.\n"
            "Calculation: Debt/EBITDA moved from 7/175 = 0.040x to 11/175 = 0.063x, an increase of 0.023x.\n"
            "Inference: The reported fact is higher debt with unchanged EBITDA; the inference is that leverage risk increased.\n"
            "Risk/tradeoff: The risk is reduced balance-sheet flexibility if EBITDA weakens, while the offset is that absolute leverage remains low.\n"
            "Conclusion: Flag the leverage increase to 0.063x, but do not infer distress from this fact alone."
        )
        checks = repair_quality_checks(
            prompt="Debt increased from $7 million to $11 million while EBITDA stayed flat at $175 million.",
            answer=answer,
            task="quantitative_qa",
            require_sections=True,
        )
        self.assertTrue(answer_admitted_to_training(checks), checks)


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