Remove ArgsDict and ArgsJson (#769)

Author: dmontagu

docs/agents.md

@@ -463,7 +463,7 @@ with capture_run_messages() as messages:  # (2)!
                 parts=[
                     ToolCallPart(
                         tool_name='calc_volume',
-                        args=ArgsDict(args_dict={'size': 6}),
+                        args={'size': 6},
                         tool_call_id=None,
                         part_kind='tool-call',
                     )
@@ -488,7 +488,7 @@ with capture_run_messages() as messages:  # (2)!
                 parts=[
                     ToolCallPart(
                         tool_name='calc_volume',
-                        args=ArgsDict(args_dict={'size': 6}),
+                        args={'size': 6},
                         tool_call_id=None,
                         part_kind='tool-call',
                     )

docs/testing-evals.md

@@ -98,7 +98,6 @@ from dirty_equals import IsNow
 from pydantic_ai import models, capture_run_messages
 from pydantic_ai.models.test import TestModel
 from pydantic_ai.messages import (
-    ArgsDict,
     ModelResponse,
     SystemPromptPart,
     TextPart,
@@ -142,12 +141,10 @@ async def test_forecast():
             parts=[
                 ToolCallPart(
                     tool_name='weather_forecast',
-                    args=ArgsDict(
-                        args_dict={
-                            'location': 'a',
-                            'forecast_date': '2024-01-01',  # (8)!
-                        }
-                    ),
+                    args={
+                        'location': 'a',
+                        'forecast_date': '2024-01-01',  # (8)!
+                    },
                     tool_call_id=None,
                 )
             ],
@@ -223,9 +220,7 @@ 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_raw_args('weather_forecast', args)]
-        )
+        return ModelResponse(parts=[ToolCallPart('weather_forecast', args)])
     else:
         # second call, return the forecast
         msg = messages[-1].parts[0]

docs/tools.md

@@ -86,10 +86,7 @@ print(dice_result.all_messages())
     ModelResponse(
         parts=[
             ToolCallPart(
-                tool_name='roll_die',
-                args=ArgsDict(args_dict={}),
-                tool_call_id=None,
-                part_kind='tool-call',
+                tool_name='roll_die', args={}, tool_call_id=None, part_kind='tool-call'
             )
         ],
         model_name='function:model_logic',
@@ -112,7 +109,7 @@ print(dice_result.all_messages())
         parts=[
             ToolCallPart(
                 tool_name='get_player_name',
-                args=ArgsDict(args_dict={}),
+                args={},
                 tool_call_id=None,
                 part_kind='tool-call',
             )

pydantic_ai_slim/pydantic_ai/_parts_manager.py

@@ -221,7 +221,7 @@ def handle_tool_call_part(
             ModelResponseStreamEvent: A `PartStartEvent` indicating that a new tool call part
             has been added to the manager, or replaced an existing part.
         """
-        new_part = ToolCallPart.from_raw_args(tool_name=tool_name, args=args, tool_call_id=tool_call_id)
+        new_part = ToolCallPart(tool_name=tool_name, args=args, tool_call_id=tool_call_id)
         if vendor_part_id is None:
             # vendor_part_id is None, so we unconditionally append a new ToolCallPart to the end of the list
             new_part_index = len(self._parts)

pydantic_ai_slim/pydantic_ai/_result.py

@@ -201,14 +201,10 @@ def validate(
         """
         try:
             pyd_allow_partial: Literal['off', 'trailing-strings'] = 'trailing-strings' if allow_partial else 'off'
-            if isinstance(tool_call.args, _messages.ArgsJson):
-                result = self.type_adapter.validate_json(
-                    tool_call.args.args_json or '', experimental_allow_partial=pyd_allow_partial
-                )
+            if isinstance(tool_call.args, str):
+                result = self.type_adapter.validate_json(tool_call.args, experimental_allow_partial=pyd_allow_partial)
             else:
-                result = self.type_adapter.validate_python(
-                    tool_call.args.args_dict, experimental_allow_partial=pyd_allow_partial
-                )
+                result = self.type_adapter.validate_python(tool_call.args, experimental_allow_partial=pyd_allow_partial)
         except ValidationError as e:
             if wrap_validation_errors:
                 m = _messages.RetryPromptPart(

pydantic_ai_slim/pydantic_ai/messages.py

@@ -6,7 +6,6 @@
 
 import pydantic
 import pydantic_core
-from typing_extensions import Self, assert_never
 
 from ._utils import now_utc as _now_utc
 from .exceptions import UnexpectedModelBehavior
@@ -168,33 +167,17 @@ def has_content(self) -> bool:
         return bool(self.content)
 
 
-@dataclass
-class ArgsJson:
-    """Tool arguments as a JSON string."""
-
-    args_json: str
-    """A JSON string of arguments."""
-
-
-@dataclass
-class ArgsDict:
-    """Tool arguments as a Python dictionary."""
-
-    args_dict: dict[str, Any]
-    """A python dictionary of arguments."""
-
-
 @dataclass
 class ToolCallPart:
     """A tool call from a model."""
 
     tool_name: str
     """The name of the tool to call."""
 
-    args: ArgsJson | ArgsDict
+    args: str | dict[str, Any]
     """The arguments to pass to the tool.
 
-    Either as JSON or a Python dictionary depending on how data was returned.
+    This is stored either as a JSON string or a Python dictionary depending on how data was received.
     """
 
     tool_call_id: str | None = None
@@ -203,24 +186,14 @@ class ToolCallPart:
     part_kind: Literal['tool-call'] = 'tool-call'
     """Part type identifier, this is available on all parts as a discriminator."""
 
-    @classmethod
-    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, converting them to `ArgsJson` or `ArgsDict`."""
-        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)
-
     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)
+        if isinstance(self.args, dict):
+            return self.args
+        args = pydantic_core.from_json(self.args)
         assert isinstance(args, dict), 'args should be a dict'
         return cast(dict[str, Any], args)
 
@@ -229,16 +202,18 @@ def args_as_json_str(self) -> str:
 
         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()
+        if isinstance(self.args, str):
+            return self.args
+        return pydantic_core.to_json(self.args).decode()
 
     def has_content(self) -> bool:
         """Return `True` if the arguments contain any data."""
-        if isinstance(self.args, ArgsDict):
-            return any(self.args.args_dict.values())
+        if isinstance(self.args, dict):
+            # TODO: This should probably return True if you have the value False, or 0, etc.
+            #   It makes sense to me to ignore empty strings, but not sure about empty lists or dicts
+            return any(self.args.values())
         else:
-            return bool(self.args.args_json)
+            return bool(self.args)
 
 
 ModelResponsePart = Annotated[Union[TextPart, ToolCallPart], pydantic.Discriminator('part_kind')]
@@ -331,7 +306,7 @@ def as_part(self) -> ToolCallPart | None:
         if self.tool_name_delta is None or self.args_delta is None:
             return None
 
-        return ToolCallPart.from_raw_args(
+        return ToolCallPart(
             self.tool_name_delta,
             self.args_delta,
             self.tool_call_id,
@@ -396,7 +371,7 @@ def _apply_to_delta(self, delta: ToolCallPartDelta) -> ToolCallPart | ToolCallPa
 
         # If we now have enough data to create a full ToolCallPart, do so
         if delta.tool_name_delta is not None and delta.args_delta is not None:
-            return ToolCallPart.from_raw_args(
+            return ToolCallPart(
                 delta.tool_name_delta,
                 delta.args_delta,
                 delta.tool_call_id,
@@ -412,15 +387,15 @@ def _apply_to_part(self, part: ToolCallPart) -> ToolCallPart:
             part = replace(part, tool_name=tool_name)
 
         if isinstance(self.args_delta, str):
-            if not isinstance(part.args, ArgsJson):
+            if not isinstance(part.args, str):
                 raise UnexpectedModelBehavior(f'Cannot apply JSON deltas to non-JSON tool arguments ({part=}, {self=})')
-            updated_json = part.args.args_json + self.args_delta
-            part = replace(part, args=ArgsJson(updated_json))
+            updated_json = part.args + self.args_delta
+            part = replace(part, args=updated_json)
         elif isinstance(self.args_delta, dict):
-            if not isinstance(part.args, ArgsDict):
+            if not isinstance(part.args, dict):
                 raise UnexpectedModelBehavior(f'Cannot apply dict deltas to non-dict tool arguments ({part=}, {self=})')
-            updated_dict = {**(part.args.args_dict or {}), **self.args_delta}
-            part = replace(part, args=ArgsDict(updated_dict))
+            updated_dict = {**(part.args or {}), **self.args_delta}
+            part = replace(part, args=updated_dict)
 
         if self.tool_call_id:
             # Replace the tool_call_id entirely if given

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -13,7 +13,6 @@
 from .. import UnexpectedModelBehavior, _utils, usage
 from .._utils import guard_tool_call_id as _guard_tool_call_id
 from ..messages import (
-    ArgsDict,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -242,7 +241,7 @@ def _process_response(self, response: AnthropicMessage) -> ModelResponse:
             else:
                 assert isinstance(item, ToolUseBlock), 'unexpected item type'
                 items.append(
-                    ToolCallPart.from_raw_args(
+                    ToolCallPart(
                         tool_name=item.name,
                         args=cast(dict[str, Any], item.input),
                         tool_call_id=item.id,
@@ -319,7 +318,6 @@ def _map_message(messages: list[ModelMessage]) -> tuple[str, list[MessageParam]]
 
 
 def _map_tool_call(t: ToolCallPart) -> ToolUseBlockParam:
-    assert isinstance(t.args, ArgsDict), f'Expected ArgsDict, got {t.args}'
     return ToolUseBlockParam(
         id=_guard_tool_call_id(t=t, model_source='Anthropic'),
         type='tool_use',

pydantic_ai_slim/pydantic_ai/models/cohere.py

@@ -191,7 +191,7 @@ def _process_response(self, response: ChatResponse) -> ModelResponse:
         for c in response.message.tool_calls or []:
             if c.function and c.function.name and c.function.arguments:
                 parts.append(
-                    ToolCallPart.from_raw_args(
+                    ToolCallPart(
                         tool_name=c.function.name,
                         args=c.function.arguments,
                         tool_call_id=c.id,

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -453,7 +453,7 @@ def _process_response_from_parts(
             items.append(TextPart(content=part['text']))
         elif 'function_call' in part:
             items.append(
-                ToolCallPart.from_raw_args(
+                ToolCallPart(
                     tool_name=part['function_call']['name'],
                     args=part['function_call']['args'],
                 )

pydantic_ai_slim/pydantic_ai/models/groq.py

@@ -224,9 +224,7 @@ def _process_response(self, response: chat.ChatCompletion) -> ModelResponse:
             items.append(TextPart(content=choice.message.content))
         if choice.message.tool_calls is not None:
             for c in choice.message.tool_calls:
-                items.append(
-                    ToolCallPart.from_raw_args(tool_name=c.function.name, args=c.function.arguments, tool_call_id=c.id)
-                )
+                items.append(ToolCallPart(tool_name=c.function.name, args=c.function.arguments, tool_call_id=c.id))
         return ModelResponse(items, model_name=self.model_name, timestamp=timestamp)
 
     async def _process_streamed_response(self, response: AsyncStream[ChatCompletionChunk]) -> GroqStreamedResponse: