pydantic-ai-slim 0.0.22__py3-none-any.whl → 0.0.23__py3-none-any.whl

pydantic_ai/models/openai.py

@@ -29,8 +29,8 @@ from ..messages import (
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
-    AgentModel,
     Model,
+    ModelRequestParameters,
     StreamedResponse,
     cached_async_http_client,
     check_allow_model_requests,
@@ -46,10 +46,16 @@ except ImportError as _import_error:
         "you can use the `openai` optional group — `pip install 'pydantic-ai-slim[openai]'`"
     ) from _import_error
 
-OpenAIModelName = Union[ChatModel, str]
+OpenAIModelName = Union[str, ChatModel]
 """
+Possible OpenAI model names.
+
+Since OpenAI supports a variety of date-stamped models, we explicitly list the latest models but
+allow any name in the type hints.
+See [the OpenAI docs](https://platform.openai.com/docs/models) for a full list.
+
 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, Deepseek)
+allows this model to be used more easily with other model types (ie, Ollama, Deepseek).
 """
 
 OpenAISystemPromptRole = Literal['system', 'developer', 'user']
@@ -58,7 +64,12 @@ 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
+    openai_reasoning_effort: chat.ChatCompletionReasoningEffort
+    """
+    Constrains effort on reasoning for [reasoning models](https://platform.openai.com/docs/guides/reasoning).
+    Currently supported values are `low`, `medium`, and `high`. Reducing reasoning effort can
+    result in faster responses and fewer tokens used on reasoning in a response.
+    """
 
 
 @dataclass(init=False)
@@ -70,10 +81,12 @@ class OpenAIModel(Model):
     Apart from `__init__`, all methods are private or match those of the base class.
     """
 
-    model_name: OpenAIModelName
     client: AsyncOpenAI = field(repr=False)
     system_prompt_role: OpenAISystemPromptRole | None = field(default=None)
 
+    _model_name: OpenAIModelName = field(repr=False)
+    _system: str | None = field(repr=False)
+
     def __init__(
         self,
         model_name: OpenAIModelName,
@@ -83,6 +96,7 @@ class OpenAIModel(Model):
         openai_client: AsyncOpenAI | None = None,
         http_client: AsyncHTTPClient | None = None,
         system_prompt_role: OpenAISystemPromptRole | None = None,
+        system: str | None = 'openai',
     ):
         """Initialize an OpenAI model.
 
@@ -100,8 +114,10 @@ class OpenAIModel(Model):
             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.
