Spaces:
Running
Running
File size: 4,605 Bytes
2b4bd40 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | 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()
|