pydantic-ai-slim 0.7.2__py3-none-any.whl → 0.7.3__py3-none-any.whl

pydantic_ai/models/openai.py

@@ -9,7 +9,7 @@ from datetime import datetime
 from typing import Any, Literal, Union, cast, overload
 
 from pydantic import ValidationError
-from typing_extensions import assert_never
+from typing_extensions import assert_never, deprecated
 
 from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage
 from .._output import DEFAULT_OUTPUT_TOOL_NAME, OutputObjectDefinition
@@ -40,7 +40,7 @@ from ..messages import (
     VideoUrl,
 )
 from ..profiles import ModelProfile, ModelProfileSpec
-from ..profiles.openai import OpenAIModelProfile
+from ..profiles.openai import OpenAIModelProfile, OpenAISystemPromptRole
 from ..providers import Provider, infer_provider
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
@@ -100,8 +100,6 @@ Using this more broad type for the model name instead of the ChatModel definitio
 allows this model to be used more easily with other model types (ie, Ollama, Deepseek).
 """
 
-OpenAISystemPromptRole = Literal['system', 'developer', 'user']
-
 
 class OpenAIModelSettings(ModelSettings, total=False):
     """Settings used for an OpenAI model request."""
@@ -196,10 +194,59 @@ class OpenAIModel(Model):
     """
 
     client: AsyncOpenAI = field(repr=False)
-    system_prompt_role: OpenAISystemPromptRole | None = field(default=None, repr=False)
 
     _model_name: OpenAIModelName = field(repr=False)
-    _system: str = field(default='openai', repr=False)
+    _provider: Provider[AsyncOpenAI] = field(repr=False)
+
+    @overload
+    def __init__(
+        self,
+        model_name: OpenAIModelName,
+        *,
+        provider: Literal[
+            'openai',
+            'deepseek',
+            'azure',
+            'openrouter',
+            'moonshotai',
+            'vercel',
+            'grok',
+            'fireworks',
+            'together',
+            'heroku',
+            'github',
+            'ollama',
+        ]
+        | Provider[AsyncOpenAI] = 'openai',
+        profile: ModelProfileSpec | None = None,
+        settings: ModelSettings | None = None,
+    ) -> None: ...
+
+    @deprecated('Set the `system_prompt_role` in the `OpenAIModelProfile` instead.')
+    @overload
+    def __init__(
+        self,
+        model_name: OpenAIModelName,
+        *,
+        provider: Literal[
+            'openai',
+            'deepseek',
+            'azure',
+            'openrouter',
+            'moonshotai',
+            'vercel',
+            'grok',
+            'fireworks',
+            'together',
+            'heroku',
+            'github',
+            'ollama',
+        ]
+        | Provider[AsyncOpenAI] = 'openai',
+        profile: ModelProfileSpec | None = None,
+        system_prompt_role: OpenAISystemPromptRole | None = None,
+        settings: ModelSettings | None = None,
+    ) -> None: ...
 
     def __init__(
         self,
@@ -240,16 +287,33 @@ class OpenAIModel(Model):
 
         if isinstance(provider, str):
             provider = infer_provider(provider)
+        self._provider = provider
         self.client = provider.client
 
-        self.system_prompt_role = system_prompt_role
-
         super().__init__(settings=settings, profile=profile or provider.model_profile)
 
+        if system_prompt_role is not None:
+            self.profile = OpenAIModelProfile(openai_system_prompt_role=system_prompt_role).update(self.profile)
+
     @property
     def base_url(self) -> str:
         return str(self.client.base_url)
 
+    @property
+    def model_name(self) -> OpenAIModelName:
+        """The model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str:
+        """The model provider."""
+        return self._provider.name
+
+    @property
+    @deprecated('Set the `system_prompt_role` in the `OpenAIModelProfile` instead.')
+    def system_prompt_role(self) -> OpenAISystemPromptRole | None:
+        return OpenAIModelProfile.from_profile(self.profile).openai_system_prompt_role
+
     async def request(
         self,
         messages: list[ModelMessage],
@@ -261,7 +325,6 @@ class OpenAIModel(Model):
             messages, False, cast(OpenAIModelSettings, model_settings or {}), model_request_parameters
         )
         model_response = self._process_response(response)
-        model_response.usage.requests = 1
         return model_response
 
     @asynccontextmanager
@@ -279,16 +342,6 @@ class OpenAIModel(Model):
         async with response:
             yield await self._process_streamed_response(response, model_request_parameters)
 
-    @property
-    def model_name(self) -> OpenAIModelName:
-        """The model name."""
-        return self._model_name
-
-    @property
-    def system(self) -> str:
-        """The system / model provider."""
-        return self._system
-
     @overload
     async def _completions_create(
         self,
@@ -445,8 +498,8 @@ class OpenAIModel(Model):
             usage=_map_usage(response),
             model_name=response.model,
             timestamp=timestamp,
-            vendor_details=vendor_details,
-            vendor_id=response.id,
+            provider_details=vendor_details,
+            provider_request_id=response.id,
         )
 
     async def _process_streamed_response(
@@ -562,9 +615,10 @@ class OpenAIModel(Model):
     async def _map_user_message(self, message: ModelRequest) -> AsyncIterable[chat.ChatCompletionMessageParam]:
         for part in message.parts:
             if isinstance(part, SystemPromptPart):
-                if self.system_prompt_role == 'developer':
+                system_prompt_role = OpenAIModelProfile.from_profile(self.profile).openai_system_prompt_role
+                if system_prompt_role == 'developer':
                     yield chat.ChatCompletionDeveloperMessageParam(role='developer', content=part.content)
-                elif self.system_prompt_role == 'user':
+                elif system_prompt_role == 'user':
                     yield chat.ChatCompletionUserMessageParam(role='user', content=part.content)
                 else:
                     yield chat.ChatCompletionSystemMessageParam(role='system', content=part.content)
@@ -660,10 +714,9 @@ class OpenAIResponsesModel(Model):
     """
 
     client: AsyncOpenAI = field(repr=False)
-    system_prompt_role: OpenAISystemPromptRole | None = field(default=None)
 
     _model_name: OpenAIModelName = field(repr=False)
-    _system: str = field(default='openai', repr=False)
+    _provider: Provider[AsyncOpenAI] = field(repr=False)
 
     def __init__(
         self,
@@ -686,6 +739,7 @@ class OpenAIResponsesModel(Model):
 
         if isinstance(provider, str):
             provider = infer_provider(provider)
+        self._provider = provider
         self.client = provider.client
 
         super().__init__(settings=settings, profile=profile or provider.model_profile)
@@ -697,8 +751,8 @@ class OpenAIResponsesModel(Model):
 
     @property
     def system(self) -> str:
-        """The system / model provider."""
-        return self._system
+        """The model provider."""
+        return self._provider.name
 
     async def request(
         self,
@@ -747,7 +801,7 @@ class OpenAIResponsesModel(Model):
             items,
             usage=_map_usage(response),
             model_name=response.model,
-            vendor_id=response.id,
+            provider_request_id=response.id,
             timestamp=timestamp,
         )
 
@@ -1265,10 +1319,10 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
         return self._timestamp
 
 
-def _map_usage(response: chat.ChatCompletion | ChatCompletionChunk | responses.Response) -> usage.Usage:
+def _map_usage(response: chat.ChatCompletion | ChatCompletionChunk | responses.Response) -> usage.RequestUsage:
     response_usage = response.usage
     if response_usage is None:
-        return usage.Usage()
+        return usage.RequestUsage()
     elif isinstance(response_usage, responses.ResponseUsage):
         details: dict[str, int] = {
             key: value
@@ -1278,29 +1332,29 @@ def _map_usage(response: chat.ChatCompletion | ChatCompletionChunk | responses.R
             if isinstance(value, int)
         }
         details['reasoning_tokens'] = response_usage.output_tokens_details.reasoning_tokens
-        details['cached_tokens'] = response_usage.input_tokens_details.cached_tokens
-        return usage.Usage(
-            request_tokens=response_usage.input_tokens,
-            response_tokens=response_usage.output_tokens,
-            total_tokens=response_usage.total_tokens,
+        return usage.RequestUsage(
+            input_tokens=response_usage.input_tokens,
+            output_tokens=response_usage.output_tokens,
+            cache_read_tokens=response_usage.input_tokens_details.cached_tokens,
             details=details,
         )
     else:
         details = {
             key: value
             for key, value in response_usage.model_dump(
-                exclude={'prompt_tokens', 'completion_tokens', 'total_tokens'}
+                exclude_none=True, exclude={'prompt_tokens', 'completion_tokens', 'total_tokens'}
             ).items()
             if isinstance(value, int)
         }
+        u = usage.RequestUsage(
+            input_tokens=response_usage.prompt_tokens,
+            output_tokens=response_usage.completion_tokens,
+            details=details,
+        )
         if response_usage.completion_tokens_details is not None:
             details.update(response_usage.completion_tokens_details.model_dump(exclude_none=True))
+            u.output_audio_tokens = response_usage.completion_tokens_details.audio_tokens or 0
         if response_usage.prompt_tokens_details is not None:
-            details.update(response_usage.prompt_tokens_details.model_dump(exclude_none=True))
-        return usage.Usage(
-            requests=1,
-            request_tokens=response_usage.prompt_tokens,
-            response_tokens=response_usage.completion_tokens,
-            total_tokens=response_usage.total_tokens,
-            details=details,
-        )
+            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

pydantic_ai/models/test.py

@@ -31,7 +31,7 @@ from ..messages import (
 from ..profiles import ModelProfileSpec
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
-from ..usage import Usage
+from ..usage import RequestUsage
 from . import Model, ModelRequestParameters, StreamedResponse
 from .function import _estimate_string_tokens, _estimate_usage  # pyright: ignore[reportPrivateUsage]
 
@@ -113,7 +113,6 @@ class TestModel(Model):
         self.last_model_request_parameters = model_request_parameters
         model_response = self._request(messages, model_settings, model_request_parameters)
         model_response.usage = _estimate_usage([*messages, model_response])
-        model_response.usage.requests = 1
         return model_response
 
     @asynccontextmanager
@@ -141,7 +140,7 @@ class TestModel(Model):
 
     @property
     def system(self) -> str:
-        """The system / model provider."""
+        """The model provider."""
         return self._system
 
     def gen_tool_args(self, tool_def: ToolDefinition) -> Any:
@@ -468,6 +467,6 @@ class _JsonSchemaTestData:
         return s
 
 
-def _get_string_usage(text: str) -> Usage:
+def _get_string_usage(text: str) -> RequestUsage:
     response_tokens = _estimate_string_tokens(text)
-    return Usage(response_tokens=response_tokens, total_tokens=response_tokens)
+    return RequestUsage(output_tokens=response_tokens)

pydantic_ai/profiles/openai.py

@@ -2,11 +2,13 @@ from __future__ import annotations as _annotations
 
 import re
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Literal
 
 from . import ModelProfile
 from ._json_schema import JsonSchema, JsonSchemaTransformer
 
+OpenAISystemPromptRole = Literal['system', 'developer', 'user']
+
 
 @dataclass
 class OpenAIModelProfile(ModelProfile):
@@ -26,8 +28,10 @@ class OpenAIModelProfile(ModelProfile):
     # safe to pass that value along.  Default is `True` to preserve existing
     # behaviour for OpenAI itself and most providers.
     openai_supports_tool_choice_required: bool = True
-    """Whether the provider accepts the value ``tool_choice='required'`` in the
-    request payload."""
+    """Whether the provider accepts the value ``tool_choice='required'`` in the request payload."""
+
+    openai_system_prompt_role: OpenAISystemPromptRole | None = None
+    """The role to use for the system prompt message. If not provided, defaults to `'system'`."""
 
 
 def openai_model_profile(model_name: str) -> ModelProfile:
@@ -36,11 +40,17 @@ def openai_model_profile(model_name: str) -> ModelProfile:
     # 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.
+
+    # The o1-mini model doesn't support the `system` role, so we default to `user`.
+    # See https://github.com/pydantic/pydantic-ai/issues/974 for more details.
+    openai_system_prompt_role = 'user' if model_name.startswith('o1-mini') else None
+
     return OpenAIModelProfile(
         json_schema_transformer=OpenAIJsonSchemaTransformer,
         supports_json_schema_output=True,
         supports_json_object_output=True,
         openai_supports_sampling_settings=not is_reasoning_model,
+        openai_system_prompt_role=openai_system_prompt_role,
     )
 
 

pydantic_ai/providers/openai.py

@@ -23,7 +23,7 @@ class OpenAIProvider(Provider[AsyncOpenAI]):
 
     @property
     def name(self) -> str:
-        return 'openai'  # pragma: no cover
+        return 'openai'
 
     @property
     def base_url(self) -> str:

pydantic_ai/result.py

@@ -27,7 +27,7 @@ from .output import (
     OutputDataT,
     ToolOutput,
 )
-from .usage import Usage, UsageLimits
+from .usage import RunUsage, UsageLimits
 
 __all__ = (
     'OutputDataT',
@@ -52,7 +52,7 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
     _tool_manager: ToolManager[AgentDepsT]
 
     _agent_stream_iterator: AsyncIterator[AgentStreamEvent] | None = field(default=None, init=False)
-    _initial_run_ctx_usage: Usage = field(init=False)
+    _initial_run_ctx_usage: RunUsage = field(init=False)
 
     def __post_init__(self):
         self._initial_run_ctx_usage = copy(self._run_ctx.usage)
@@ -110,7 +110,7 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
         """Get the current state of the response."""
         return self._raw_stream_response.get()
 
-    def usage(self) -> Usage:
+    def usage(self) -> RunUsage:
         """Return the usage of the whole run.
 
         !!! note
@@ -382,7 +382,7 @@ class StreamedRunResult(Generic[AgentDepsT, OutputDataT]):
         await self._marked_completed(self._stream_response.get())
         return output
 
-    def usage(self) -> Usage:
+    def usage(self) -> RunUsage:
         """Return the usage of the whole run.
 
         !!! note
@@ -425,7 +425,7 @@ class FinalResult(Generic[OutputDataT]):
 def _get_usage_checking_stream_response(
     stream_response: models.StreamedResponse,
     limits: UsageLimits | None,
-    get_usage: Callable[[], Usage],
+    get_usage: Callable[[], RunUsage],
 ) -> AsyncIterator[AgentStreamEvent]:
     if limits is not None and limits.has_token_limits():
 

pydantic_ai/run.py

@@ -66,9 +66,7 @@ class AgentRun(Generic[AgentDepsT, OutputDataT]):
             CallToolsNode(
                 model_response=ModelResponse(
                     parts=[TextPart(content='The capital of France is Paris.')],
-                    usage=Usage(
-                        requests=1, request_tokens=56, response_tokens=7, total_tokens=63
-                    ),
+                    usage=RequestUsage(input_tokens=56, output_tokens=7),
                     model_name='gpt-4o',
                     timestamp=datetime.datetime(...),
                 )
@@ -203,12 +201,7 @@ class AgentRun(Generic[AgentDepsT, OutputDataT]):
                     CallToolsNode(
                         model_response=ModelResponse(
                             parts=[TextPart(content='The capital of France is Paris.')],
-                            usage=Usage(
-                                requests=1,
-                                request_tokens=56,
-                                response_tokens=7,
-                                total_tokens=63,
-                            ),
+                            usage=RequestUsage(input_tokens=56, output_tokens=7),
                             model_name='gpt-4o',
                             timestamp=datetime.datetime(...),
                         )
@@ -235,7 +228,7 @@ class AgentRun(Generic[AgentDepsT, OutputDataT]):
         assert isinstance(next_node, End), f'Unexpected node type: {type(next_node)}'
         return next_node
 
-    def usage(self) -> _usage.Usage:
+    def usage(self) -> _usage.RunUsage:
         """Get usage statistics for the run so far, including token usage, model requests, and so on."""
         return self._graph_run.state.usage
 
@@ -352,6 +345,6 @@ class AgentRunResult(Generic[OutputDataT]):
             self.new_messages(output_tool_return_content=output_tool_return_content)
         )
 
-    def usage(self) -> _usage.Usage:
+    def usage(self) -> _usage.RunUsage:
         """Return the usage of the whole run."""
         return self._state.usage