+            system: The model provider used, defaults to `openai`. This is for observability purposes, you must
+                customize the `base_url` and `api_key` to use a different provider.
         """
-        self.model_name: OpenAIModelName = model_name
+        self._model_name = model_name
         # This is a workaround for the OpenAI client requiring an API key, whilst locally served,
         # openai compatible models do not always need an API key.
         if api_key is None and 'OPENAI_API_KEY' not in os.environ and base_url is not None and openai_client is None:
@@ -116,84 +132,70 @@ class OpenAIModel(Model):
         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,
-        *,
-        function_tools: list[ToolDefinition],
-        allow_text_result: bool,
-        result_tools: list[ToolDefinition],
-    ) -> AgentModel:
-        check_allow_model_requests()
-        tools = [self._map_tool_definition(r) for r in function_tools]
-        if result_tools:
-            tools += [self._map_tool_definition(r) for r in result_tools]
-        return OpenAIAgentModel(
-            self.client,
-            self.model_name,
-            allow_text_result,
-            tools,
-            self.system_prompt_role,
-        )
+        self._system = system
 
     def name(self) -> str:
-        return f'openai:{self.model_name}'
-
-    @staticmethod
-    def _map_tool_definition(f: ToolDefinition) -> chat.ChatCompletionToolParam:
-        return {
-            'type': 'function',
-            'function': {
-                'name': f.name,
-                'description': f.description,
-                'parameters': f.parameters_json_schema,
-            },
-        }
-
-
-@dataclass
-class OpenAIAgentModel(AgentModel):
-    """Implementation of `AgentModel` for OpenAI models."""
-
-    client: AsyncOpenAI
-    model_name: OpenAIModelName
-    allow_text_result: bool
-    tools: list[chat.ChatCompletionToolParam]
-    system_prompt_role: OpenAISystemPromptRole | None
+        return f'openai:{self._model_name}'
 
     async def request(
-        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
     ) -> tuple[ModelResponse, usage.Usage]:
-        response = await self._completions_create(messages, False, cast(OpenAIModelSettings, model_settings or {}))
+        check_allow_model_requests()
+        response = await self._completions_create(
+            messages, False, cast(OpenAIModelSettings, model_settings or {}), model_request_parameters
+        )
         return self._process_response(response), _map_usage(response)
 
     @asynccontextmanager
     async def request_stream(
-        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
     ) -> AsyncIterator[StreamedResponse]:
-        response = await self._completions_create(messages, True, cast(OpenAIModelSettings, model_settings or {}))
+        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)
 
     @overload
     async def _completions_create(
-        self, messages: list[ModelMessage], stream: Literal[True], model_settings: OpenAIModelSettings
+        self,
+        messages: list[ModelMessage],
+        stream: Literal[True],
+        model_settings: OpenAIModelSettings,
+        model_request_parameters: ModelRequestParameters,
     ) -> AsyncStream[ChatCompletionChunk]:
         pass
 
     @overload
     async def _completions_create(
-        self, messages: list[ModelMessage], stream: Literal[False], model_settings: OpenAIModelSettings
+        self,
+        messages: list[ModelMessage],
+        stream: Literal[False],
+        model_settings: OpenAIModelSettings,
+        model_request_parameters: ModelRequestParameters,
     ) -> chat.ChatCompletion:
         pass
 
     async def _completions_create(
-        self, messages: list[ModelMessage], stream: bool, model_settings: OpenAIModelSettings
+        self,
+        messages: list[ModelMessage],
+        stream: bool,
+        model_settings: OpenAIModelSettings,
+        model_request_parameters: ModelRequestParameters,
     ) -> chat.ChatCompletion | AsyncStream[ChatCompletionChunk]:
+        tools = self._get_tools(model_request_parameters)
+
         # standalone function to make it easier to override
-        if not self.tools:
+        if not tools:
             tool_choice: Literal['none', 'required', 'auto'] | None = None
-        elif not self.allow_text_result:
+        elif not model_request_parameters.allow_text_result:
             tool_choice = 'required'
         else:
             tool_choice = 'auto'
@@ -201,11 +203,11 @@ class OpenAIAgentModel(AgentModel):
         openai_messages = list(chain(*(self._map_message(m) for m in messages)))
 
         return await self.client.chat.completions.create(
-            model=self.model_name,
+            model=self._model_name,
             messages=openai_messages,
             n=1,
             parallel_tool_calls=model_settings.get('parallel_tool_calls', NOT_GIVEN),
-            tools=self.tools or NOT_GIVEN,
+            tools=tools or NOT_GIVEN,
             tool_choice=tool_choice or NOT_GIVEN,
             stream=stream,
             stream_options={'include_usage': True} if stream else NOT_GIVEN,
@@ -217,6 +219,7 @@ class OpenAIAgentModel(AgentModel):
             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),
+            reasoning_effort=model_settings.get('openai_reasoning_effort', NOT_GIVEN),
         )
 
     def _process_response(self, response: chat.ChatCompletion) -> ModelResponse:
@@ -229,7 +232,7 @@ class OpenAIAgentModel(AgentModel):
         if choice.message.tool_calls is not None:
             for c in choice.message.tool_calls:
                 items.append(ToolCallPart(c.function.name, c.function.arguments, c.id))
-        return ModelResponse(items, model_name=self.model_name, timestamp=timestamp)
+        return ModelResponse(items, model_name=self._model_name, timestamp=timestamp)
 
     async def _process_streamed_response(self, response: AsyncStream[ChatCompletionChunk]) -> OpenAIStreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
@@ -239,11 +242,17 @@ class OpenAIAgentModel(AgentModel):
             raise UnexpectedModelBehavior('Streamed response ended without content or tool calls')
 
         return OpenAIStreamedResponse(
-            _model_name=self.model_name,
+            _model_name=self._model_name,
             _response=peekable_response,
             _timestamp=datetime.fromtimestamp(first_chunk.created, tz=timezone.utc),
         )
 
+    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.result_tools:
+            tools += [self._map_tool_definition(r) for r in model_request_parameters.result_tools]
+        return tools
+
     def _map_message(self, message: ModelMessage) -> Iterable[chat.ChatCompletionMessageParam]:
         """Just maps a `pydantic_ai.Message` to a `openai.types.ChatCompletionMessageParam`."""
         if isinstance(message, ModelRequest):
@@ -255,7 +264,7 @@ class OpenAIAgentModel(AgentModel):
                 if isinstance(item, TextPart):
                     texts.append(item.content)
                 elif isinstance(item, ToolCallPart):
-                    tool_calls.append(_map_tool_call(item))
+                    tool_calls.append(self._map_tool_call(item))
                 else:
                     assert_never(item)
             message_param = chat.ChatCompletionAssistantMessageParam(role='assistant')
@@ -269,6 +278,25 @@ class OpenAIAgentModel(AgentModel):
         else:
             assert_never(message)
 
+    @staticmethod
+    def _map_tool_call(t: ToolCallPart) -> chat.ChatCompletionMessageToolCallParam:
+        return chat.ChatCompletionMessageToolCallParam(
+            id=_guard_tool_call_id(t=t, model_source='OpenAI'),
+            type='function',
+            function={'name': t.tool_name, 'arguments': t.args_as_json_str()},
+        )
+
+    @staticmethod
+    def _map_tool_definition(f: ToolDefinition) -> chat.ChatCompletionToolParam:
+        return {
+            'type': 'function',
+            'function': {
+                'name': f.name,
+                'description': f.description,
+                'parameters': f.parameters_json_schema,
+            },
+        }
+
     def _map_user_message(self, message: ModelRequest) -> Iterable[chat.ChatCompletionMessageParam]:
         for part in message.parts:
             if isinstance(part, SystemPromptPart):
@@ -334,14 +362,6 @@ class OpenAIStreamedResponse(StreamedResponse):
         return self._timestamp
 
 
-def _map_tool_call(t: ToolCallPart) -> chat.ChatCompletionMessageToolCallParam:
-    return chat.ChatCompletionMessageToolCallParam(
-        id=_guard_tool_call_id(t=t, model_source='OpenAI'),
-        type='function',
-        function={'name': t.tool_name, 'arguments': t.args_as_json_str()},
-    )
-
-
 def _map_usage(response: chat.ChatCompletion | ChatCompletionChunk) -> usage.Usage:
     response_usage = response.usage
     if response_usage is None:

pydantic_ai/models/test.py

@@ -26,8 +26,8 @@ from ..result import Usage
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
-    AgentModel,
     Model,
+    ModelRequestParameters,
     StreamedResponse,
 )
 from .function import _estimate_string_tokens, _estimate_usage  # pyright: ignore[reportPrivateUsage]
@@ -71,102 +71,92 @@ class TestModel(Model):
     """If set, these args will be passed to the result tool."""
     seed: int = 0
     """Seed for generating random data."""
-    agent_model_function_tools: list[ToolDefinition] | None = field(default=None, init=False)
-    """Definition of function tools passed to the model.
+    last_model_request_parameters: ModelRequestParameters | None = field(default=None, init=False)
+    """The last ModelRequestParameters passed to the model in a request.
 
-    This is set when the model is called, so will reflect the function tools from the last step of the last run.
-    """
-    agent_model_allow_text_result: bool | None = field(default=None, init=False)
-    """Whether plain text responses from the model are allowed.
+    The ModelRequestParameters contains information about the function and result tools available during request handling.
 
-    This is set when the model is called, so will reflect the value from the last step of the last run.
+    This is set when a request is made, so will reflect the function tools from the last step of the last run.
     """
