fixing remaining tests, breaking TestModel

Author: samuelcolvin

pydantic_ai_slim/pydantic_ai/_utils.py

@@ -126,6 +126,12 @@ def is_left(self) -> bool:
     def whichever(self) -> Left | Right:
         return self._left.value if self._left is not None else self.right
 
+    def __repr__(self):
+        if left := self._left:
+            return f'Either(left={left.value!r})'
+        else:
+            return f'Either(right={self.right!r})'
+
 
 @asynccontextmanager
 async def group_by_temporal(

pydantic_ai_slim/pydantic_ai/models/test.py

@@ -59,7 +59,7 @@ class TestModel(Model):
     agent_model_function_tools: dict[str, ToolDefinition] | None = field(default=None, init=False)
     agent_model_allow_text_result: bool | None = field(default=None, init=False)
     agent_model_result_tools: dict[str, ToolDefinition] | None = field(default=None, init=False)
-    _agent_model: TestAgentModel | None = field(default=None, init=False)
+    agent_model: TestAgentModel | None = field(default=None, init=False)
 
     async def prepare(
         self,
@@ -71,40 +71,43 @@ async def prepare(
         self.agent_model_function_tools = function_tools
         self.agent_model_allow_text_result = allow_text_result
         self.agent_model_result_tools = result_tools
-        if self._agent_model is None:
-            self._agent_model = self._build_agent()
-        return self._agent_model
+        if self.agent_model is None:
+            self.agent_model = self._build_agent(function_tools, allow_text_result, result_tools)
+        return self.agent_model
 
-    def _build_agent(self):
+    def _build_agent(
+        self,
+        function_tools: dict[str, ToolDefinition],
+        allow_text_result: bool,
+        result_tools: dict[str, ToolDefinition] | None,
+    ) -> TestAgentModel:
         if self.call_tools == 'all':
-            tool_calls = [(r.name, r) for r in self.agent_model_function_tools.values()]
+            tool_calls = [(r.name, r) for r in function_tools.values()]
         else:
-            tools_to_call = (self.agent_model_function_tools[name] for name in self.call_tools)
+            tools_to_call = (function_tools[name] for name in self.call_tools)
             tool_calls = [(r.name, r) for r in tools_to_call]
 
         if self.custom_result_text is not None:
-            assert self.agent_model_allow_text_result, 'Plain response not allowed, but `custom_result_text` is set.'
+            assert allow_text_result, 'Plain response not allowed, but `custom_result_text` is set.'
             assert self.custom_result_args is None, 'Cannot set both `custom_result_text` and `custom_result_args`.'
             result: _utils.Either[str | None, Any | None] = _utils.Either(left=self.custom_result_text)
         elif self.custom_result_args is not None:
-            assert (
-                self.agent_model_result_tools is not None
-            ), 'No result tools provided, but `custom_result_args` is set.'
-            result_tool = next(iter(self.agent_model_result_tools.values()))
+            assert result_tools is not None, 'No result tools provided, but `custom_result_args` is set.'
+            result_tool = next(iter(result_tools.values()))
 
             if k := result_tool.outer_typed_dict_key:
                 result = _utils.Either(right={k: self.custom_result_args})
             else:
                 result = _utils.Either(right=self.custom_result_args)
-        elif self.agent_model_allow_text_result:
+        elif allow_text_result:
             result = _utils.Either(left=None)
-        elif self.agent_model_result_tools is not None:
+        elif result_tools is not None:
             result = _utils.Either(right=None)
         else:
             result = _utils.Either(left=None)
 
-        result_tools_list = list(self.agent_model_function_tools.values()) if self.agent_model_function_tools else None
-        return TestAgentModel(tool_calls, result, result_tools_list, self.seed)
+        result_tools = list(result_tools.values()) if result_tools is not None else None
+        return TestAgentModel(tool_calls, result, result_tools, self.seed)
 
     def name(self) -> str:
         return 'test-model'

tests/test_agent.py

@@ -24,6 +24,7 @@
 from pydantic_ai.models.function import AgentInfo, FunctionModel
 from pydantic_ai.models.test import TestModel
 from pydantic_ai.result import Cost, RunResult
+from pydantic_ai.tools import ToolDefinition
 
 from .conftest import IsNow, TestEnv
 
@@ -34,7 +35,7 @@ def test_result_tuple(set_event_loop: None):
     def return_tuple(_: list[Message], info: AgentInfo) -> ModelAnyResponse:
         assert info.result_tools is not None
         args_json = '{"response": ["foo", "bar"]}'
-        return ModelStructuredResponse(calls=[ToolCall.from_json(info.result_tools[0].name, args_json)])
+        return ModelStructuredResponse(calls=[ToolCall.from_json(info.result_tools['final_result'].name, args_json)])
 
     agent = Agent(FunctionModel(return_tuple), result_type=tuple[str, str])
 
@@ -51,7 +52,7 @@ def test_result_pydantic_model(set_event_loop: None):
     def return_model(_: list[Message], info: AgentInfo) -> ModelAnyResponse:
         assert info.result_tools is not None
         args_json = '{"a": 1, "b": "foo"}'
-        return ModelStructuredResponse(calls=[ToolCall.from_json(info.result_tools[0].name, args_json)])
+        return ModelStructuredResponse(calls=[ToolCall.from_json(info.result_tools['final_result'].name, args_json)])
 
     agent = Agent(FunctionModel(return_model), result_type=Foo)
 
@@ -67,7 +68,7 @@ def return_model(messages: list[Message], info: AgentInfo) -> ModelAnyResponse:
             args_json = '{"a": "wrong", "b": "foo"}'
         else:
             args_json = '{"a": 42, "b": "foo"}'
-        return ModelStructuredResponse(calls=[ToolCall.from_json(info.result_tools[0].name, args_json)])
+        return ModelStructuredResponse(calls=[ToolCall.from_json(info.result_tools['final_result'].name, args_json)])
 
     agent = Agent(FunctionModel(return_model), result_type=Foo)
 
@@ -112,7 +113,7 @@ def return_model(messages: list[Message], info: AgentInfo) -> ModelAnyResponse:
             args_json = '{"a": 41, "b": "foo"}'
         else:
             args_json = '{"a": 42, "b": "foo"}'
-        return ModelStructuredResponse(calls=[ToolCall.from_json(info.result_tools[0].name, args_json)])
+        return ModelStructuredResponse(calls=[ToolCall.from_json(info.result_tools['final_result'].name, args_json)])
 
     agent = Agent(FunctionModel(return_model), result_type=Foo)
 
@@ -153,7 +154,9 @@ def return_tuple(_: list[Message], info: AgentInfo) -> ModelAnyResponse:
             return ModelTextResponse(content='hello')
         else:
             args_json = '{"response": ["foo", "bar"]}'
-            return ModelStructuredResponse(calls=[ToolCall.from_json(info.result_tools[0].name, args_json)])
+            return ModelStructuredResponse(
+                calls=[ToolCall.from_json(info.result_tools['final_result'].name, args_json)]
+            )
 
     agent = Agent(FunctionModel(return_tuple), result_type=tuple[str, str])
 
@@ -191,27 +194,26 @@ def test_response_tuple(set_event_loop: None):
     assert m.agent_model_result_tools is not None
     assert len(m.agent_model_result_tools) == 1
 
-    # to match the protocol, we just extract the attributes we care about
-    fields = 'name', 'description', 'parameters_json_schema', 'outer_typed_dict_key'
-    agent_model_result_tool = {f: getattr(m.agent_model_result_tools[0], f) for f in fields}
-    assert agent_model_result_tool == snapshot(
+    assert m.agent_model_result_tools == snapshot(
         {
-            'name': 'final_result',
-            'description': 'The final response which ends this conversation',
-            'parameters_json_schema': {
-                'properties': {
-                    'response': {
-                        'maxItems': 2,
-                        'minItems': 2,
-                        'prefixItems': [{'type': 'string'}, {'type': 'string'}],
-                        'title': 'Response',
-                        'type': 'array',
-                    }
+            'final_result': ToolDefinition(
+                name='final_result',
+                description='The final response which ends this conversation',
+                parameters_json_schema={
+                    'properties': {
+                        'response': {
+                            'maxItems': 2,
+                            'minItems': 2,
+                            'prefixItems': [{'type': 'string'}, {'type': 'string'}],
+                            'title': 'Response',
+                            'type': 'array',
+                        }
+                    },
+                    'required': ['response'],
+                    'type': 'object',
                 },
-                'required': ['response'],
-                'type': 'object',
-            },
-            'outer_typed_dict_key': 'response',
+                outer_typed_dict_key='response',
+            )
         }
     )
 
@@ -250,20 +252,21 @@ def validate_result(ctx: RunContext[None], r: Any) -> Any:
     assert m.agent_model_result_tools is not None
     assert len(m.agent_model_result_tools) == 1
 
-    # to match the protocol, we just extract the attributes we care about
-    fields = 'name', 'description', 'parameters_json_schema', 'outer_typed_dict_key'
-    agent_model_result_tool = {f: getattr(m.agent_model_result_tools[0], f) for f in fields}
-    assert agent_model_result_tool == snapshot(
+    assert m.agent_model_result_tools == snapshot(
         {
-            'name': 'final_result',
-            'description': 'The final response which ends this conversation',
-            'parameters_json_schema': {
-                'properties': {'a': {'title': 'A', 'type': 'integer'}, 'b': {'title': 'B', 'type': 'string'}},
-                'title': 'Foo',
-                'required': ['a', 'b'],
-                'type': 'object',
-            },
-            'outer_typed_dict_key': None,
+            'final_result': ToolDefinition(
+                name='final_result',
+                description='The final response which ends this conversation',
+                parameters_json_schema={
+                    'properties': {
+                        'a': {'title': 'A', 'type': 'integer'},
+                        'b': {'title': 'B', 'type': 'string'},
+                    },
+                    'required': ['a', 'b'],
+                    'title': 'Foo',
+                    'type': 'object',
+                },
+            )
         }
     )
 
@@ -324,15 +327,12 @@ def validate_result(ctx: RunContext[None], r: Any) -> Any:
     assert m.agent_model_result_tools is not None
     assert len(m.agent_model_result_tools) == 2
 
-    # to match the protocol, we just extract the attributes we care about
-    fields = 'name', 'description', 'parameters_json_schema', 'outer_typed_dict_key'
-    agent_model_result_tools = [{f: getattr(t, f) for f in fields} for t in m.agent_model_result_tools]
-    assert agent_model_result_tools == snapshot(
-        [
-            {
-                'name': 'final_result_Foo',
-                'description': 'Foo: The final response which ends this conversation',
-                'parameters_json_schema': {
+    assert m.agent_model_result_tools == snapshot(
+        {
+            'final_result_Foo': ToolDefinition(
+                name='final_result_Foo',
+                description='Foo: The final response which ends this conversation',
+                parameters_json_schema={
                     'properties': {
                         'a': {'title': 'A', 'type': 'integer'},
                         'b': {'title': 'B', 'type': 'string'},
@@ -341,20 +341,18 @@ def validate_result(ctx: RunContext[None], r: Any) -> Any:
                     'title': 'Foo',
                     'type': 'object',
                 },
-                'outer_typed_dict_key': None,
-            },
-            {
-                'name': 'final_result_Bar',
-                'description': 'This is a bar model.',
-                'parameters_json_schema': {
+            ),
+            'final_result_Bar': ToolDefinition(
+                name='final_result_Bar',
+                description='This is a bar model.',
+                parameters_json_schema={
                     'properties': {'b': {'title': 'B', 'type': 'string'}},
                     'required': ['b'],
                     'title': 'Bar',
                     'type': 'object',
                 },
-                'outer_typed_dict_key': None,
-            },
-        ]
+            ),
+        }
     )
 
     result = agent.run_sync('Hello', model=TestModel(seed=1))
@@ -383,6 +381,7 @@ async def ret_a(x: str) -> str:
             ModelTextResponse(content='{"ret_a":"a-apple"}', timestamp=IsNow(tz=timezone.utc)),
         ]
     )
+    m.agent_model = None
 
     # if we pass new_messages, system prompt is inserted before the message_history messages
     result2 = agent.run_sync('Hello again', message_history=result1.new_messages())
@@ -415,6 +414,8 @@ async def ret_a(x: str) -> str:
     assert new_msg_roles == snapshot(['user', 'model-structured-response', 'tool-return', 'model-text-response'])
     assert result2.new_messages_json().startswith(b'[{"content":"Hello again",')
 
+    m.agent_model = None
+
     # if we pass all_messages, system prompt is NOT inserted before the message_history messages,
     # so only one system prompt
     result3 = agent.run_sync('Hello again', message_history=result1.all_messages())
@@ -540,6 +541,7 @@ async def make_request() -> str:
     for _ in range(2):
         result = agent.run_sync('Hello')
         assert result.data == '{"make_request":"200"}'
+        agent.model.agent_model = None
 
 
 async def test_agent_name():

tests/test_deps.py

@@ -21,19 +21,23 @@ async def example_tool(ctx: RunContext[MyDeps]) -> str:
 def test_deps_used(set_event_loop: None):
     result = agent.run_sync('foobar', deps=MyDeps(foo=1, bar=2))
     assert result.data == '{"example_tool":"MyDeps(foo=1, bar=2)"}'
+    agent.model.agent_model = None
 
 
 def test_deps_override(set_event_loop: None):
     with agent.override(deps=MyDeps(foo=3, bar=4)):
         result = agent.run_sync('foobar', deps=MyDeps(foo=1, bar=2))
         assert result.data == '{"example_tool":"MyDeps(foo=3, bar=4)"}'
+        agent.model.agent_model = None
 
         with agent.override(deps=MyDeps(foo=5, bar=6)):
             result = agent.run_sync('foobar', deps=MyDeps(foo=1, bar=2))
             assert result.data == '{"example_tool":"MyDeps(foo=5, bar=6)"}'
+            agent.model.agent_model = None
 
         result = agent.run_sync('foobar', deps=MyDeps(foo=1, bar=2))
         assert result.data == '{"example_tool":"MyDeps(foo=3, bar=4)"}'
+        agent.model.agent_model = None
 
     result = agent.run_sync('foobar', deps=MyDeps(foo=1, bar=2))
     assert result.data == '{"example_tool":"MyDeps(foo=1, bar=2)"}'