pydantic-ai-slim 0.0.19__py3-none-any.whl → 0.0.21__py3-none-any.whl

pydantic_ai/models/mistral.py

@@ -6,7 +6,7 @@ from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from itertools import chain
-from typing import Any, Callable, Literal, Union
+from typing import Any, Callable, Literal, Union, cast
 
 import pydantic_core
 from httpx import AsyncClient as AsyncHTTPClient, Timeout
@@ -15,7 +15,6 @@ from typing_extensions import assert_never
 from .. import UnexpectedModelBehavior, _utils
 from .._utils import now_utc as _now_utc
 from ..messages import (
-    ArgsJson,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -36,6 +35,7 @@ from . import (
     Model,
     StreamedResponse,
     cached_async_http_client,
+    check_allow_model_requests,
 )
 
 try:
@@ -84,6 +84,12 @@ Since [the Mistral docs](https://docs.mistral.ai/getting-started/models/models_o
 """
 
 
+class MistralModelSettings(ModelSettings):
+    """Settings used for a Mistral model request."""
+
+    # This class is a placeholder for any future mistral-specific settings
+
+
 @dataclass(init=False)
 class MistralModel(Model):
     """A model that uses Mistral.
@@ -130,6 +136,7 @@ class MistralModel(Model):
         result_tools: list[ToolDefinition],
     ) -> AgentModel:
         """Create an agent model, this is called for each step of an agent run from Pydantic AI call."""
+        check_allow_model_requests()
         return MistralAgentModel(
             self.client,
             self.model_name,
@@ -147,7 +154,7 @@ class MistralAgentModel(AgentModel):
     """Implementation of `AgentModel` for Mistral models."""
 
     client: Mistral
-    model_name: str
+    model_name: MistralModelName
     allow_text_result: bool
     function_tools: list[ToolDefinition]
     result_tools: list[ToolDefinition]
@@ -157,7 +164,7 @@ class MistralAgentModel(AgentModel):
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
     ) -> tuple[ModelResponse, Usage]:
         """Make a non-streaming request to the model from Pydantic AI call."""
-        response = await self._completions_create(messages, model_settings)
+        response = await self._completions_create(messages, cast(MistralModelSettings, model_settings or {}))
         return self._process_response(response), _map_usage(response)
 
     @asynccontextmanager
@@ -165,15 +172,14 @@ class MistralAgentModel(AgentModel):
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
     ) -> AsyncIterator[StreamedResponse]:
         """Make a streaming request to the model from Pydantic AI call."""
-        response = await self._stream_completions_create(messages, model_settings)
+        response = await self._stream_completions_create(messages, cast(MistralModelSettings, model_settings or {}))
         async with response:
             yield await self._process_streamed_response(self.result_tools, response)
 
     async def _completions_create(
-        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+        self, messages: list[ModelMessage], model_settings: MistralModelSettings
     ) -> MistralChatCompletionResponse:
         """Make a non-streaming request to the model."""
-        model_settings = model_settings or {}
         response = await self.client.chat.complete_async(
             model=str(self.model_name),
             messages=list(chain(*(self._map_message(m) for m in messages))),
@@ -185,6 +191,7 @@ class MistralAgentModel(AgentModel):
             temperature=model_settings.get('temperature', UNSET),
             top_p=model_settings.get('top_p', 1),
             timeout_ms=self._get_timeout_ms(model_settings.get('timeout')),
+            random_seed=model_settings.get('seed', UNSET),
         )
         assert response, 'A unexpected empty response from Mistral.'
         return response
@@ -192,12 +199,11 @@ class MistralAgentModel(AgentModel):
     async def _stream_completions_create(
         self,
         messages: list[ModelMessage],
-        model_settings: ModelSettings | None,
+        model_settings: MistralModelSettings,
     ) -> MistralEventStreamAsync[MistralCompletionEvent]:
         """Create a streaming completion request to the Mistral model."""
         response: MistralEventStreamAsync[MistralCompletionEvent] | None
         mistral_messages = list(chain(*(self._map_message(m) for m in messages)))
-        model_settings = model_settings or {}
 
         if self.result_tools and self.function_tools or self.function_tools:
             # Function Calling
