pydantic-ai-slim 1.0.2__py3-none-any.whl → 1.0.4__py3-none-any.whl

pydantic_ai/models/openai.py

@@ -4,7 +4,7 @@ import base64
 import warnings
 from collections.abc import AsyncIterable, AsyncIterator, Sequence
 from contextlib import asynccontextmanager
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from datetime import datetime
 from typing import Any, Literal, cast, overload
 
@@ -31,6 +31,7 @@ from ..messages import (
     ModelResponse,
     ModelResponsePart,
     ModelResponseStreamEvent,
+    PartStartEvent,
     RetryPromptPart,
     SystemPromptPart,
     TextPart,
@@ -73,6 +74,7 @@ try:
     )
     from openai.types.responses import ComputerToolParam, FileSearchToolParam, WebSearchToolParam
     from openai.types.responses.response_input_param import FunctionCallOutput, Message
+    from openai.types.responses.response_reasoning_item_param import Summary
     from openai.types.responses.response_status import ResponseStatus
     from openai.types.shared import ReasoningEffort
     from openai.types.shared_params import Reasoning
@@ -491,9 +493,17 @@ class OpenAIChatModel(Model):
 
         choice = response.choices[0]
         items: list[ModelResponsePart] = []
-        # The `reasoning_content` is only present in DeepSeek models.
+        # The `reasoning_content` field is only present in DeepSeek models.
+        # https://api-docs.deepseek.com/guides/reasoning_model
         if reasoning_content := getattr(choice.message, 'reasoning_content', None):
-            items.append(ThinkingPart(content=reasoning_content))
+            items.append(ThinkingPart(id='reasoning_content', content=reasoning_content, provider_name=self.system))
+
+        # NOTE: We don't currently handle OpenRouter `reasoning_details`:
+        # - https://openrouter.ai/docs/use-cases/reasoning-tokens#preserving-reasoning-blocks
+        # NOTE: We don't currently handle OpenRouter/gpt-oss `reasoning`:
+        # - https://cookbook.openai.com/articles/gpt-oss/handle-raw-cot#chat-completions-api
+        # - https://openrouter.ai/docs/use-cases/reasoning-tokens#basic-usage-with-reasoning-tokens
+        # If you need this, please file an issue.
 
         vendor_details: dict[str, Any] = {}
 
@@ -513,7 +523,10 @@ class OpenAIChatModel(Model):
             ]
 
         if choice.message.content is not None:
-            items.extend(split_content_into_text_and_thinking(choice.message.content, self.profile.thinking_tags))
+            items.extend(
+                (replace(part, id='content', provider_name=self.system) if isinstance(part, ThinkingPart) else part)
+                for part in split_content_into_text_and_thinking(choice.message.content, self.profile.thinking_tags)
+            )
         if choice.message.tool_calls is not None:
             for c in choice.message.tool_calls:
                 if isinstance(c, ChatCompletionMessageFunctionToolCall):
@@ -527,10 +540,9 @@ class OpenAIChatModel(Model):
                 part.tool_call_id = _guard_tool_call_id(part)
                 items.append(part)
 
