tests/test_agent_system.py

import os
import unittest
from unittest.mock import patch

from agent_system import (
    AgentSettings,
    LocalAgentSystem,
    ValidationDecision,
    clean_submission_value,
)
from model_config import DEFAULT_OLLAMA_MODEL


class AgentSettingsTests(unittest.TestCase):
    def test_gemma_4_is_the_shared_default(self) -> None:
        with patch.dict(os.environ, {}, clear=True):
            settings = AgentSettings.from_env()
        self.assertEqual(settings.text_model, DEFAULT_OLLAMA_MODEL)
        self.assertEqual(settings.multimodal_model, DEFAULT_OLLAMA_MODEL)
        self.assertEqual(settings.max_research_steps, 6)
        self.assertEqual(settings.max_validation_retries, 2)


class ValidationDecisionTests(unittest.TestCase):
    def test_pass_requires_consistent_supported_evidence(self) -> None:
        supported = ValidationDecision.from_payload(
            {
                "status": "pass",
                "answer": "42",
                "supporting_evidence": ["The report directly establishes 42."],
                "issues": [],
                "required_research": [],
                "rerun_plan": False,
            }
        )
        inconsistent = ValidationDecision.from_payload(
            {
                "status": "pass",
                "answer": "42",
                "supporting_evidence": ["One source says 42."],
                "issues": ["Another source says 43."],
                "required_research": [],
                "rerun_plan": False,
            }
        )
        self.assertTrue(supported.passed)
        self.assertFalse(inconsistent.passed)


class _FakeStructuredAgent:
    def __init__(self, responses):
        self.responses = iter(responses)

    def run(self, prompt, schema):
        return next(self.responses)


class _FakeResearcher:
    def __init__(self, responses):
        self.responses = iter(responses)
        self.calls = 0

    def run(self, prompt, reset):
        self.calls += 1
        return next(self.responses)


class RetryLoopTests(unittest.TestCase):
    def test_validation_feedback_triggers_research_retry(self) -> None:
        system = LocalAgentSystem.__new__(LocalAgentSystem)
        system.settings = AgentSettings(
            ollama_base_url="http://localhost:11434",
            text_model=DEFAULT_OLLAMA_MODEL,
            multimodal_model=DEFAULT_OLLAMA_MODEL,
            context_size=8192,
            max_research_steps=4,
            max_validation_retries=2,
        )
        plan = {
            "answer_format": "integer",
            "facts_to_verify": ["the exact count"],
            "research_queries": ["authoritative count"],
            "calculations": [],
            "attachment_use": "none",
        }
        system.planner = _FakeStructuredAgent([plan])
        system.researcher = _FakeResearcher(
            ["Conflicting evidence: 41 or 42", "Two sources establish 42"]
        )
        system.validator = _FakeStructuredAgent(
            [
                {
                    "status": "retry",
                    "answer": "",
                    "supporting_evidence": [],
                    "issues": ["The count is inconsistent."],
                    "required_research": ["Resolve 41 versus 42."],
                    "rerun_plan": False,
                },
                {
                    "status": "pass",
                    "answer": "42",
                    "supporting_evidence": ["Two sources establish 42."],
                    "issues": [],
                    "required_research": [],
                    "rerun_plan": False,
                },
            ]
        )

        answer = system.solve("task", "What is the count?", "No attachment.")

        self.assertEqual(answer, "42")
        self.assertEqual(system.researcher.calls, 2)


class CleanSubmissionValueTests(unittest.TestCase):
    def test_extracts_validator_marker(self) -> None:
        self.assertEqual(
            clean_submission_value("SUBMISSION_VALUE: 42"),
            "42",
        )

    def test_removes_thinking_and_quotes(self) -> None:
        raw = '<think>private reasoning</think>\nSUBMISSION_VALUE: "Ada Lovelace"'
        self.assertEqual(clean_submission_value(raw), "Ada Lovelace")

    def test_rejects_multiline_answer(self) -> None:
        with self.assertRaises(ValueError):
            clean_submission_value("first line\nsecond line")

    def test_rejects_forbidden_phrase(self) -> None:
        with self.assertRaises(ValueError):
            clean_submission_value("The final answer is 42")


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