-    agent_model_result_tools: list[ToolDefinition] | None = field(default=None, init=False)
-    """Definition of result tools passed to the model.
+    _model_name: str = field(default='test', repr=False)
+    _system: str | None = field(default=None, repr=False)
 
-    This is set when the model is called, so will reflect the result tools from the last step of the last run.
-    """
+    async def request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> tuple[ModelResponse, Usage]:
+        self.last_model_request_parameters = model_request_parameters
 
-    async def agent_model(
+        model_response = self._request(messages, model_settings, model_request_parameters)
+        usage = _estimate_usage([*messages, model_response])
+        return model_response, usage
+
+    @asynccontextmanager
+    async def request_stream(
         self,
-        *,
-        function_tools: list[ToolDefinition],
-        allow_text_result: bool,
-        result_tools: list[ToolDefinition],
-    ) -> AgentModel:
-        self.agent_model_function_tools = function_tools
-        self.agent_model_allow_text_result = allow_text_result
-        self.agent_model_result_tools = result_tools
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> AsyncIterator[StreamedResponse]:
+        self.last_model_request_parameters = model_request_parameters
+
+        model_response = self._request(messages, model_settings, model_request_parameters)
+        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()
+
+    def _get_tool_calls(self, model_request_parameters: ModelRequestParameters) -> list[tuple[str, ToolDefinition]]:
         if self.call_tools == 'all':
-            tool_calls = [(r.name, r) for r in function_tools]
+            return [(r.name, r) for r in model_request_parameters.function_tools]
         else:
-            function_tools_lookup = {t.name: t for t in function_tools}
+            function_tools_lookup = {t.name: t for t in model_request_parameters.function_tools}
             tools_to_call = (function_tools_lookup[name] for name in self.call_tools)
-            tool_calls = [(r.name, r) for r in tools_to_call]
+            return [(r.name, r) for r in tools_to_call]
 
+    def _get_result(self, model_request_parameters: ModelRequestParameters) -> _TextResult | _FunctionToolResult:
         if self.custom_result_text is not None:
-            assert allow_text_result, 'Plain response not allowed, but `custom_result_text` is set.'
+            assert (
+                model_request_parameters.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: _TextResult | _FunctionToolResult = _TextResult(self.custom_result_text)
+            return _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]
+            assert (
+                model_request_parameters.result_tools is not None
+            ), 'No result tools provided, but `custom_result_args` is set.'
+            result_tool = model_request_parameters.result_tools[0]
 
             if k := result_tool.outer_typed_dict_key:
-                result = _FunctionToolResult({k: self.custom_result_args})
+                return _FunctionToolResult({k: self.custom_result_args})
             else:
-                result = _FunctionToolResult(self.custom_result_args)
-        elif allow_text_result:
-            result = _TextResult(None)
-        elif result_tools:
-            result = _FunctionToolResult(None)
+                return _FunctionToolResult(self.custom_result_args)
+        elif model_request_parameters.allow_text_result:
+            return _TextResult(None)
+        elif model_request_parameters.result_tools:
+            return _FunctionToolResult(None)
         else:
-            result = _TextResult(None)
-
-        return TestAgentModel(tool_calls, result, result_tools, self.seed)
-
-    def name(self) -> str:
-        return 'test-model'
-
+            return _TextResult(None)
 
-@dataclass
-class TestAgentModel(AgentModel):
-    """Implementation of `AgentModel` for testing purposes."""
-
-    # NOTE: Avoid test discovery by pytest.
-    __test__ = False
-
-    tool_calls: list[tuple[str, ToolDefinition]]
-    # left means the text is plain text; right means it's a function call
-    result: _TextResult | _FunctionToolResult
-    result_tools: list[ToolDefinition]
-    seed: int
-    model_name: str = 'test'
-
-    async def request(
-        self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> tuple[ModelResponse, Usage]:
-        model_response = self._request(messages, model_settings)
-        usage = _estimate_usage([*messages, model_response])
-        return model_response, usage
-
-    @asynccontextmanager
-    async def request_stream(
-        self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> AsyncIterator[StreamedResponse]:
-        model_response = self._request(messages, model_settings)
-        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()
+    def _request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> ModelResponse:
+        tool_calls = self._get_tool_calls(model_request_parameters)
+        result = self._get_result(model_request_parameters)
+        result_tools = model_request_parameters.result_tools
 
-    def _request(self, messages: list[ModelMessage], model_settings: ModelSettings | None) -> ModelResponse:
         # 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):
+        if tool_calls and not any(isinstance(m, ModelResponse) for m in messages):
             return ModelResponse(
-                parts=[ToolCallPart(name, self.gen_tool_args(args)) for name, args in self.tool_calls],
-                model_name=self.model_name,
+                parts=[ToolCallPart(name, self.gen_tool_args(args)) for name, args in tool_calls],
+                model_name=self._model_name,
             )
 
         if messages:
@@ -179,28 +169,26 @@ class TestAgentModel(AgentModel):
                 # Handle retries for both function tools and result tools
                 # Check function tools first
                 retry_parts: list[ModelResponsePart] = [
-                    ToolCallPart(name, self.gen_tool_args(args))
-                    for name, args in self.tool_calls
-                    if name in new_retry_names
+                    ToolCallPart(name, self.gen_tool_args(args)) for name, args in tool_calls if name in new_retry_names
                 ]
                 # Check result tools
-                if self.result_tools:
+                if result_tools:
                     retry_parts.extend(
                         [
                             ToolCallPart(
                                 tool.name,
-                                self.result.value
-                                if isinstance(self.result, _FunctionToolResult) and self.result.value is not None
+                                result.value
+                                if isinstance(result, _FunctionToolResult) and result.value is not None
                                 else self.gen_tool_args(tool),
                             )
-                            for tool in self.result_tools
+                            for tool in result_tools
                             if tool.name in new_retry_names
                         ]
                     )
-                return ModelResponse(parts=retry_parts, model_name=self.model_name)
+                return ModelResponse(parts=retry_parts, model_name=self._model_name)
 
-        if isinstance(self.result, _TextResult):
-            if (response_text := self.result.value) is None:
+        if isinstance(result, _TextResult):
+            if (response_text := result.value) is None:
                 # build up details of tool responses
                 output: dict[str, Any] = {}
                 for message in messages:
@@ -210,23 +198,23 @@ class TestAgentModel(AgentModel):
                                 output[part.tool_name] = part.content
                 if output:
                     return ModelResponse(
-                        parts=[TextPart(pydantic_core.to_json(output).decode())], model_name=self.model_name
+                        parts=[TextPart(pydantic_core.to_json(output).decode())], model_name=self._model_name
                     )
                 else:
-                    return ModelResponse(parts=[TextPart('success (no tool calls)')], model_name=self.model_name)
+                    return ModelResponse(parts=[TextPart('success (no tool calls)')], model_name=self._model_name)
             else:
-                return ModelResponse(parts=[TextPart(response_text)], model_name=self.model_name)
+                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.value
-            result_tool = self.result_tools[self.seed % len(self.result_tools)]
+            assert result_tools, 'No result tools provided'
+            custom_result_args = result.value
+            result_tool = result_tools[self.seed % len(result_tools)]
             if custom_result_args is not None:
                 return ModelResponse(
-                    parts=[ToolCallPart(result_tool.name, custom_result_args)], model_name=self.model_name
+                    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(result_tool.name, response_args)], model_name=self.model_name)
+                return ModelResponse(parts=[ToolCallPart(result_tool.name, response_args)], model_name=self._model_name)
 
 
 @dataclass