-        finish_reason: FinishReason | None = None
-        if raw_finish_reason := choice.finish_reason:  # pragma: no branch
-            vendor_details['finish_reason'] = raw_finish_reason
-            finish_reason = _CHAT_FINISH_REASON_MAP.get(raw_finish_reason)
+        raw_finish_reason = choice.finish_reason
+        vendor_details['finish_reason'] = raw_finish_reason
+        finish_reason = _CHAT_FINISH_REASON_MAP.get(raw_finish_reason)
 
         return ModelResponse(
             parts=items,
@@ -556,7 +568,7 @@ class OpenAIChatModel(Model):
 
         return OpenAIStreamedResponse(
             model_request_parameters=model_request_parameters,
-            _model_name=self._model_name,
+            _model_name=first_chunk.model,
             _model_profile=self.profile,
             _response=peekable_response,
             _timestamp=number_to_datetime(first_chunk.created),
@@ -569,6 +581,12 @@ class OpenAIChatModel(Model):
     def _get_web_search_options(self, model_request_parameters: ModelRequestParameters) -> WebSearchOptions | None:
         for tool in model_request_parameters.builtin_tools:
             if isinstance(tool, WebSearchTool):  # pragma: no branch
+                if not OpenAIModelProfile.from_profile(self.profile).openai_chat_supports_web_search:
+                    raise UserError(
+                        f'WebSearchTool is not supported with `OpenAIChatModel` and model {self.model_name!r}. '
+                        f'Please use `OpenAIResponsesModel` instead.'
+                    )
+
                 if tool.user_location:
                     return WebSearchOptions(
                         search_context_size=tool.search_context_size,
@@ -580,7 +598,7 @@ class OpenAIChatModel(Model):
                 return WebSearchOptions(search_context_size=tool.search_context_size)
             else:
                 raise UserError(
-                    f'`{tool.__class__.__name__}` is not supported by `OpenAIModel`. If it should be, please file an issue.'
+                    f'`{tool.__class__.__name__}` is not supported by `OpenAIChatModel`. If it should be, please file an issue.'
                 )
 
     async def _map_messages(self, messages: list[ModelMessage]) -> list[chat.ChatCompletionMessageParam]:
@@ -597,10 +615,11 @@ class OpenAIChatModel(Model):
                     if isinstance(item, TextPart):
                         texts.append(item.content)
                     elif isinstance(item, ThinkingPart):
-                        # NOTE: We don't send ThinkingPart to the providers yet. If you are unsatisfied with this,
-                        # please open an issue. The below code is the code to send thinking to the provider.
-                        # texts.append(f'<think>\n{item.content}\n</think>')
-                        pass
+                        # NOTE: DeepSeek `reasoning_content` field should NOT be sent back per https://api-docs.deepseek.com/guides/reasoning_model,
+                        # but we currently just send it in `<think>` tags anyway as we don't want DeepSeek-specific checks here.
+                        # If you need this changed, please file an issue.
+                        start_tag, end_tag = self.profile.thinking_tags
+                        texts.append('\n'.join([start_tag, item.content, end_tag]))
                     elif isinstance(item, ToolCallPart):
                         tool_calls.append(self._map_tool_call(item))
                     # OpenAI doesn't return built-in tool calls
@@ -838,17 +857,30 @@ class OpenAIResponsesModel(Model):
         timestamp = number_to_datetime(response.created_at)
         items: list[ModelResponsePart] = []
         for item in response.output:
-            if item.type == 'reasoning':
+            if isinstance(item, responses.ResponseReasoningItem):
+                signature = item.encrypted_content
                 for summary in item.summary:
-                    # NOTE: We use the same id for all summaries because we can merge them on the round trip.
-                    # The providers don't force the signature to be unique.
-                    items.append(ThinkingPart(content=summary.text, id=item.id))
-            elif item.type == 'message':
+                    # We use the same id for all summaries so that we can merge them on the round trip.
+                    # We only need to store the signature once.
+                    items.append(
+                        ThinkingPart(
+                            content=summary.text,
+                            id=item.id,
+                            signature=signature,
+                            provider_name=self.system if signature else None,
+                        )
+                    )
+                    signature = None
+                # NOTE: We don't currently handle the raw CoT from gpt-oss `reasoning_text`: https://cookbook.openai.com/articles/gpt-oss/handle-raw-cot
+                # If you need this, please file an issue.
+            elif isinstance(item, responses.ResponseOutputMessage):
                 for content in item.content:
-                    if content.type == 'output_text':  # pragma: no branch
+                    if isinstance(content, responses.ResponseOutputText):  # pragma: no branch
                         items.append(TextPart(content.text))
-            elif item.type == 'function_call':
-                items.append(ToolCallPart(item.name, item.arguments, tool_call_id=item.call_id))
+            elif isinstance(item, responses.ResponseFunctionToolCall):
+                items.append(
+                    ToolCallPart(item.name, item.arguments, tool_call_id=_combine_tool_call_ids(item.call_id, item.id))
+                )
 
         finish_reason: FinishReason | None = None
         provider_details: dict[str, Any] | None = None
@@ -882,7 +914,7 @@ class OpenAIResponsesModel(Model):
         assert isinstance(first_chunk, responses.ResponseCreatedEvent)
         return OpenAIResponsesStreamedResponse(
             model_request_parameters=model_request_parameters,
-            _model_name=self._model_name,
+            _model_name=first_chunk.response.model,
             _response=peekable_response,
             _timestamp=number_to_datetime(first_chunk.response.created_at),
             _provider_name=self._provider.name,
@@ -950,10 +982,15 @@ class OpenAIResponsesModel(Model):
             text = text or {}
             text['verbosity'] = verbosity
 
-        unsupported_model_settings = OpenAIModelProfile.from_profile(self.profile).openai_unsupported_model_settings
+        profile = OpenAIModelProfile.from_profile(self.profile)
+        unsupported_model_settings = profile.openai_unsupported_model_settings
         for setting in unsupported_model_settings:
             model_settings.pop(setting, None)
 
+        include: list[responses.ResponseIncludable] | None = None
+        if profile.openai_supports_encrypted_reasoning_content:
+            include = ['reasoning.encrypted_content']
+
         try:
             extra_headers = model_settings.get('extra_headers', {})
             extra_headers.setdefault('User-Agent', get_user_agent())
@@ -974,6 +1011,7 @@ class OpenAIResponsesModel(Model):
                 reasoning=reasoning,
                 user=model_settings.get('openai_user', NOT_GIVEN),
                 text=text or NOT_GIVEN,
+                include=include or NOT_GIVEN,
                 extra_headers=extra_headers,
                 extra_body=model_settings.get('extra_body'),
             )
@@ -1035,7 +1073,7 @@ class OpenAIResponsesModel(Model):
             ),
         }
 
-    async def _map_messages(
+    async def _map_messages(  # noqa: C901
         self, messages: list[ModelMessage]
     ) -> tuple[str | NotGiven, list[responses.ResponseInputItemParam]]:
         """Just maps a `pydantic_ai.Message` to a `openai.types.responses.ResponseInputParam`."""
@@ -1048,13 +1086,14 @@ class OpenAIResponsesModel(Model):
                     elif isinstance(part, UserPromptPart):
                         openai_messages.append(await self._map_user_prompt(part))
                     elif isinstance(part, ToolReturnPart):
-                        openai_messages.append(
-                            FunctionCallOutput(
-                                type='function_call_output',
-                                call_id=_guard_tool_call_id(t=part),
-                                output=part.model_response_str(),
-                            )
+                        call_id = _guard_tool_call_id(t=part)
+                        call_id, _ = _split_combined_tool_call_id(call_id)
+                        item = FunctionCallOutput(
+                            type='function_call_output',
+                            call_id=call_id,
+                            output=part.model_response_str(),
                         )
+                        openai_messages.append(item)
                     elif isinstance(part, RetryPromptPart):
                         # TODO(Marcelo): How do we test this conditional branch?
                         if part.tool_name is None:  # pragma: no cover
@@ -1062,43 +1101,41 @@ class OpenAIResponsesModel(Model):
                                 Message(role='user', content=[{'type': 'input_text', 'text': part.model_response()}])
                             )
                         else:
-                            openai_messages.append(
-                                FunctionCallOutput(
-                                    type='function_call_output',
-                                    call_id=_guard_tool_call_id(t=part),
-                                    output=part.model_response(),
-                                )
+                            call_id = _guard_tool_call_id(t=part)
+                            call_id, _ = _split_combined_tool_call_id(call_id)
+                            item = FunctionCallOutput(
+                                type='function_call_output',
+                                call_id=call_id,
+                                output=part.model_response(),
                             )
+                            openai_messages.append(item)
                     else:
                         assert_never(part)
             elif isinstance(message, ModelResponse):
-                # last_thinking_part_idx: int | None = None
+                reasoning_item: responses.ResponseReasoningItemParam | None = None
                 for item in message.parts:
                     if isinstance(item, TextPart):
                         openai_messages.append(responses.EasyInputMessageParam(role='assistant', content=item.content))
                     elif isinstance(item, ToolCallPart):
                         openai_messages.append(self._map_tool_call(item))
-                    # OpenAI doesn't return built-in tool calls
                     elif isinstance(item, BuiltinToolCallPart | BuiltinToolReturnPart):
+                        # We don't currently track built-in tool calls from OpenAI
                         pass
                     elif isinstance(item, ThinkingPart):
-                        # NOTE: We don't send ThinkingPart to the providers yet. If you are unsatisfied with this,
-                        # please open an issue. The below code is the code to send thinking to the provider.
-                        # if last_thinking_part_idx is not None:
-                        #     reasoning_item = cast(responses.ResponseReasoningItemParam, openai_messages[last_thinking_part_idx])  # fmt: skip
-                        #     if item.id == reasoning_item['id']:
-                        #         assert isinstance(reasoning_item['summary'], list)
-                        #         reasoning_item['summary'].append(Summary(text=item.content, type='summary_text'))
-                        #         continue
-                        # last_thinking_part_idx = len(openai_messages)
-                        # openai_messages.append(
-                        #     responses.ResponseReasoningItemParam(
-                        #         id=item.id or generate_tool_call_id(),
-                        #         summary=[Summary(text=item.content, type='summary_text')],
-                        #         type='reasoning',
-                        #     )
-                        # )
-                        pass
+                        if reasoning_item is not None and item.id == reasoning_item['id']:
+                            reasoning_item['summary'] = [
+                                *reasoning_item['summary'],
+                                Summary(text=item.content, type='summary_text'),
+                            ]
+                            continue
+
+                        reasoning_item = responses.ResponseReasoningItemParam(
+                            id=item.id or _utils.generate_tool_call_id(),
+                            summary=[Summary(text=item.content, type='summary_text')],
+                            encrypted_content=item.signature if item.provider_name == self.system else None,
+                            type='reasoning',
+                        )
+                        openai_messages.append(reasoning_item)
                     else:
                         assert_never(item)
             else:
@@ -1108,12 +1145,18 @@ class OpenAIResponsesModel(Model):
 
     @staticmethod
     def _map_tool_call(t: ToolCallPart) -> responses.ResponseFunctionToolCallParam:
-        return responses.ResponseFunctionToolCallParam(
-            arguments=t.args_as_json_str(),
-            call_id=_guard_tool_call_id(t=t),
+        call_id = _guard_tool_call_id(t=t)
+        call_id, id = _split_combined_tool_call_id(call_id)
+
+        param = responses.ResponseFunctionToolCallParam(
             name=t.tool_name,
+            arguments=t.args_as_json_str(),
+            call_id=call_id,
             type='function_call',
         )
+        if id:  # pragma: no branch
+            param['id'] = id
+        return param
 
     def _map_json_schema(self, o: OutputObjectDefinition) -> responses.ResponseFormatTextJSONSchemaConfigParam:
         response_format_param: responses.ResponseFormatTextJSONSchemaConfigParam = {
@@ -1231,12 +1274,19 @@ class OpenAIStreamedResponse(StreamedResponse):
                     ignore_leading_whitespace=self._model_profile.ignore_streamed_leading_whitespace,
                 )
                 if maybe_event is not None:  # pragma: no branch
+                    if isinstance(maybe_event, PartStartEvent) and isinstance(maybe_event.part, ThinkingPart):
+                        maybe_event.part.id = 'content'
+                        maybe_event.part.provider_name = self.provider_name
                     yield maybe_event
 
-            # Handle reasoning part of the response, present in DeepSeek models
+            # The `reasoning_content` field is only present in DeepSeek models.
+            # https://api-docs.deepseek.com/guides/reasoning_model
             if reasoning_content := getattr(choice.delta, 'reasoning_content', None):
                 yield self._parts_manager.handle_thinking_delta(
-                    vendor_part_id='reasoning_content', content=reasoning_content
+                    vendor_part_id='reasoning_content',
+                    id='reasoning_content',
+                    content=reasoning_content,
+                    provider_name=self.provider_name,
                 )
 
             for dtc in choice.delta.tool_calls or []:
@@ -1325,7 +1375,7 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
                         vendor_part_id=chunk.item.id,
                         tool_name=chunk.item.name,
                         args=chunk.item.arguments,
-                        tool_call_id=chunk.item.call_id,
+                        tool_call_id=_combine_tool_call_ids(chunk.item.call_id, chunk.item.id),
                     )
                 elif isinstance(chunk.item, responses.ResponseReasoningItem):
                     pass
@@ -1340,7 +1390,15 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
                     )
 
             elif isinstance(chunk, responses.ResponseOutputItemDoneEvent):
-                # NOTE: We only need this if the tool call deltas don't include the final info.
+                if isinstance(chunk.item, responses.ResponseReasoningItem):
+                    # Add the signature to the part corresponding to the first summary item
+                    signature = chunk.item.encrypted_content
+                    yield self._parts_manager.handle_thinking_delta(
+                        vendor_part_id=f'{chunk.item.id}-0',
+                        id=chunk.item.id,
+                        signature=signature,
+                        provider_name=self.provider_name if signature else None,
+                    )
                 pass
 
             elif isinstance(chunk, responses.ResponseReasoningSummaryPartAddedEvent):
@@ -1458,3 +1516,17 @@ def _map_usage(response: chat.ChatCompletion | ChatCompletionChunk | responses.R
             u.input_audio_tokens = response_usage.prompt_tokens_details.audio_tokens or 0
             u.cache_read_tokens = response_usage.prompt_tokens_details.cached_tokens or 0
         return u
+
+
+def _combine_tool_call_ids(call_id: str, id: str | None) -> str:
+    # When reasoning, the Responses API requires the `ResponseFunctionToolCall` to be returned with both the `call_id` and `id` fields.
+    # Our `ToolCallPart` has only the `call_id` field, so we combine the two fields into a single string.
+    return f'{call_id}|{id}' if id else call_id
+
+
+def _split_combined_tool_call_id(combined_id: str) -> tuple[str, str | None]:
+    if '|' in combined_id:
+        call_id, id = combined_id.split('|', 1)
+        return call_id, id
+    else:
+        return combined_id, None  # pragma: no cover

pydantic_ai/profiles/openai.py

@@ -38,6 +38,12 @@ class OpenAIModelProfile(ModelProfile):
     openai_system_prompt_role: OpenAISystemPromptRole | None = None
     """The role to use for the system prompt message. If not provided, defaults to `'system'`."""
 
+    openai_chat_supports_web_search: bool = False
+    """Whether the model supports web search in Chat Completions API."""
+
+    openai_supports_encrypted_reasoning_content: bool = False
+    """Whether the model supports including encrypted reasoning content in the response."""
+
     def __post_init__(self):  # pragma: no cover
         if not self.openai_supports_sampling_settings:
             warnings.warn(
@@ -50,6 +56,9 @@ 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') or model_name.startswith('gpt-5')
+    # Check if the model supports web search (only specific search-preview models)
+    supports_web_search = '-search-preview' in model_name
+
     # 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.
@@ -77,6 +86,8 @@ def openai_model_profile(model_name: str) -> ModelProfile:
         supports_json_object_output=True,
         openai_unsupported_model_settings=openai_unsupported_model_settings,
         openai_system_prompt_role=openai_system_prompt_role,
+        openai_chat_supports_web_search=supports_web_search,
+        openai_supports_encrypted_reasoning_content=is_reasoning_model,
     )
 
 

pydantic_ai/providers/__init__.py

@@ -47,6 +47,9 @@ class Provider(ABC, Generic[InterfaceClient]):
         """The model profile for the named model, if available."""
         return None  # pragma: no cover
 
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}(name={self.name}, base_url={self.base_url})'
+
 
 def infer_provider_class(provider: str) -> type[Provider[Any]]:  # noqa: C901
     """Infers the provider class from the provider name."""

pydantic_ai/providers/anthropic.py

@@ -45,12 +45,15 @@ class AnthropicProvider(Provider[AsyncAnthropicClient]):
     def __init__(self, *, anthropic_client: AsyncAnthropicClient | None = None) -> None: ...
 
     @overload
-    def __init__(self, *, api_key: str | None = None, http_client: httpx.AsyncClient | None = None) -> None: ...
+    def __init__(
+        self, *, api_key: str | None = None, base_url: str | None = None, http_client: httpx.AsyncClient | None = None
+    ) -> None: ...
 
     def __init__(
         self,
         *,
         api_key: str | None = None,
+        base_url: str | None = None,
         anthropic_client: AsyncAnthropicClient | None = None,
         http_client: httpx.AsyncClient | None = None,
     ) -> None:
@@ -59,6 +62,7 @@ class AnthropicProvider(Provider[AsyncAnthropicClient]):
         Args:
             api_key: The API key to use for authentication, if not provided, the `ANTHROPIC_API_KEY` environment variable
                 will be used if available.
+            base_url: The base URL to use for the Anthropic API.
             anthropic_client: An existing [`AsyncAnthropic`](https://github.com/anthropics/anthropic-sdk-python)
                 client to use. If provided, the `api_key` and `http_client` arguments will be ignored.
             http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
@@ -68,14 +72,14 @@ class AnthropicProvider(Provider[AsyncAnthropicClient]):
             assert api_key is None, 'Cannot provide both `anthropic_client` and `api_key`'
             self._client = anthropic_client
         else:
-            api_key = api_key or os.environ.get('ANTHROPIC_API_KEY')
+            api_key = api_key or os.getenv('ANTHROPIC_API_KEY')
             if not api_key:
                 raise UserError(
                     'Set the `ANTHROPIC_API_KEY` environment variable or pass it via `AnthropicProvider(api_key=...)`'
                     'to use the Anthropic provider.'
                 )
             if http_client is not None:
-                self._client = AsyncAnthropic(api_key=api_key, http_client=http_client)
+                self._client = AsyncAnthropic(api_key=api_key, base_url=base_url, http_client=http_client)
             else:
                 http_client = cached_async_http_client(provider='anthropic')
-                self._client = AsyncAnthropic(api_key=api_key, http_client=http_client)
+                self._client = AsyncAnthropic(api_key=api_key, base_url=base_url, http_client=http_client)

pydantic_ai/providers/bedrock.py

@@ -48,6 +48,14 @@ def bedrock_amazon_model_profile(model_name: str) -> ModelProfile | None:
     return profile
 
 
+def bedrock_deepseek_model_profile(model_name: str) -> ModelProfile | None:
+    """Get the model profile for a DeepSeek model used via Bedrock."""
+    profile = deepseek_model_profile(model_name)
+    if 'r1' in model_name:
+        return BedrockModelProfile(bedrock_send_back_thinking_parts=True).update(profile)
+    return profile  # pragma: no cover
+
+
 class BedrockProvider(Provider[BaseClient]):
     """Provider for AWS Bedrock."""
 
@@ -74,7 +82,7 @@ class BedrockProvider(Provider[BaseClient]):
             'cohere': cohere_model_profile,
             'amazon': bedrock_amazon_model_profile,
             'meta': meta_model_profile,
-            'deepseek': deepseek_model_profile,
+            'deepseek': bedrock_deepseek_model_profile,
         }
 
         # Split the model name into parts

pydantic_ai/providers/cohere.py

@@ -60,14 +60,14 @@ class CohereProvider(Provider[AsyncClientV2]):
             assert api_key is None, 'Cannot provide both `cohere_client` and `api_key`'
             self._client = cohere_client
         else:
-            api_key = api_key or os.environ.get('CO_API_KEY')
+            api_key = api_key or os.getenv('CO_API_KEY')
             if not api_key:
                 raise UserError(
                     'Set the `CO_API_KEY` environment variable or pass it via `CohereProvider(api_key=...)`'
                     'to use the Cohere provider.'
                 )
 
-            base_url = os.environ.get('CO_BASE_URL')
+            base_url = os.getenv('CO_BASE_URL')
             if http_client is not None:
                 self._client = AsyncClientV2(api_key=api_key, httpx_client=http_client, base_url=base_url)
             else: