Spaces:
Sleeping
Sleeping
File size: 12,985 Bytes
116524e | 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 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 | """Tests for RRStep — PydanticAI-based Recursive Reflector."""
from __future__ import annotations
from unittest.mock import patch
import pytest
from pydantic_ai.settings import ModelSettings
from ace.implementations.rr.config import RecursiveConfig
from ace.core.context import ACEStepContext, SkillbookView
from ace.core.outputs import AgentOutput, ReflectorOutput
from ace.core.skillbook import Skillbook
from ace.steps.rr_step import RRStep, RRConfig
from ace.implementations.rr.tools import RRDeps
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_ctx(
question: str = "test",
answer: str = "a",
reasoning: str = "r",
ground_truth: str | None = None,
feedback: str | None = None,
) -> ACEStepContext:
"""Build an ACEStepContext suitable for RRStep.__call__."""
trace: dict = {
"question": question,
"steps": [
{"role": "agent", "reasoning": reasoning, "answer": answer, "skill_ids": []}
],
}
if ground_truth is not None:
trace["ground_truth"] = ground_truth
if feedback is not None:
trace["feedback"] = feedback
return ACEStepContext(trace=trace, skillbook=SkillbookView(Skillbook()))
_RUN_SYNC = "ace.core.recursive_agent.run_agent_sync"
def _mock_compaction_result(
*,
reasoning: str = "mock reasoning",
key_insight: str = "mock insight",
correct_approach: str = "mock approach",
) -> tuple[ReflectorOutput, dict]:
"""Create a mock return value for run_agent_sync."""
output = ReflectorOutput(
reasoning=reasoning,
error_identification="none",
root_cause_analysis="mock root cause",
correct_approach=correct_approach,
key_insight=key_insight,
raw={},
)
metadata = {
"usage": {
"input_tokens": 100,
"output_tokens": 50,
"total_tokens": 150,
"requests": 3,
},
"compactions": 0,
"depth": 0,
"iterations": 2,
"timed_out": False,
}
return output, metadata
# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
@pytest.mark.unit
class TestRRStep:
"""Test RRStep construction and StepProtocol."""
def test_step_protocol_attributes(self):
rr = RRStep("test-model", config=RRConfig())
assert "trace" in rr.requires
assert "skillbook" in rr.requires
assert "reflections" in rr.provides
assert "reflection" not in rr.provides
def test_call_produces_reflection_on_context(self):
"""RRStep.__call__ populates ctx.reflections."""
rr = RRStep("test-model", config=RRConfig())
reflection, metadata = _mock_compaction_result(key_insight="step test")
with patch(_RUN_SYNC, return_value=(reflection, metadata)):
ctx = _make_ctx(
question="What is 2+2?",
answer="4",
reasoning="2+2=4",
ground_truth="4",
feedback="Correct!",
)
result_ctx = rr(ctx)
assert len(result_ctx.reflections) == 1
assert isinstance(result_ctx.reflections[0], ReflectorOutput)
assert result_ctx.reflections[0].key_insight == "step test"
def test_rr_trace_metadata_populated(self):
"""Successful reflection populates rr_trace in raw."""
rr = RRStep("test-model", config=RRConfig())
reflection, metadata = _mock_compaction_result()
with patch(_RUN_SYNC, return_value=(reflection, metadata)):
result_ctx = rr(_make_ctx())
result = result_ctx.reflections[0]
assert "rr_trace" in result.raw
assert result.raw["rr_trace"]["timed_out"] is False
assert "usage" in result.raw
def test_thoughts_are_exposed_in_raw(self):
"""RRStep preserves think-tool notes recorded during evidence gathering."""
rr = RRStep("test-model", config=RRConfig())
reflection, metadata = _mock_compaction_result()
def _run_with_thought(*args, **kwargs):
deps = kwargs["deps"]
deps.thoughts.append(
{
"thought": "The selected flights satisfy the requested dates.",
"evidence_refs": ["messages[5]", "messages[9]"],
}
)
return reflection, metadata
with patch(_RUN_SYNC, side_effect=_run_with_thought):
result_ctx = rr(_make_ctx())
thoughts = result_ctx.reflections[0].raw["thoughts"]
assert thoughts == [
{
"thought": "The selected flights satisfy the requested dates.",
"evidence_refs": ["messages[5]", "messages[9]"],
}
]
def test_timeout_produces_output(self):
"""Budget exhaustion produces a timeout ReflectorOutput."""
from ace.core.recursive_agent import BudgetExhausted
rr = RRStep("test-model", config=RRConfig())
with patch(_RUN_SYNC, side_effect=BudgetExhausted(compaction_count=0)):
result_ctx = rr(_make_ctx())
assert len(result_ctx.reflections) == 1
output = result_ctx.reflections[0]
assert isinstance(output, ReflectorOutput)
assert "budget limit" in output.reasoning.lower()
assert output.raw.get("timeout") is True
def test_timeout_with_ground_truth_correct(self):
"""Timeout correctly detects correct answer."""
from ace.core.recursive_agent import BudgetExhausted
rr = RRStep("test-model", config=RRConfig())
with patch(_RUN_SYNC, side_effect=BudgetExhausted(compaction_count=0)):
output = rr.reflect(
question="What is 2+2?",
agent_output=AgentOutput(reasoning="r", final_answer="4"),
ground_truth="4",
)
assert isinstance(output, ReflectorOutput)
assert "correct" in output.reasoning.lower()
def test_error_produces_safe_output(self):
"""General exception produces a safe fallback output."""
rr = RRStep("test-model", config=RRConfig())
with patch(_RUN_SYNC, side_effect=RuntimeError("unexpected error")):
result_ctx = rr(_make_ctx())
assert len(result_ctx.reflections) == 1
output = result_ctx.reflections[0]
assert "failed" in output.reasoning.lower()
@pytest.mark.unit
class TestRRStepProtocol:
"""Test that RRStep satisfies structural protocols."""
def test_satisfies_reflector_like(self):
"""RRStep satisfies ReflectorLike protocol."""
from ace.protocols import ReflectorLike
rr = RRStep("test-model", config=RRConfig())
assert isinstance(rr, ReflectorLike)
def test_reflect_method(self):
"""reflect() delegates to the PydanticAI agent."""
rr = RRStep("test-model", config=RRConfig())
reflection, metadata = _mock_compaction_result(key_insight="reflected")
with patch(_RUN_SYNC, return_value=(reflection, metadata)):
result = rr.reflect(
question="What is 2+2?",
agent_output=AgentOutput(reasoning="r", final_answer="4"),
ground_truth="4",
feedback="Correct!",
)
assert isinstance(result, ReflectorOutput)
assert result.key_insight == "reflected"
@pytest.mark.unit
class TestMeteredModel:
"""``MeteredModel`` fires the usage callback from the pydantic-ai model layer."""
def test_callback_invoked_with_request_usage_and_model_name(self):
from pydantic_ai import Agent
from pydantic_ai.models.test import TestModel
from pydantic_ai.usage import RequestUsage
from ace.core.metered_model import MeteredModel
calls: list[tuple[RequestUsage, str]] = []
def _cb(usage, model_id):
calls.append((usage, model_id))
inner = TestModel()
agent = Agent(MeteredModel(inner, _cb), output_type=str)
result = agent.run_sync("hello")
assert result.output
assert len(calls) >= 1
reported_usage, model_id = calls[-1]
assert isinstance(reported_usage, RequestUsage)
assert reported_usage.input_tokens > 0
assert model_id == inner.model_name
def test_callback_exception_does_not_break_agent_run(self):
from pydantic_ai import Agent
from pydantic_ai.models.test import TestModel
from ace.core.metered_model import MeteredModel
def _cb(usage, model_id):
raise RuntimeError("boom")
agent = Agent(MeteredModel(TestModel(), _cb), output_type=str)
result = agent.run_sync("hello")
assert result.output
def test_rrstep_accepts_prebuilt_model_instance(self):
"""Passing a pre-built ``Model`` flows through ``RRStep`` unchanged."""
from pydantic_ai.models.test import TestModel
test_model = TestModel()
rr = RRStep(test_model, config=RRConfig())
assert rr._model is test_model
assert rr._agent.model is test_model
def test_rrstep_wraps_model_when_usage_callback_set(self):
"""``RRStep.__init__`` routes the agent model through ``MeteredModel``."""
from ace.core.metered_model import MeteredModel
rr = RRStep(
"test-model",
config=RRConfig(usage_callback=lambda u, n: None),
)
assert isinstance(rr._agent.model, MeteredModel)
def test_rrstep_does_not_wrap_when_no_callback(self):
"""Without a callback there's no wrapper overhead."""
from ace.core.metered_model import MeteredModel
rr = RRStep("test-model", config=RRConfig())
assert not isinstance(rr._agent.model, MeteredModel)
def test_rrstep_uses_prompted_reflector_output(self):
"""RR should gather evidence with tools and return structured output directly."""
rr = RRStep("test-model", config=RRConfig())
assert rr._agent._output_schema.mode == "prompted"
assert rr._agent._output_schema.allows_text is True
def test_rrstep_defaults_to_deterministic_temperature(self):
"""RR defaults to deterministic evidence analysis unless overridden."""
rr = RRStep("test-model", config=RRConfig())
assert rr._agent.model_settings["temperature"] == 0.0
def test_rrstep_preserves_explicit_model_settings(self):
"""Callers can still override RR model settings explicitly."""
rr = RRStep(
"test-model",
config=RRConfig(),
model_settings=ModelSettings(temperature=0.7),
)
assert rr._agent.model_settings["temperature"] == 0.7
def test_rrstep_specializes_execute_code_tool_description(self):
"""RR should present execute_code as an evidence tool, not a prose channel."""
rr = RRStep("test-model", config=RRConfig())
tool = rr._agent._function_toolset.tools["execute_code"]
assert "evidence workbench" in tool.description
assert "think" in tool.description
assert "store strings/snippets" in tool.description
assert tool.function_schema.description == tool.description
code_schema = tool.function_schema.json_schema["properties"]["code"]
assert "short snippet" in code_schema["description"]
def test_small_trace_summary_includes_effort_guidance(self):
"""Small traces should discourage transcript walkthroughs."""
rr = RRStep("test-model", config=RRConfig())
summary = rr._build_data_summary(
{
"question": "q",
"feedback": "Task PASSED",
"messages": [{"role": "user", "content": "hello"}],
}
)
assert "Expected effort" in summary
assert "2-4 focused execute_code checks" in summary
assert "Do not produce a transcript walkthrough" in summary
def test_prebuilt_model_and_callback_compose(self):
"""Pre-built Model + usage_callback both apply — meter wraps the instance."""
from pydantic_ai.models.test import TestModel
from ace.core.metered_model import MeteredModel
inner = TestModel()
rr = RRStep(
inner,
config=RRConfig(usage_callback=lambda u, n: None),
)
assert isinstance(rr._agent.model, MeteredModel)
assert rr._agent.model.wrapped is inner
|