FlyingNunchucks/07-tool-using-agent

1from types import SimpleNamespace2 3import src.agent as agent_module4 5 6class FakeToolCall:7    def __init__(self, call_id: str, name: str, arguments: str):8        self.id = call_id9        self.function = SimpleNamespace(10            name=name,11            arguments=arguments,12        )13 14    def model_dump(self):15        return {16            "id": self.id,17            "type": "function",18            "function": {19                "name": self.function.name,20                "arguments": self.function.arguments,21            },22        }23 24 25class FakeMessage:26    def __init__(self, content=None, tool_calls=None):27        self.content = content28        self.tool_calls = tool_calls29 30 31class FakeResponse:32    def __init__(self, message):33        self.choices = [SimpleNamespace(message=message)]34 35 36class FakeCompletions:37    def __init__(self, responses):38        self._responses = iter(responses)39 40    def create(self, **_kwargs):41        return next(self._responses)42 43 44class FakeClient:45    def __init__(self, responses):46        self.chat = SimpleNamespace(47            completions=FakeCompletions(responses)48        )49 50 51def _patch_client(monkeypatch, responses):52    monkeypatch.setattr(53        agent_module,54        "get_client",55        lambda: FakeClient(responses),56    )57    monkeypatch.setattr(58        agent_module,59        "write_audit_record",60        lambda call, result: None,61    )62 63 64def test_run_agent_iter_emits_real_authorization_and_execution_events(monkeypatch):65    _patch_client(66        monkeypatch,67        [68            FakeResponse(69                FakeMessage(70                    tool_calls=[71                        FakeToolCall(72                            "call-1",73                            "calculator",74                            '{"expression": "12 * 4"}',75                        )76                    ]77                )78            ),79            FakeResponse(80                FakeMessage(81                    content="The result is 48.",82                    tool_calls=None,83                )84            ),85        ],86    )87 88    events = list(89        agent_module.run_agent_iter("What is 12 multiplied by 4?")90    )91 92    event_names = [event.event for event in events]93 94    assert event_names == [95        "run_started",96        "model_thinking",97        "tool_requested",98        "tool_approved",99        "tool_succeeded",100        "model_thinking",101        "final_answer",102    ]103 104    approved = next(event for event in events if event.event == "tool_approved")105    succeeded = next(event for event in events if event.event == "tool_succeeded")106    final = events[-1]107 108    assert approved.authorization is not None109    assert approved.authorization.status == "approved"110    assert succeeded.result is not None111    assert succeeded.result.status == "success"112    assert succeeded.result.output["result"] == 48113    assert final.final_result is not None114    assert final.final_result.final_answer == "The result is 48."115 116 117def test_run_agent_iter_reports_blocked_call_before_execution(monkeypatch):118    _patch_client(119        monkeypatch,120        [121            FakeResponse(122                FakeMessage(123                    tool_calls=[124                        FakeToolCall(125                            "call-blocked",126                            "run_shell_command",127                            '{"command": "whoami"}',128                        )129                    ]130                )131            ),132            FakeResponse(133                FakeMessage(134                    content="That capability is not available.",135                    tool_calls=None,136                )137            ),138        ],139    )140 141    events = list(142        agent_module.run_agent_iter("Run a shell command.")143    )144 145    blocked = next(event for event in events if event.event == "tool_blocked")146 147    assert blocked.authorization is not None148    assert blocked.authorization.status == "blocked"149    assert blocked.result is not None150    assert blocked.result.status == "blocked"151    assert "not registered" in blocked.message152    assert not any(event.event == "tool_succeeded" for event in events)153 154 155def test_run_agent_preserves_final_result_api(monkeypatch):156    _patch_client(157        monkeypatch,158        [159            FakeResponse(160                FakeMessage(161                    content="No tool was needed.",162                    tool_calls=None,163                )164            )165        ],166    )167 168    result = agent_module.run_agent("Say hello without tools.")169 170    assert result.final_answer == "No tool was needed."171    assert result.tool_calls == []172    assert result.tool_results == []173