@@ -211,6 +217,8 @@ class MistralAgentModel(AgentModel):
                 top_p=model_settings.get('top_p', 1),
                 max_tokens=model_settings.get('max_tokens', UNSET),
                 timeout_ms=self._get_timeout_ms(model_settings.get('timeout')),
+                presence_penalty=model_settings.get('presence_penalty'),
+                frequency_penalty=model_settings.get('frequency_penalty'),
             )
 
         elif self.result_tools:
@@ -265,8 +273,7 @@ class MistralAgentModel(AgentModel):
         ]
         return tools if tools else None
 
-    @staticmethod
-    def _process_response(response: MistralChatCompletionResponse) -> ModelResponse:
+    def _process_response(self, response: MistralChatCompletionResponse) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
         assert response.choices, 'Unexpected empty response choice.'
 
@@ -288,10 +295,10 @@ class MistralAgentModel(AgentModel):
                 tool = _map_mistral_to_pydantic_tool_call(tool_call=tool_call)
                 parts.append(tool)
 
-        return ModelResponse(parts, timestamp=timestamp)
+        return ModelResponse(parts, model_name=self.model_name, timestamp=timestamp)
 
-    @staticmethod
     async def _process_streamed_response(
+        self,
         result_tools: list[ToolDefinition],
         response: MistralEventStreamAsync[MistralCompletionEvent],
     ) -> StreamedResponse:
@@ -306,23 +313,21 @@ class MistralAgentModel(AgentModel):
         else:
             timestamp = datetime.now(tz=timezone.utc)
 
-        return MistralStreamedResponse(peekable_response, timestamp, {c.name: c for c in result_tools})
+        return MistralStreamedResponse(
+            _response=peekable_response,
+            _model_name=self.model_name,
+            _timestamp=timestamp,
+            _result_tools={c.name: c for c in result_tools},
+        )
 
     @staticmethod
     def _map_to_mistral_tool_call(t: ToolCallPart) -> MistralToolCall:
         """Maps a pydantic-ai ToolCall to a MistralToolCall."""
-        if isinstance(t.args, ArgsJson):
-            return MistralToolCall(
-                id=t.tool_call_id,
-                type='function',
-                function=MistralFunctionCall(name=t.tool_name, arguments=t.args.args_json),
-            )
-        else:
-            return MistralToolCall(
-                id=t.tool_call_id,
-                type='function',
-                function=MistralFunctionCall(name=t.tool_name, arguments=t.args.args_dict),
-            )
+        return MistralToolCall(
+            id=t.tool_call_id,
+            type='function',
+            function=MistralFunctionCall(name=t.tool_name, arguments=t.args),
+        )
 
     def _generate_user_output_format(self, schemas: list[dict[str, Any]]) -> MistralUserMessage:
         """Get a message with an example of the expected output format."""
@@ -505,7 +510,7 @@ class MistralStreamedResponse(StreamedResponse):
                     continue
 
                 # The following part_id will be thrown away
-                return ToolCallPart.from_raw_args(tool_name=result_tool.name, args=output_json)
+                return ToolCallPart(tool_name=result_tool.name, args=output_json)
 
     @staticmethod
     def _validate_required_json_schema(json_dict: dict[str, Any], json_schema: dict[str, Any]) -> bool:
@@ -563,7 +568,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
 
-    return ToolCallPart.from_raw_args(func_call.name, func_call.arguments, tool_call_id)
+    return ToolCallPart(func_call.name, func_call.arguments, tool_call_id)
 
 
 def _map_usage(response: MistralChatCompletionResponse | MistralCompletionChunk) -> Usage:
@@ -594,7 +599,7 @@ def _map_content(content: MistralOptionalNullable[MistralContent]) -> str | None
     elif isinstance(content, str):
         result = content
 
-    # Note: Check len to handle potential mismatch between function calls and responses from the API. (`msg: not the same number of function class and reponses`)
+    # Note: Check len to handle potential mismatch between function calls and responses from the API. (`msg: not the same number of function class and responses`)
     if result and len(result) == 0:
         result = None
 

pydantic_ai/models/openai.py

@@ -5,7 +5,7 @@ from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from itertools import chain
-from typing import Literal, Union, overload
+from typing import Literal, Union, cast, overload
 
 from httpx import AsyncClient as AsyncHTTPClient
 from typing_extensions import assert_never
@@ -48,9 +48,17 @@ except ImportError as _import_error:
 OpenAIModelName = Union[ChatModel, str]
 """
 Using this more broad type for the model name instead of the ChatModel definition
-allows this model to be used more easily with other model types (ie, Ollama)
+allows this model to be used more easily with other model types (ie, Ollama, Deepseek)
 """
 
+OpenAISystemPromptRole = Literal['system', 'developer', 'user']
+
+
+class OpenAIModelSettings(ModelSettings):
+    """Settings used for an OpenAI model request."""
+
+    # This class is a placeholder for any future openai-specific settings
+
 
 @dataclass(init=False)
 class OpenAIModel(Model):
@@ -63,6 +71,7 @@ class OpenAIModel(Model):
 
     model_name: OpenAIModelName
     client: AsyncOpenAI = field(repr=False)
+    system_prompt_role: OpenAISystemPromptRole | None = field(default=None)
 
     def __init__(
         self,
@@ -72,6 +81,7 @@ class OpenAIModel(Model):
         api_key: str | None = None,
         openai_client: AsyncOpenAI | None = None,
         http_client: AsyncHTTPClient | None = None,
+        system_prompt_role: OpenAISystemPromptRole | None = None,
     ):
         """Initialize an OpenAI model.
 
@@ -87,6 +97,8 @@ class OpenAIModel(Model):
                 [`AsyncOpenAI`](https://github.com/openai/openai-python?tab=readme-ov-file#async-usage)
                 client to use. If provided, `base_url`, `api_key`, and `http_client` must be `None`.
             http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
+            system_prompt_role: The role to use for the system prompt message. If not provided, defaults to `'system'`.
+                In the future, this may be inferred from the model name.
         """
         self.model_name: OpenAIModelName = model_name
         if openai_client is not None:
@@ -98,6 +110,7 @@ class OpenAIModel(Model):
             self.client = AsyncOpenAI(base_url=base_url, api_key=api_key, http_client=http_client)
         else:
             self.client = AsyncOpenAI(base_url=base_url, api_key=api_key, http_client=cached_async_http_client())
+        self.system_prompt_role = system_prompt_role
 
     async def agent_model(
         self,
@@ -115,6 +128,7 @@ class OpenAIModel(Model):
             self.model_name,
             allow_text_result,
             tools,
+            self.system_prompt_role,
         )
 
     def name(self) -> str:
@@ -140,35 +154,36 @@ class OpenAIAgentModel(AgentModel):
     model_name: OpenAIModelName
     allow_text_result: bool
     tools: list[chat.ChatCompletionToolParam]
+    system_prompt_role: OpenAISystemPromptRole | None
 
     async def request(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
     ) -> tuple[ModelResponse, usage.Usage]:
-        response = await self._completions_create(messages, False, model_settings)
+        response = await self._completions_create(messages, False, cast(OpenAIModelSettings, model_settings or {}))
         return self._process_response(response), _map_usage(response)
 
     @asynccontextmanager
     async def request_stream(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
     ) -> AsyncIterator[StreamedResponse]:
-        response = await self._completions_create(messages, True, model_settings)
+        response = await self._completions_create(messages, True, cast(OpenAIModelSettings, model_settings or {}))
         async with response:
             yield await self._process_streamed_response(response)
 
     @overload
     async def _completions_create(
-        self, messages: list[ModelMessage], stream: Literal[True], model_settings: ModelSettings | None
+        self, messages: list[ModelMessage], stream: Literal[True], model_settings: OpenAIModelSettings
     ) -> AsyncStream[ChatCompletionChunk]:
         pass
 
     @overload
     async def _completions_create(
-        self, messages: list[ModelMessage], stream: Literal[False], model_settings: ModelSettings | None
+        self, messages: list[ModelMessage], stream: Literal[False], model_settings: OpenAIModelSettings
     ) -> chat.ChatCompletion:
         pass
 
     async def _completions_create(
-        self, messages: list[ModelMessage], stream: bool, model_settings: ModelSettings | None
+        self, messages: list[ModelMessage], stream: bool, model_settings: OpenAIModelSettings
     ) -> chat.ChatCompletion | AsyncStream[ChatCompletionChunk]:
         # standalone function to make it easier to override
         if not self.tools:
@@ -180,13 +195,11 @@ class OpenAIAgentModel(AgentModel):
 
         openai_messages = list(chain(*(self._map_message(m) for m in messages)))
 
-        model_settings = model_settings or {}
-
         return await self.client.chat.completions.create(
             model=self.model_name,
             messages=openai_messages,
             n=1,
-            parallel_tool_calls=True if self.tools else NOT_GIVEN,
+            parallel_tool_calls=model_settings.get('parallel_tool_calls', NOT_GIVEN),
             tools=self.tools or NOT_GIVEN,
             tool_choice=tool_choice or NOT_GIVEN,
             stream=stream,
@@ -195,10 +208,13 @@ class OpenAIAgentModel(AgentModel):
             temperature=model_settings.get('temperature', NOT_GIVEN),
             top_p=model_settings.get('top_p', NOT_GIVEN),
             timeout=model_settings.get('timeout', NOT_GIVEN),
+            seed=model_settings.get('seed', NOT_GIVEN),
+            presence_penalty=model_settings.get('presence_penalty', NOT_GIVEN),
+            frequency_penalty=model_settings.get('frequency_penalty', NOT_GIVEN),
+            logit_bias=model_settings.get('logit_bias', NOT_GIVEN),
         )
 
-    @staticmethod
-    def _process_response(response: chat.ChatCompletion) -> ModelResponse:
+    def _process_response(self, response: chat.ChatCompletion) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
         timestamp = datetime.fromtimestamp(response.created, tz=timezone.utc)
         choice = response.choices[0]
@@ -207,24 +223,26 @@ class OpenAIAgentModel(AgentModel):
             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_raw_args(c.function.name, c.function.arguments, c.id))
-        return ModelResponse(items, timestamp=timestamp)
+                items.append(ToolCallPart(c.function.name, c.function.arguments, c.id))
+        return ModelResponse(items, model_name=self.model_name, timestamp=timestamp)
 
-    @staticmethod
-    async def _process_streamed_response(response: AsyncStream[ChatCompletionChunk]) -> OpenAIStreamedResponse:
+    async def _process_streamed_response(self, response: AsyncStream[ChatCompletionChunk]) -> OpenAIStreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
         peekable_response = _utils.PeekableAsyncStream(response)
         first_chunk = await peekable_response.peek()
         if isinstance(first_chunk, _utils.Unset):
             raise UnexpectedModelBehavior('Streamed response ended without content or tool calls')
 
-        return OpenAIStreamedResponse(peekable_response, datetime.fromtimestamp(first_chunk.created, tz=timezone.utc))
+        return OpenAIStreamedResponse(
+            _model_name=self.model_name,
+            _response=peekable_response,
+            _timestamp=datetime.fromtimestamp(first_chunk.created, tz=timezone.utc),
+        )
 
-    @classmethod
-    def _map_message(cls, message: ModelMessage) -> Iterable[chat.ChatCompletionMessageParam]:
+    def _map_message(self, message: ModelMessage) -> Iterable[chat.ChatCompletionMessageParam]:
         """Just maps a `pydantic_ai.Message` to a `openai.types.ChatCompletionMessageParam`."""
         if isinstance(message, ModelRequest):
-            yield from cls._map_user_message(message)
+            yield from self._map_user_message(message)
         elif isinstance(message, ModelResponse):
             texts: list[str] = []
             tool_calls: list[chat.ChatCompletionMessageToolCallParam] = []
@@ -246,11 +264,15 @@ class OpenAIAgentModel(AgentModel):
         else:
             assert_never(message)
 
-    @classmethod
-    def _map_user_message(cls, message: ModelRequest) -> Iterable[chat.ChatCompletionMessageParam]:
+    def _map_user_message(self, message: ModelRequest) -> Iterable[chat.ChatCompletionMessageParam]:
         for part in message.parts:
             if isinstance(part, SystemPromptPart):
