pydantic-ai-slim 0.6.2__py3-none-any.whl → 0.7.0__py3-none-any.whl

pydantic_ai/models/openai.py

@@ -11,13 +11,13 @@ from typing import Any, Literal, Union, cast, overload
 from pydantic import ValidationError
 from typing_extensions import assert_never
 
-from pydantic_ai.exceptions import UserError
-
 from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage
 from .._output import DEFAULT_OUTPUT_TOOL_NAME, OutputObjectDefinition
+from .._run_context import RunContext
 from .._thinking_part import split_content_into_text_and_thinking
 from .._utils import guard_tool_call_id as _guard_tool_call_id, now_utc as _now_utc, number_to_datetime
 from ..builtin_tools import CodeExecutionTool, WebSearchTool
+from ..exceptions import UserError
 from ..messages import (
     AudioUrl,
     BinaryContent,
@@ -256,13 +256,14 @@ class OpenAIModel(Model):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
         check_allow_model_requests()
         response = await self._completions_create(
             messages, True, cast(OpenAIModelSettings, model_settings or {}), model_request_parameters
         )
         async with response:
-            yield await self._process_streamed_response(response)
+            yield await self._process_streamed_response(response, model_request_parameters)
 
     @property
     def model_name(self) -> OpenAIModelName:
@@ -427,7 +428,9 @@ class OpenAIModel(Model):
             vendor_id=response.id,
         )
 
-    async def _process_streamed_response(self, response: AsyncStream[ChatCompletionChunk]) -> OpenAIStreamedResponse:
+    async def _process_streamed_response(
+        self, response: AsyncStream[ChatCompletionChunk], model_request_parameters: ModelRequestParameters
+    ) -> OpenAIStreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
         peekable_response = _utils.PeekableAsyncStream(response)
         first_chunk = await peekable_response.peek()
@@ -437,6 +440,7 @@ class OpenAIModel(Model):
             )
 
         return OpenAIStreamedResponse(
+            model_request_parameters=model_request_parameters,
             _model_name=self._model_name,
             _model_profile=self.profile,
             _response=peekable_response,
@@ -444,10 +448,7 @@ class OpenAIModel(Model):
         )
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[chat.ChatCompletionToolParam]:
-        tools = [self._map_tool_definition(r) for r in model_request_parameters.function_tools]
-        if model_request_parameters.output_tools:
-            tools += [self._map_tool_definition(r) for r in model_request_parameters.output_tools]
-        return tools
+        return [self._map_tool_definition(r) for r in model_request_parameters.tool_defs.values()]
 
     def _get_web_search_options(self, model_request_parameters: ModelRequestParameters) -> WebSearchOptions | None:
         for tool in model_request_parameters.builtin_tools:
@@ -461,8 +462,10 @@ class OpenAIModel(Model):
                         ),
                     )
                 return WebSearchOptions(search_context_size=tool.search_context_size)
-            elif isinstance(tool, CodeExecutionTool):  # pragma: no branch
-                raise UserError('`CodeExecutionTool` is not supported by OpenAI')
+            else:
+                raise UserError(
+                    f'`{tool.__class__.__name__}` is not supported by `OpenAIModel`. If it should be, please file an issue.'
+                )
 
     async def _map_messages(self, messages: list[ModelMessage]) -> list[chat.ChatCompletionMessageParam]:
         """Just maps a `pydantic_ai.Message` to a `openai.types.ChatCompletionMessageParam`."""
@@ -631,14 +634,6 @@ class OpenAIResponsesModel(Model):
     The [OpenAI Responses API](https://platform.openai.com/docs/api-reference/responses) is the
     new API for OpenAI models.
 
-    The Responses API has built-in tools, that you can use instead of building your own:
-
-    - [Web search](https://platform.openai.com/docs/guides/tools-web-search)
-    - [File search](https://platform.openai.com/docs/guides/tools-file-search)
-    - [Computer use](https://platform.openai.com/docs/guides/tools-computer-use)
-
-    Use the `openai_builtin_tools` setting to add these tools to your model.
-
     If you are interested in the differences between the Responses API and the Chat Completions API,
     see the [OpenAI API docs](https://platform.openai.com/docs/guides/responses-vs-chat-completions).
     """
@@ -702,13 +697,14 @@ class OpenAIResponsesModel(Model):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
         check_allow_model_requests()
         response = await self._responses_create(
             messages, True, cast(OpenAIResponsesModelSettings, model_settings or {}), model_request_parameters
         )
         async with response:
-            yield await self._process_streamed_response(response)
+            yield await self._process_streamed_response(response, model_request_parameters)
 
     def _process_response(self, response: responses.Response) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
@@ -735,7 +731,9 @@ class OpenAIResponsesModel(Model):
         )
 
     async def _process_streamed_response(
-        self, response: AsyncStream[responses.ResponseStreamEvent]
+        self,
+        response: AsyncStream[responses.ResponseStreamEvent],
+        model_request_parameters: ModelRequestParameters,
     ) -> OpenAIResponsesStreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
         peekable_response = _utils.PeekableAsyncStream(response)
@@ -745,6 +743,7 @@ class OpenAIResponsesModel(Model):
 
         assert isinstance(first_chunk, responses.ResponseCreatedEvent)
         return OpenAIResponsesStreamedResponse(
+            model_request_parameters=model_request_parameters,
             _model_name=self._model_name,
             _response=peekable_response,
             _timestamp=number_to_datetime(first_chunk.response.created_at),
@@ -775,8 +774,11 @@ class OpenAIResponsesModel(Model):
         model_settings: OpenAIResponsesModelSettings,
         model_request_parameters: ModelRequestParameters,
     ) -> responses.Response | AsyncStream[responses.ResponseStreamEvent]:
-        tools = self._get_tools(model_request_parameters)
-        tools = self._get_builtin_tools(model_request_parameters) + tools
+        tools = (
+            self._get_builtin_tools(model_request_parameters)
+            + list(model_settings.get('openai_builtin_tools', []))
+            + self._get_tools(model_request_parameters)
+        )
 
         if not tools:
             tool_choice: Literal['none', 'required', 'auto'] | None = None
@@ -859,10 +861,7 @@ class OpenAIResponsesModel(Model):
         return Reasoning(effort=reasoning_effort, summary=reasoning_summary)
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[responses.FunctionToolParam]:
-        tools = [self._map_tool_definition(r) for r in model_request_parameters.function_tools]
-        if model_request_parameters.output_tools:
-            tools += [self._map_tool_definition(r) for r in model_request_parameters.output_tools]
-        return tools
+        return [self._map_tool_definition(r) for r in model_request_parameters.tool_defs.values()]
 
     def _get_builtin_tools(self, model_request_parameters: ModelRequestParameters) -> list[responses.ToolParam]:
         tools: list[responses.ToolParam] = []
@@ -878,6 +877,10 @@ class OpenAIResponsesModel(Model):
                 tools.append(web_search_tool)
             elif isinstance(tool, CodeExecutionTool):  # pragma: no branch
                 tools.append({'type': 'code_interpreter', 'container': {'type': 'auto'}})
+            else:
+                raise UserError(  # pragma: no cover
+                    f'`{tool.__class__.__name__}` is not supported by `OpenAIResponsesModel`. If it should be, please file an issue.'
+                )
         return tools
 
     def _map_tool_definition(self, f: ToolDefinition) -> responses.FunctionToolParam:

pydantic_ai/models/test.py

@@ -12,6 +12,7 @@ import pydantic_core
 from typing_extensions import assert_never
 
 from .. import _utils
+from .._run_context import RunContext
 from ..exceptions import UserError
 from ..messages import (
     BuiltinToolCallPart,
@@ -121,11 +122,13 @@ class TestModel(Model):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
         self.last_model_request_parameters = model_request_parameters
 
         model_response = self._request(messages, model_settings, model_request_parameters)
         yield TestStreamedResponse(
+            model_request_parameters=model_request_parameters,
             _model_name=self._model_name,
             _structured_response=model_response,
             _messages=messages,

pydantic_ai/models/wrapper.py

@@ -6,6 +6,7 @@ from dataclasses import dataclass
 from functools import cached_property
 from typing import Any
 
+from .._run_context import RunContext
 from ..messages import ModelMessage, ModelResponse
 from ..profiles import ModelProfile
 from ..settings import ModelSettings
@@ -35,8 +36,11 @@ class WrapperModel(Model):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
-        async with self.wrapped.request_stream(messages, model_settings, model_request_parameters) as response_stream:
+        async with self.wrapped.request_stream(
+            messages, model_settings, model_request_parameters, run_context
+        ) as response_stream:
             yield response_stream
 
     def customize_request_parameters(self, model_request_parameters: ModelRequestParameters) -> ModelRequestParameters:
@@ -60,4 +64,4 @@ class WrapperModel(Model):
         return self.wrapped.settings
 
     def __getattr__(self, item: str):
-        return getattr(self.wrapped, item)  # pragma: no cover
+        return getattr(self.wrapped, item)

pydantic_ai/profiles/openai.py

@@ -32,7 +32,7 @@ class OpenAIModelProfile(ModelProfile):
 
 def openai_model_profile(model_name: str) -> ModelProfile:
     """Get the model profile for an OpenAI model."""
-    is_reasoning_model = model_name.startswith('o')
+    is_reasoning_model = model_name.startswith('o') or model_name.startswith('gpt-5')
     # Structured Outputs (output mode 'native') is only supported with the gpt-4o-mini, gpt-4o-mini-2024-07-18, and gpt-4o-2024-08-06 model snapshots and later.
     # We leave it in here for all models because the `default_structured_output_mode` is `'tool'`, so `native` is only used
     # when the user specifically uses the `NativeOutput` marker, so an error from the API is acceptable.

pydantic_ai/providers/google.py

@@ -12,8 +12,8 @@ from pydantic_ai.profiles.google import google_model_profile
 from pydantic_ai.providers import Provider
 
 try:
-    from google import genai
     from google.auth.credentials import Credentials
+    from google.genai import Client
     from google.genai.types import HttpOptionsDict
 except ImportError as _import_error:
     raise ImportError(
@@ -22,7 +22,7 @@ except ImportError as _import_error:
     ) from _import_error
 
 
-class GoogleProvider(Provider[genai.Client]):
+class GoogleProvider(Provider[Client]):
     """Provider for Google."""
 
     @property
@@ -34,7 +34,7 @@ class GoogleProvider(Provider[genai.Client]):
         return str(self._client._api_client._http_options.base_url)  # type: ignore[reportPrivateUsage]
 
     @property
-    def client(self) -> genai.Client:
+    def client(self) -> Client:
         return self._client
 
     def model_profile(self, model_name: str) -> ModelProfile | None:
@@ -53,7 +53,7 @@ class GoogleProvider(Provider[genai.Client]):
     ) -> None: ...
 
     @overload
-    def __init__(self, *, client: genai.Client) -> None: ...
+    def __init__(self, *, client: Client) -> None: ...
 
     @overload
     def __init__(self, *, vertexai: bool = False) -> None: ...
@@ -65,7 +65,7 @@ class GoogleProvider(Provider[genai.Client]):
         credentials: Credentials | None = None,
         project: str | None = None,
         location: VertexAILocation | Literal['global'] | None = None,
-        client: genai.Client | None = None,
+        client: Client | None = None,
         vertexai: bool | None = None,
     ) -> None:
         """Create a new Google provider.
@@ -102,9 +102,9 @@ class GoogleProvider(Provider[genai.Client]):
                         'Set the `GOOGLE_API_KEY` environment variable or pass it via `GoogleProvider(api_key=...)`'
                         'to use the Google Generative Language API.'
                     )
-                self._client = genai.Client(vertexai=vertexai, api_key=api_key, http_options=http_options)
+                self._client = Client(vertexai=vertexai, api_key=api_key, http_options=http_options)
             else:
-                self._client = genai.Client(
+                self._client = Client(
                     vertexai=vertexai,
                     project=project or os.environ.get('GOOGLE_CLOUD_PROJECT'),
                     # From https://github.com/pydantic/pydantic-ai/pull/2031/files#r2169682149:

pydantic_ai/result.py

@@ -1,6 +1,6 @@
 from __future__ import annotations as _annotations
 
-from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable
+from collections.abc import AsyncIterator, Awaitable, Callable
 from copy import copy
 from dataclasses import dataclass, field
 from datetime import datetime
@@ -22,7 +22,7 @@ from ._output import (
     ToolOutputSchema,
 )
 from ._run_context import AgentDepsT, RunContext
-from .messages import AgentStreamEvent, FinalResultEvent
+from .messages import AgentStreamEvent
 from .output import (
     OutputDataT,
     ToolOutput,
@@ -45,13 +45,13 @@ T = TypeVar('T')
 class AgentStream(Generic[AgentDepsT, OutputDataT]):
     _raw_stream_response: models.StreamedResponse
     _output_schema: OutputSchema[OutputDataT]
+    _model_request_parameters: models.ModelRequestParameters
     _output_validators: list[OutputValidator[AgentDepsT, OutputDataT]]
     _run_ctx: RunContext[AgentDepsT]
     _usage_limits: UsageLimits | None
     _tool_manager: ToolManager[AgentDepsT]
 
     _agent_stream_iterator: AsyncIterator[AgentStreamEvent] | None = field(default=None, init=False)
-    _final_result_event: FinalResultEvent | None = field(default=None, init=False)
     _initial_run_ctx_usage: Usage = field(init=False)
 
     def __post_init__(self):
@@ -60,12 +60,12 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
     async def stream_output(self, *, debounce_by: float | None = 0.1) -> AsyncIterator[OutputDataT]:
         """Asynchronously stream the (validated) agent outputs."""
         async for response in self.stream_responses(debounce_by=debounce_by):
-            if self._final_result_event is not None:
+            if self._raw_stream_response.final_result_event is not None:
                 try:
                     yield await self._validate_response(response, allow_partial=True)
                 except ValidationError:
                     pass
-        if self._final_result_event is not None:  # pragma: no branch
+        if self._raw_stream_response.final_result_event is not None:  # pragma: no branch
             yield await self._validate_response(self._raw_stream_response.get())
 
     async def stream_responses(self, *, debounce_by: float | None = 0.1) -> AsyncIterator[_messages.ModelResponse]:
@@ -131,10 +131,11 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
 
     async def _validate_response(self, message: _messages.ModelResponse, *, allow_partial: bool = False) -> OutputDataT:
         """Validate a structured result message."""
-        if self._final_result_event is None:
+        final_result_event = self._raw_stream_response.final_result_event
+        if final_result_event is None:
             raise exceptions.UnexpectedModelBehavior('Invalid response, unable to find output')  # pragma: no cover
 
-        output_tool_name = self._final_result_event.tool_name
+        output_tool_name = final_result_event.tool_name
 
         if isinstance(self._output_schema, ToolOutputSchema) and output_tool_name is not None:
             tool_call = next(
@@ -195,7 +196,7 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
                     and isinstance(event.part, _messages.TextPart)
                     and event.part.content
                 ):
-                    yield event.part.content, event.index  # pragma: no cover
+                    yield event.part.content, event.index
                 elif (  # pragma: no branch
                     isinstance(event, _messages.PartDeltaEvent)
                     and isinstance(event.delta, _messages.TextPartDelta)
@@ -221,52 +222,12 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
                 yield ''.join(deltas)
 
     def __aiter__(self) -> AsyncIterator[AgentStreamEvent]:
-        """Stream [`AgentStreamEvent`][pydantic_ai.messages.AgentStreamEvent]s.
-
-        This proxies the _raw_stream_response and sends all events to the agent stream, while also checking for matches
-        on the result schema and emitting a [`FinalResultEvent`][pydantic_ai.messages.FinalResultEvent] if/when the
-        first match is found.
-        """
-        if self._agent_stream_iterator is not None:
-            return self._agent_stream_iterator
-
-        async def aiter():
-            output_schema = self._output_schema
-
-            def _get_final_result_event(e: _messages.ModelResponseStreamEvent) -> _messages.FinalResultEvent | None:
-                """Return an appropriate FinalResultEvent if `e` corresponds to a part that will produce a final result."""
-                if isinstance(e, _messages.PartStartEvent):
-                    new_part = e.part
-                    if isinstance(new_part, _messages.TextPart) and isinstance(
-                        output_schema, TextOutputSchema
-                    ):  # pragma: no branch
-                        return _messages.FinalResultEvent(tool_name=None, tool_call_id=None)
-                    elif isinstance(new_part, _messages.ToolCallPart) and (
-                        tool_def := self._tool_manager.get_tool_def(new_part.tool_name)
-                    ):
-                        if tool_def.kind == 'output':
-                            return _messages.FinalResultEvent(
-                                tool_name=new_part.tool_name, tool_call_id=new_part.tool_call_id
-                            )
-                        elif tool_def.kind == 'deferred':
-                            return _messages.FinalResultEvent(tool_name=None, tool_call_id=None)
-
-            usage_checking_stream = _get_usage_checking_stream_response(
+        """Stream [`AgentStreamEvent`][pydantic_ai.messages.AgentStreamEvent]s."""
+        if self._agent_stream_iterator is None:
+            self._agent_stream_iterator = _get_usage_checking_stream_response(
                 self._raw_stream_response, self._usage_limits, self.usage
             )
-            async for event in usage_checking_stream:
-                yield event
-                if (final_result_event := _get_final_result_event(event)) is not None:
-                    self._final_result_event = final_result_event
-                    yield final_result_event
-                    break
-
-            # If we broke out of the above loop, we need to yield the rest of the events
-            # If we didn't, this will just be a no-op
-            async for event in usage_checking_stream:
-                yield event
-
-        self._agent_stream_iterator = aiter()
+
         return self._agent_stream_iterator
 
 
@@ -462,10 +423,10 @@ class FinalResult(Generic[OutputDataT]):
 
 
 def _get_usage_checking_stream_response(
-    stream_response: AsyncIterable[_messages.ModelResponseStreamEvent],
+    stream_response: models.StreamedResponse,
     limits: UsageLimits | None,
     get_usage: Callable[[], Usage],
-) -> AsyncIterable[_messages.ModelResponseStreamEvent]:
+) -> AsyncIterator[AgentStreamEvent]:
     if limits is not None and limits.has_token_limits():
 
         async def _usage_checking_iterator():
@@ -475,4 +436,9 @@ def _get_usage_checking_stream_response(
 
         return _usage_checking_iterator()
     else:
-        return stream_response
+        # TODO: Use `return aiter(stream_response)` once we drop support for Python 3.9
+        async def _iterator():
+            async for item in stream_response:
+                yield item
+
+        return _iterator()