Make args handling more robust (#489)

Author: dmontagu

docs/testing-evals.md

@@ -219,7 +219,9 @@ def call_weather_forecast(  # (1)!
         m = re.search(r'\d{4}-\d{2}-\d{2}', user_prompt.content)
         assert m is not None
         args = {'location': 'London', 'forecast_date': m.group()}  # (2)!
-        return ModelResponse(parts=[ToolCallPart.from_dict('weather_forecast', args)])
+        return ModelResponse(
+            parts=[ToolCallPart.from_raw_args('weather_forecast', args)]
+        )
     else:
         # second call, return the forecast
         msg = messages[-1].parts[0]

pydantic_ai_slim/pydantic_ai/messages.py

@@ -2,11 +2,11 @@
 
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Annotated, Any, Literal, Union
+from typing import Annotated, Any, Literal, Union, cast
 
 import pydantic
 import pydantic_core
-from typing_extensions import Self
+from typing_extensions import Self, assert_never
 
 from ._utils import now_utc as _now_utc
 
@@ -190,12 +190,34 @@ class ToolCallPart:
     """Part type identifier, this is available on all parts as a discriminator."""
 
     @classmethod
-    def from_json(cls, tool_name: str, args_json: str, tool_call_id: str | None = None) -> Self:
-        return cls(tool_name, ArgsJson(args_json), tool_call_id)
+    def from_raw_args(cls, tool_name: str, args: str | dict[str, Any], tool_call_id: str | None = None) -> Self:
+        """Create a `ToolCallPart` from raw arguments."""
+        if isinstance(args, str):
+            return cls(tool_name, ArgsJson(args), tool_call_id)
+        elif isinstance(args, dict):
+            return cls(tool_name, ArgsDict(args), tool_call_id)
+        else:
+            assert_never(args)
 
-    @classmethod
-    def from_dict(cls, tool_name: str, args_dict: dict[str, Any], tool_call_id: str | None = None) -> Self:
-        return cls(tool_name, ArgsDict(args_dict), tool_call_id)
+    def args_as_dict(self) -> dict[str, Any]:
+        """Return the arguments as a Python dictionary.
+
+        This is just for convenience with models that require dicts as input.
+        """
+        if isinstance(self.args, ArgsDict):
+            return self.args.args_dict
+        args = pydantic_core.from_json(self.args.args_json)
+        assert isinstance(args, dict), 'args should be a dict'
+        return cast(dict[str, Any], args)
+
+    def args_as_json_str(self) -> str:
+        """Return the arguments as a JSON string.
+
+        This is just for convenience with models that require JSON strings as input.
+        """
+        if isinstance(self.args, ArgsJson):
+            return self.args.args_json
+        return pydantic_core.to_json(self.args.args_dict).decode()
 
     def has_content(self) -> bool:
         if isinstance(self.args, ArgsDict):

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -220,7 +220,7 @@ def _process_response(response: AnthropicMessage) -> ModelResponse:
             else:
                 assert isinstance(item, ToolUseBlock), 'unexpected item type'
                 items.append(
-                    ToolCallPart.from_dict(
+                    ToolCallPart.from_raw_args(
                         item.name,
                         cast(dict[str, Any], item.input),
                         item.id,
@@ -311,7 +311,7 @@ def _map_tool_call(t: ToolCallPart) -> ToolUseBlockParam:
         id=_guard_tool_call_id(t=t, model_source='Anthropic'),
         type='tool_use',
         name=t.tool_name,
-        input=t.args.args_dict,
+        input=t.args_as_dict(),
     )
 
 

pydantic_ai_slim/pydantic_ai/models/function.py

@@ -9,12 +9,10 @@
 from itertools import chain
 from typing import Callable, Union, cast
 
-import pydantic_core
 from typing_extensions import TypeAlias, assert_never, overload
 
 from .. import _utils, result
 from ..messages import (
-    ArgsJson,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -232,7 +230,7 @@ def get(self, *, final: bool = False) -> ModelResponse:
         calls: list[ModelResponsePart] = []
         for c in self._delta_tool_calls.values():
             if c.name is not None and c.json_args is not None:
-                calls.append(ToolCallPart.from_json(c.name, c.json_args))
+                calls.append(ToolCallPart.from_raw_args(c.name, c.json_args))
 
         return ModelResponse(calls, timestamp=self._timestamp)
 
@@ -268,11 +266,7 @@ def _estimate_usage(messages: Iterable[ModelMessage]) -> result.Usage:
                     response_tokens += _estimate_string_usage(part.content)
                 elif isinstance(part, ToolCallPart):
                     call = part
-                    if isinstance(call.args, ArgsJson):
-                        args_str = call.args.args_json
-                    else:
-                        args_str = pydantic_core.to_json(call.args.args_dict).decode()
-                    response_tokens += 1 + _estimate_string_usage(args_str)
+                    response_tokens += 1 + _estimate_string_usage(call.args_as_json_str())
                 else:
                     assert_never(part)
         else:

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -16,7 +16,6 @@
 
 from .. import UnexpectedModelBehavior, _utils, exceptions, result
 from ..messages import (
-    ArgsDict,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -460,8 +459,7 @@ class _GeminiFunctionCallPart(TypedDict):
 
 
 def _function_call_part_from_call(tool: ToolCallPart) -> _GeminiFunctionCallPart:
-    assert isinstance(tool.args, ArgsDict), f'Expected ArgsObject, got {tool.args}'
-    return _GeminiFunctionCallPart(function_call=_GeminiFunctionCall(name=tool.tool_name, args=tool.args.args_dict))
+    return _GeminiFunctionCallPart(function_call=_GeminiFunctionCall(name=tool.tool_name, args=tool.args_as_dict()))
 
 
 def _process_response_from_parts(parts: Sequence[_GeminiPartUnion], timestamp: datetime | None = None) -> ModelResponse:
@@ -470,7 +468,7 @@ def _process_response_from_parts(parts: Sequence[_GeminiPartUnion], timestamp: d
         if 'text' in part:
             items.append(TextPart(part['text']))
         elif 'function_call' in part:
-            items.append(ToolCallPart.from_dict(part['function_call']['name'], part['function_call']['args']))
+            items.append(ToolCallPart.from_raw_args(part['function_call']['name'], part['function_call']['args']))
         elif 'function_response' in part:
             raise exceptions.UnexpectedModelBehavior(
                 f'Unsupported response from Gemini, expected all parts to be function calls or text, got: {part!r}'

pydantic_ai_slim/pydantic_ai/models/groq.py

@@ -13,7 +13,6 @@
 from .. import UnexpectedModelBehavior, _utils, result
 from .._utils import guard_tool_call_id as _guard_tool_call_id
 from ..messages import (
-    ArgsJson,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -221,7 +220,7 @@ def _process_response(response: chat.ChatCompletion) -> ModelResponse:
             items.append(TextPart(choice.message.content))
         if choice.message.tool_calls is not None:
             for c in choice.message.tool_calls:
-                items.append(ToolCallPart.from_json(c.function.name, c.function.arguments, c.id))
+                items.append(ToolCallPart.from_raw_args(c.function.name, c.function.arguments, c.id))
         return ModelResponse(items, timestamp=timestamp)
 
     @staticmethod
@@ -380,7 +379,7 @@ def get(self, *, final: bool = False) -> ModelResponse:
         for c in self._delta_tool_calls.values():
             if f := c.function:
                 if f.name is not None and f.arguments is not None:
-                    items.append(ToolCallPart.from_json(f.name, f.arguments, c.id))
+                    items.append(ToolCallPart.from_raw_args(f.name, f.arguments, c.id))
 
         return ModelResponse(items, timestamp=self._timestamp)
 
@@ -392,11 +391,10 @@ def timestamp(self) -> datetime:
 
 
 def _map_tool_call(t: ToolCallPart) -> chat.ChatCompletionMessageToolCallParam:
-    assert isinstance(t.args, ArgsJson), f'Expected ArgsJson, got {t.args}'
     return chat.ChatCompletionMessageToolCallParam(
         id=_guard_tool_call_id(t=t, model_source='Groq'),
         type='function',
-        function={'name': t.tool_name, 'arguments': t.args.args_json},
+        function={'name': t.tool_name, 'arguments': t.args_as_json_str()},
     )
 
 

pydantic_ai_slim/pydantic_ai/models/mistral.py

@@ -560,13 +560,10 @@ def get(self, *, final: bool = False) -> ModelResponse:
                     # when `{"response":` then `repair_json` sets `{"response": ""}` (type not found default str)
                     # when `{"response": {` then `repair_json` sets `{"response": {}}` (type found)
                     # This ensures it's corrected to `{"response": {}}` and other required parameters and type.
-                    if not self._validate_required_json_shema(output_json, result_tool.parameters_json_schema):
+                    if not self._validate_required_json_schema(output_json, result_tool.parameters_json_schema):
                         continue
 
-                    tool = ToolCallPart.from_dict(
-                        tool_name=result_tool.name,
-                        args_dict=output_json,
-                    )
+                    tool = ToolCallPart.from_raw_args(result_tool.name, output_json)
                     calls.append(tool)
 
         return ModelResponse(calls, timestamp=self._timestamp)
@@ -578,7 +575,7 @@ def timestamp(self) -> datetime:
         return self._timestamp
 
     @staticmethod
-    def _validate_required_json_shema(json_dict: dict[str, Any], json_schema: dict[str, Any]) -> bool:
+    def _validate_required_json_schema(json_dict: dict[str, Any], json_schema: dict[str, Any]) -> bool:
         """Validate that all required parameters in the JSON schema are present in the JSON dictionary."""
         required_params = json_schema.get('required', [])
         properties = json_schema.get('properties', {})
@@ -602,7 +599,7 @@ def _validate_required_json_shema(json_dict: dict[str, Any], json_schema: dict[s
 
             if isinstance(json_dict[param], dict) and 'properties' in param_schema:
                 nested_schema = param_schema
-                if not MistralStreamStructuredResponse._validate_required_json_shema(json_dict[param], nested_schema):
+                if not MistralStreamStructuredResponse._validate_required_json_schema(json_dict[param], nested_schema):
                     return False
 
         return True
@@ -633,16 +630,7 @@ def _map_mistral_to_pydantic_tool_call(tool_call: MistralToolCall) -> ToolCallPa
     tool_call_id = tool_call.id or None
     func_call = tool_call.function
 
-    if isinstance(func_call.arguments, str):
-        return ToolCallPart.from_json(
-            tool_name=func_call.name,
-            args_json=func_call.arguments,
-            tool_call_id=tool_call_id,
-        )
-    else:
-        return ToolCallPart.from_dict(
-            tool_name=func_call.name, args_dict=func_call.arguments, tool_call_id=tool_call_id
-        )
+    return ToolCallPart.from_raw_args(func_call.name, func_call.arguments, tool_call_id)
 
 
 def _map_usage(response: MistralChatCompletionResponse | MistralCompletionChunk) -> Usage:

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -13,7 +13,6 @@
 from .. import UnexpectedModelBehavior, _utils, result
 from .._utils import guard_tool_call_id as _guard_tool_call_id
 from ..messages import (
-    ArgsJson,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -211,7 +210,7 @@ def _process_response(response: chat.ChatCompletion) -> ModelResponse:
             items.append(TextPart(choice.message.content))
         if choice.message.tool_calls is not None:
             for c in choice.message.tool_calls:
-                items.append(ToolCallPart.from_json(c.function.name, c.function.arguments, c.id))
+                items.append(ToolCallPart.from_raw_args(c.function.name, c.function.arguments, c.id))
         return ModelResponse(items, timestamp=timestamp)
 
     @staticmethod
@@ -372,7 +371,7 @@ def get(self, *, final: bool = False) -> ModelResponse:
         for c in self._delta_tool_calls.values():
             if f := c.function:
                 if f.name is not None and f.arguments is not None:
-                    items.append(ToolCallPart.from_json(f.name, f.arguments, c.id))
+                    items.append(ToolCallPart.from_raw_args(f.name, f.arguments, c.id))
 
         return ModelResponse(items, timestamp=self._timestamp)
 
@@ -384,11 +383,10 @@ def timestamp(self) -> datetime:
 
 
 def _map_tool_call(t: ToolCallPart) -> chat.ChatCompletionMessageToolCallParam:
-    assert isinstance(t.args, ArgsJson), f'Expected ArgsJson, got {t.args}'
     return chat.ChatCompletionMessageToolCallParam(
         id=_guard_tool_call_id(t=t, model_source='OpenAI'),
         type='function',
-        function={'name': t.tool_name, 'arguments': t.args.args_json},
+        function={'name': t.tool_name, 'arguments': t.args_as_json_str()},
     )
 
 

pydantic_ai_slim/pydantic_ai/models/test.py

@@ -167,7 +167,7 @@ def _request(self, messages: list[ModelMessage], model_settings: ModelSettings |
         # if there are tools, the first thing we want to do is call all of them
         if self.tool_calls and not any(isinstance(m, ModelResponse) for m in messages):
             return ModelResponse(
-                parts=[ToolCallPart.from_dict(name, self.gen_tool_args(args)) for name, args in self.tool_calls]
+                parts=[ToolCallPart.from_raw_args(name, self.gen_tool_args(args)) for name, args in self.tool_calls]
             )
 
         if messages:
@@ -179,7 +179,7 @@ def _request(self, messages: list[ModelMessage], model_settings: ModelSettings |
             if new_retry_names:
                 return ModelResponse(
                     parts=[
-                        ToolCallPart.from_dict(name, self.gen_tool_args(args))
+                        ToolCallPart.from_raw_args(name, self.gen_tool_args(args))
                         for name, args in self.tool_calls
                         if name in new_retry_names
                     ]
@@ -205,10 +205,10 @@ def _request(self, messages: list[ModelMessage], model_settings: ModelSettings |
             custom_result_args = self.result.right
             result_tool = self.result_tools[self.seed % len(self.result_tools)]
             if custom_result_args is not None:
-                return ModelResponse(parts=[ToolCallPart.from_dict(result_tool.name, custom_result_args)])
+                return ModelResponse(parts=[ToolCallPart.from_raw_args(result_tool.name, custom_result_args)])
             else:
                 response_args = self.gen_tool_args(result_tool)
-                return ModelResponse(parts=[ToolCallPart.from_dict(result_tool.name, response_args)])
+                return ModelResponse(parts=[ToolCallPart.from_raw_args(result_tool.name, response_args)])
 
 
 @dataclass

tests/models/test_gemini.py

@@ -413,7 +413,7 @@ async def test_text_success(get_gemini_client: GetGeminiClient):
 async def test_request_structured_response(get_gemini_client: GetGeminiClient):
     response = gemini_response(
         _content_model_response(
-            ModelResponse(parts=[ToolCallPart.from_dict('final_result', {'response': [1, 2, 123]})])
+            ModelResponse(parts=[ToolCallPart.from_raw_args('final_result', {'response': [1, 2, 123]})])
         )
     )
     gemini_client = get_gemini_client(response)
@@ -449,12 +449,12 @@ async def test_request_tool_call(get_gemini_client: GetGeminiClient):
     responses = [
         gemini_response(
             _content_model_response(
-                ModelResponse(parts=[ToolCallPart.from_dict('get_location', {'loc_name': 'San Fransisco'})])
+                ModelResponse(parts=[ToolCallPart.from_raw_args('get_location', {'loc_name': 'San Fransisco'})])
             )
         ),
         gemini_response(
             _content_model_response(
-                ModelResponse(parts=[ToolCallPart.from_dict('get_location', {'loc_name': 'London'})])
+                ModelResponse(parts=[ToolCallPart.from_raw_args('get_location', {'loc_name': 'London'})])
             )
         ),
         gemini_response(_content_model_response(ModelResponse.from_text('final response'))),
@@ -596,7 +596,7 @@ async def test_stream_structured(get_gemini_client: GetGeminiClient):
     responses = [
         gemini_response(
             _content_model_response(
-                ModelResponse(parts=[ToolCallPart.from_dict('final_result', {'response': [1, 2]})])
+                ModelResponse(parts=[ToolCallPart.from_raw_args('final_result', {'response': [1, 2]})])
             ),
         ),
     ]
@@ -615,10 +615,10 @@ async def test_stream_structured(get_gemini_client: GetGeminiClient):
 async def test_stream_structured_tool_calls(get_gemini_client: GetGeminiClient):
     first_responses = [
         gemini_response(
-            _content_model_response(ModelResponse(parts=[ToolCallPart.from_dict('foo', {'x': 'a'})])),
+            _content_model_response(ModelResponse(parts=[ToolCallPart.from_raw_args('foo', {'x': 'a'})])),
         ),
         gemini_response(
-            _content_model_response(ModelResponse(parts=[ToolCallPart.from_dict('bar', {'y': 'b'})])),
+            _content_model_response(ModelResponse(parts=[ToolCallPart.from_raw_args('bar', {'y': 'b'})])),
         ),
     ]
     d1 = _gemini_streamed_response_ta.dump_json(first_responses, by_alias=True)
@@ -627,7 +627,7 @@ async def test_stream_structured_tool_calls(get_gemini_client: GetGeminiClient):
     second_responses = [
         gemini_response(
             _content_model_response(
-                ModelResponse(parts=[ToolCallPart.from_dict('final_result', {'response': [1, 2]})])
+                ModelResponse(parts=[ToolCallPart.from_raw_args('final_result', {'response': [1, 2]})])
             ),
         ),
     ]