-                yield chat.ChatCompletionSystemMessageParam(role='system', content=part.content)
+                if self.system_prompt_role == 'developer':
+                    yield chat.ChatCompletionDeveloperMessageParam(role='developer', content=part.content)
+                elif self.system_prompt_role == 'user':
+                    yield chat.ChatCompletionUserMessageParam(role='user', content=part.content)
+                else:
+                    yield chat.ChatCompletionSystemMessageParam(role='system', content=part.content)
             elif isinstance(part, UserPromptPart):
                 yield chat.ChatCompletionUserMessageParam(role='user', content=part.content)
             elif isinstance(part, ToolReturnPart):

pydantic_ai/models/test.py

@@ -12,7 +12,6 @@ import pydantic_core
 
 from .. import _utils
 from ..messages import (
-    ArgsJson,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -34,6 +33,20 @@ from . import (
 from .function import _estimate_string_tokens, _estimate_usage  # pyright: ignore[reportPrivateUsage]
 
 
+@dataclass
+class _TextResult:
+    """A private wrapper class to tag a result that came from the custom_result_text field."""
+
+    value: str | None
+
+
+@dataclass
+class _FunctionToolResult:
+    """A wrapper class to tag a result that came from the custom_result_args field."""
+
+    value: Any | None
+
+
 @dataclass
 class TestModel(Model):
     """A model specifically for testing purposes.
@@ -53,7 +66,7 @@ class TestModel(Model):
     call_tools: list[str] | Literal['all'] = 'all'
     """List of tools to call. If `'all'`, all tools will be called."""
     custom_result_text: str | None = None
-    """If set, this text is return as the final result."""
+    """If set, this text is returned as the final result."""
     custom_result_args: Any | None = None
     """If set, these args will be passed to the result tool."""
     seed: int = 0
@@ -95,21 +108,21 @@ class TestModel(Model):
         if self.custom_result_text is not None:
             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)
+            result: _TextResult | _FunctionToolResult = _TextResult(self.custom_result_text)
         elif self.custom_result_args is not None:
             assert result_tools is not None, 'No result tools provided, but `custom_result_args` is set.'
             result_tool = result_tools[0]
 
             if k := result_tool.outer_typed_dict_key:
-                result = _utils.Either(right={k: self.custom_result_args})
+                result = _FunctionToolResult({k: self.custom_result_args})
             else:
-                result = _utils.Either(right=self.custom_result_args)
+                result = _FunctionToolResult(self.custom_result_args)
         elif allow_text_result:
-            result = _utils.Either(left=None)
+            result = _TextResult(None)
         elif result_tools:
-            result = _utils.Either(right=None)
+            result = _FunctionToolResult(None)
         else:
-            result = _utils.Either(left=None)
+            result = _TextResult(None)
 
         return TestAgentModel(tool_calls, result, result_tools, self.seed)
 
@@ -126,9 +139,10 @@ class TestAgentModel(AgentModel):
 
     tool_calls: list[tuple[str, ToolDefinition]]
     # left means the text is plain text; right means it's a function call
-    result: _utils.Either[str | None, Any | None]
+    result: _TextResult | _FunctionToolResult
     result_tools: list[ToolDefinition]
     seed: int
+    model_name: str = 'test'
 
     async def request(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
@@ -142,7 +156,7 @@ class TestAgentModel(AgentModel):
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
     ) -> AsyncIterator[StreamedResponse]:
         model_response = self._request(messages, model_settings)
-        yield TestStreamedResponse(model_response, messages)
+        yield TestStreamedResponse(_model_name=self.model_name, _structured_response=model_response, _messages=messages)
 
     def gen_tool_args(self, tool_def: ToolDefinition) -> Any:
         return _JsonSchemaTestData(tool_def.parameters_json_schema, self.seed).generate()
@@ -151,7 +165,8 @@ class TestAgentModel(AgentModel):
         # 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_raw_args(name, self.gen_tool_args(args)) for name, args in self.tool_calls]
+                parts=[ToolCallPart(name, self.gen_tool_args(args)) for name, args in self.tool_calls],
+                model_name=self.model_name,
             )
 
         if messages:
@@ -164,7 +179,7 @@ class TestAgentModel(AgentModel):
                 # Handle retries for both function tools and result tools
                 # Check function tools first
                 retry_parts: list[ModelResponsePart] = [
-                    ToolCallPart.from_raw_args(name, self.gen_tool_args(args))
+                    ToolCallPart(name, self.gen_tool_args(args))
                     for name, args in self.tool_calls
                     if name in new_retry_names
                 ]
@@ -172,15 +187,20 @@ class TestAgentModel(AgentModel):
                 if self.result_tools:
                     retry_parts.extend(
                         [
-                            ToolCallPart.from_raw_args(tool.name, self.gen_tool_args(tool))
+                            ToolCallPart(
+                                tool.name,
+                                self.result.value
+                                if isinstance(self.result, _FunctionToolResult) and self.result.value is not None
+                                else self.gen_tool_args(tool),
+                            )
                             for tool in self.result_tools
                             if tool.name in new_retry_names
                         ]
                     )
-                return ModelResponse(parts=retry_parts)
+                return ModelResponse(parts=retry_parts, model_name=self.model_name)
 
-        if response_text := self.result.left:
-            if response_text.value is None:
+        if isinstance(self.result, _TextResult):
+            if (response_text := self.result.value) is None:
                 # build up details of tool responses
                 output: dict[str, Any] = {}
                 for message in messages:
@@ -189,20 +209,24 @@ class TestAgentModel(AgentModel):
                             if isinstance(part, ToolReturnPart):
                                 output[part.tool_name] = part.content
                 if output:
-                    return ModelResponse.from_text(pydantic_core.to_json(output).decode())
+                    return ModelResponse(
+                        parts=[TextPart(pydantic_core.to_json(output).decode())], model_name=self.model_name
+                    )
                 else:
-                    return ModelResponse.from_text('success (no tool calls)')
+                    return ModelResponse(parts=[TextPart('success (no tool calls)')], model_name=self.model_name)
             else:
-                return ModelResponse.from_text(response_text.value)
+                return ModelResponse(parts=[TextPart(response_text)], model_name=self.model_name)
         else:
             assert self.result_tools, 'No result tools provided'
-            custom_result_args = self.result.right
+            custom_result_args = self.result.value
             result_tool = self.result_tools[self.seed % len(self.result_tools)]
             if custom_result_args is not None:
-                return ModelResponse(parts=[ToolCallPart.from_raw_args(result_tool.name, custom_result_args)])
+                return ModelResponse(
+                    parts=[ToolCallPart(result_tool.name, custom_result_args)], model_name=self.model_name
+                )
             else:
                 response_args = self.gen_tool_args(result_tool)
-                return ModelResponse(parts=[ToolCallPart.from_raw_args(result_tool.name, response_args)])
+                return ModelResponse(parts=[ToolCallPart(result_tool.name, response_args)], model_name=self.model_name)
 
 
 @dataclass
@@ -233,9 +257,8 @@ class TestStreamedResponse(StreamedResponse):
                     self._usage += _get_string_usage(word)
                     yield self._parts_manager.handle_text_delta(vendor_part_id=i, content=word)
             else:
-                args = part.args.args_json if isinstance(part.args, ArgsJson) else part.args.args_dict
                 yield self._parts_manager.handle_tool_call_part(
-                    vendor_part_id=i, tool_name=part.tool_name, args=args, tool_call_id=part.tool_call_id
+                    vendor_part_id=i, tool_name=part.tool_name, args=part.args, tool_call_id=part.tool_call_id
                 )
 
     def timestamp(self) -> datetime:

pydantic_ai/models/vertexai.py

@@ -10,7 +10,7 @@ from httpx import AsyncClient as AsyncHTTPClient
 from .._utils import run_in_executor
 from ..exceptions import UserError
 from ..tools import ToolDefinition
-from . import Model, cached_async_http_client
+from . import Model, cached_async_http_client, check_allow_model_requests
 from .gemini import GeminiAgentModel, GeminiModelName
 
 try:
@@ -114,6 +114,7 @@ class VertexAIModel(Model):
         allow_text_result: bool,
         result_tools: list[ToolDefinition],
     ) -> GeminiAgentModel:
+        check_allow_model_requests()
         url, auth = await self.ainit()
         return GeminiAgentModel(
             http_client=self.http_client,