Adding commentary and tests re heterogenous behavior (#517)

Author: sydney-runkle

pydantic_ai_examples/weather_agent.py

@@ -35,6 +35,8 @@ class Deps:
 
 weather_agent = Agent(
     'openai:gpt-4o',
+    # 'Be concise, reply with one sentence.' is enough for some models (like openai) to use
+    # the below tools appropriately, but others like anthropic and gemini require a bit more direction.
     system_prompt=(
         'Be concise, reply with one sentence.'
         'Use the `get_lat_lng` tool to get the latitude and longitude of the locations, '

tests/models/test_gemini.py

@@ -19,7 +19,6 @@
     ModelResponse,
     RetryPromptPart,
     SystemPromptPart,
-    TextPart,
     ToolCallPart,
     ToolReturnPart,
     UserPromptPart,
@@ -537,61 +536,6 @@ def handler(_: httpx.Request):
     assert str(exc_info.value) == snapshot('Unexpected response from gemini 401, body:\ninvalid request')
 
 
-async def test_heterogeneous_responses(get_gemini_client: GetGeminiClient):
-    """Indicates that tool calls are prioritized over text in heterogeneous responses."""
-    responses = [
-        gemini_response(
-            _content_model_response(
-                ModelResponse(
-                    parts=[TextPart(content='foo'), ToolCallPart.from_raw_args('get_location', {'loc_name': 'London'})]
-                )
-            )
-        ),
-        gemini_response(_content_model_response(ModelResponse.from_text('final response'))),
-    ]
-
-    gemini_client = get_gemini_client(responses)
-    m = GeminiModel('gemini-1.5-flash', http_client=gemini_client)
-    agent = Agent(m)
-
-    @agent.tool_plain
-    async def get_location(loc_name: str) -> str:
-        if loc_name == 'London':
-            return json.dumps({'lat': 51, 'lng': 0})
-        else:
-            raise ModelRetry('Wrong location, please try again')
-
-    result = await agent.run('Hello')
-    assert result.data == 'final response'
-    assert result.all_messages() == snapshot(
-        [
-            ModelRequest(
-                parts=[
-                    UserPromptPart(content='Hello', timestamp=IsNow(tz=timezone.utc)),
-                ]
-            ),
-            ModelResponse(
-                parts=[
-                    TextPart(content='foo'),
-                    ToolCallPart(
-                        tool_name='get_location',
-                        args=ArgsDict(args_dict={'loc_name': 'London'}),
-                    ),
-                ],
-                timestamp=IsNow(tz=timezone.utc),
-            ),
-            ModelRequest(
-                parts=[
-                    ToolReturnPart(
-                        tool_name='get_location', content='{"lat": 51, "lng": 0}', timestamp=IsNow(tz=timezone.utc)
-                    )
-                ]
-            ),
-            ModelResponse.from_text(content='final response', timestamp=IsNow(tz=timezone.utc)),
-        ]
-    )
-
-
 async def test_stream_text(get_gemini_client: GetGeminiClient):
     responses = [
         gemini_response(_content_model_response(ModelResponse.from_text('Hello '))),

tests/test_agent.py

@@ -1,3 +1,4 @@
+import json
 import sys
 from datetime import timezone
 from typing import Any, Callable, Union
@@ -16,6 +17,7 @@
     ModelMessage,
     ModelRequest,
     ModelResponse,
+    ModelResponsePart,
     RetryPromptPart,
     SystemPromptPart,
     TextPart,
@@ -1124,3 +1126,58 @@ def return_empty_text(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
 
     result = await agent.run('Hello')
     assert result.data == ('foo', 'bar')
+
+
+def test_heterogenous_reponses_non_streaming(set_event_loop: None) -> None:
+    """Indicates that tool calls are prioritized over text in heterogeneous responses."""
+
+    def return_model(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+        assert info.result_tools is not None
+        parts: list[ModelResponsePart] = []
+        if len(messages) == 1:
+            parts = [
+                TextPart(content='foo'),
+                ToolCallPart.from_raw_args('get_location', {'loc_name': 'London'}),
+            ]
+        else:
+            parts = [TextPart(content='final response')]
+        return ModelResponse(parts=parts)
+
+    agent = Agent(FunctionModel(return_model))
+
+    @agent.tool_plain
+    async def get_location(loc_name: str) -> str:
+        if loc_name == 'London':
+            return json.dumps({'lat': 51, 'lng': 0})
+        else:
+            raise ModelRetry('Wrong location, please try again')
+
+    result = agent.run_sync('Hello')
+    assert result.data == 'final response'
+    assert result.all_messages() == snapshot(
+        [
+            ModelRequest(
+                parts=[
+                    UserPromptPart(content='Hello', timestamp=IsNow(tz=timezone.utc)),
+                ]
+            ),
+            ModelResponse(
+                parts=[
+                    TextPart(content='foo'),
+                    ToolCallPart(
+                        tool_name='get_location',
+                        args=ArgsDict(args_dict={'loc_name': 'London'}),
+                    ),
+                ],
+                timestamp=IsNow(tz=timezone.utc),
+            ),
+            ModelRequest(
+                parts=[
+                    ToolReturnPart(
+                        tool_name='get_location', content='{"lat": 51, "lng": 0}', timestamp=IsNow(tz=timezone.utc)
+                    )
+                ]
+            ),
+            ModelResponse.from_text(content='final response', timestamp=IsNow(tz=timezone.utc)),
+        ]
+    )