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

pydantic_ai/models/groq.py

@@ -197,7 +197,7 @@ class GroqAgentModel(AgentModel):
             model=str(self.model_name),
             messages=groq_messages,
             n=1,
-            parallel_tool_calls=True if self.tools else NOT_GIVEN,
+            parallel_tool_calls=model_settings.get('parallel_tool_calls', True if self.tools else NOT_GIVEN),
             tools=self.tools or NOT_GIVEN,
             tool_choice=tool_choice or NOT_GIVEN,
             stream=stream,
@@ -207,8 +207,7 @@ class GroqAgentModel(AgentModel):
             timeout=model_settings.get('timeout', NOT_GIVEN),
         )
 
-    @staticmethod
-    def _process_response(response: chat.ChatCompletion) -> ModelResponse:
+    def _process_response(self, response: chat.ChatCompletion) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
         timestamp = datetime.fromtimestamp(response.created, tz=timezone.utc)
         choice = response.choices[0]
@@ -220,17 +219,20 @@ class GroqAgentModel(AgentModel):
                 items.append(
                     ToolCallPart.from_raw_args(tool_name=c.function.name, args=c.function.arguments, tool_call_id=c.id)
                 )
-        return ModelResponse(items, timestamp=timestamp)
+        return ModelResponse(items, model_name=self.model_name, timestamp=timestamp)
 
-    @staticmethod
-    async def _process_streamed_response(response: AsyncStream[ChatCompletionChunk]) -> GroqStreamedResponse:
+    async def _process_streamed_response(self, response: AsyncStream[ChatCompletionChunk]) -> GroqStreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
         peekable_response = _utils.PeekableAsyncStream(response)
         first_chunk = await peekable_response.peek()
         if isinstance(first_chunk, _utils.Unset):
             raise UnexpectedModelBehavior('Streamed response ended without content or tool calls')
 
-        return GroqStreamedResponse(peekable_response, datetime.fromtimestamp(first_chunk.created, tz=timezone.utc))
+        return GroqStreamedResponse(
+            _response=peekable_response,
+            _model_name=self.model_name,
+            _timestamp=datetime.fromtimestamp(first_chunk.created, tz=timezone.utc),
+        )
 
     @classmethod
     def _map_message(cls, message: ModelMessage) -> Iterable[chat.ChatCompletionMessageParam]:

pydantic_ai/models/mistral.py

@@ -36,6 +36,7 @@ from . import (
     Model,
     StreamedResponse,
     cached_async_http_client,
+    check_allow_model_requests,
 )
 
 try:
@@ -130,6 +131,7 @@ class MistralModel(Model):
         result_tools: list[ToolDefinition],
     ) -> AgentModel:
         """Create an agent model, this is called for each step of an agent run from Pydantic AI call."""
+        check_allow_model_requests()
         return MistralAgentModel(
             self.client,
             self.model_name,
@@ -147,7 +149,7 @@ class MistralAgentModel(AgentModel):
     """Implementation of `AgentModel` for Mistral models."""
 
     client: Mistral
-    model_name: str
+    model_name: MistralModelName
     allow_text_result: bool
     function_tools: list[ToolDefinition]
     result_tools: list[ToolDefinition]
@@ -265,8 +267,7 @@ class MistralAgentModel(AgentModel):
         ]
         return tools if tools else None
 
-    @staticmethod
-    def _process_response(response: MistralChatCompletionResponse) -> ModelResponse:
+    def _process_response(self, response: MistralChatCompletionResponse) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
         assert response.choices, 'Unexpected empty response choice.'
 
@@ -288,10 +289,10 @@ class MistralAgentModel(AgentModel):
                 tool = _map_mistral_to_pydantic_tool_call(tool_call=tool_call)
                 parts.append(tool)
 
-        return ModelResponse(parts, timestamp=timestamp)
+        return ModelResponse(parts, model_name=self.model_name, timestamp=timestamp)
 
-    @staticmethod
     async def _process_streamed_response(
+        self,
         result_tools: list[ToolDefinition],
         response: MistralEventStreamAsync[MistralCompletionEvent],
     ) -> StreamedResponse:
@@ -306,7 +307,12 @@ class MistralAgentModel(AgentModel):
         else:
             timestamp = datetime.now(tz=timezone.utc)
 
-        return MistralStreamedResponse(peekable_response, timestamp, {c.name: c for c in result_tools})
+        return MistralStreamedResponse(
+            _response=peekable_response,
+            _model_name=self.model_name,
+            _timestamp=timestamp,
+            _result_tools={c.name: c for c in result_tools},
+        )
 
     @staticmethod
     def _map_to_mistral_tool_call(t: ToolCallPart) -> MistralToolCall:

pydantic_ai/models/ollama.py

@@ -10,6 +10,7 @@ from . import (
     AgentModel,
     Model,
     cached_async_http_client,
+    check_allow_model_requests,
 )
 
 try:
@@ -25,6 +26,7 @@ from .openai import OpenAIModel
 
 CommonOllamaModelNames = Literal[
     'codellama',
+    'deepseek-r1',
     'gemma',
     'gemma2',
     'llama3',
@@ -110,6 +112,7 @@ class OllamaModel(Model):
         allow_text_result: bool,
         result_tools: list[ToolDefinition],
     ) -> AgentModel:
+        check_allow_model_requests()
         return await self.openai_model.agent_model(
             function_tools=function_tools,
             allow_text_result=allow_text_result,

pydantic_ai/models/openai.py

@@ -51,6 +51,8 @@ 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)
 """
 
+OpenAISystemPromptRole = Literal['system', 'developer', 'user']
+
 
 @dataclass(init=False)
 class OpenAIModel(Model):
@@ -63,6 +65,7 @@ class OpenAIModel(Model):
 
     model_name: OpenAIModelName
     client: AsyncOpenAI = field(repr=False)
+    system_prompt_role: OpenAISystemPromptRole | None = field(default=None)
 
     def __init__(
         self,
@@ -72,6 +75,7 @@ class OpenAIModel(Model):
         api_key: str | None = None,
         openai_client: AsyncOpenAI | None = None,
         http_client: AsyncHTTPClient | None = None,
+        system_prompt_role: OpenAISystemPromptRole | None = None,
     ):
         """Initialize an OpenAI model.
 
@@ -87,6 +91,8 @@ class OpenAIModel(Model):
                 [`AsyncOpenAI`](https://github.com/openai/openai-python?tab=readme-ov-file#async-usage)
                 client to use. If provided, `base_url`, `api_key`, and `http_client` must be `None`.
             http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
+            system_prompt_role: The role to use for the system prompt message. If not provided, defaults to `'system'`.
+                In the future, this may be inferred from the model name.
         """
         self.model_name: OpenAIModelName = model_name
         if openai_client is not None:
@@ -98,6 +104,7 @@ class OpenAIModel(Model):
             self.client = AsyncOpenAI(base_url=base_url, api_key=api_key, http_client=http_client)
         else:
             self.client = AsyncOpenAI(base_url=base_url, api_key=api_key, http_client=cached_async_http_client())
+        self.system_prompt_role = system_prompt_role
 
     async def agent_model(
         self,
@@ -115,6 +122,7 @@ class OpenAIModel(Model):
             self.model_name,
             allow_text_result,
             tools,
+            self.system_prompt_role,
         )
 
     def name(self) -> str:
@@ -140,6 +148,7 @@ class OpenAIAgentModel(AgentModel):
     model_name: OpenAIModelName
     allow_text_result: bool
     tools: list[chat.ChatCompletionToolParam]
+    system_prompt_role: OpenAISystemPromptRole | None
 
     async def request(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
@@ -186,7 +195,7 @@ class OpenAIAgentModel(AgentModel):
             model=self.model_name,
             messages=openai_messages,
             n=1,
-            parallel_tool_calls=True if self.tools else NOT_GIVEN,
+            parallel_tool_calls=model_settings.get('parallel_tool_calls', True if self.tools else NOT_GIVEN),
             tools=self.tools or NOT_GIVEN,
             tool_choice=tool_choice or NOT_GIVEN,
             stream=stream,
@@ -197,8 +206,7 @@ class OpenAIAgentModel(AgentModel):
             timeout=model_settings.get('timeout', NOT_GIVEN),
         )
 
-    @staticmethod
-    def _process_response(response: chat.ChatCompletion) -> ModelResponse:
+    def _process_response(self, response: chat.ChatCompletion) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
         timestamp = datetime.fromtimestamp(response.created, tz=timezone.utc)
         choice = response.choices[0]
@@ -208,23 +216,25 @@ class OpenAIAgentModel(AgentModel):
         if choice.message.tool_calls is not None:
             for c in choice.message.tool_calls:
                 items.append(ToolCallPart.from_raw_args(c.function.name, c.function.arguments, c.id))
-        return ModelResponse(items, timestamp=timestamp)
+        return ModelResponse(items, model_name=self.model_name, timestamp=timestamp)
 
-    @staticmethod
-    async def _process_streamed_response(response: AsyncStream[ChatCompletionChunk]) -> OpenAIStreamedResponse:
+    async def _process_streamed_response(self, response: AsyncStream[ChatCompletionChunk]) -> OpenAIStreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
         peekable_response = _utils.PeekableAsyncStream(response)
         first_chunk = await peekable_response.peek()
         if isinstance(first_chunk, _utils.Unset):
             raise UnexpectedModelBehavior('Streamed response ended without content or tool calls')
 
-        return OpenAIStreamedResponse(peekable_response, datetime.fromtimestamp(first_chunk.created, tz=timezone.utc))
+        return OpenAIStreamedResponse(
+            _model_name=self.model_name,
+            _response=peekable_response,
+            _timestamp=datetime.fromtimestamp(first_chunk.created, tz=timezone.utc),
+        )
 
-    @classmethod
-    def _map_message(cls, message: ModelMessage) -> Iterable[chat.ChatCompletionMessageParam]:
+    def _map_message(self, message: ModelMessage) -> Iterable[chat.ChatCompletionMessageParam]:
         """Just maps a `pydantic_ai.Message` to a `openai.types.ChatCompletionMessageParam`."""
         if isinstance(message, ModelRequest):
-            yield from cls._map_user_message(message)
+            yield from self._map_user_message(message)
         elif isinstance(message, ModelResponse):
             texts: list[str] = []
             tool_calls: list[chat.ChatCompletionMessageToolCallParam] = []
@@ -246,11 +256,15 @@ class OpenAIAgentModel(AgentModel):
         else:
             assert_never(message)
 
-    @classmethod
-    def _map_user_message(cls, message: ModelRequest) -> Iterable[chat.ChatCompletionMessageParam]:
+    def _map_user_message(self, message: ModelRequest) -> Iterable[chat.ChatCompletionMessageParam]:
         for part in message.parts:
             if isinstance(part, SystemPromptPart):
-                yield chat.ChatCompletionSystemMessageParam(role='system', content=part.content)
+                if self.system_prompt_role == 'developer':
+                    yield chat.ChatCompletionDeveloperMessageParam(role='developer', content=part.content)
+                elif self.system_prompt_role == 'user':
+                    yield chat.ChatCompletionUserMessageParam(role='user', content=part.content)
+                else:
+                    yield chat.ChatCompletionSystemMessageParam(role='system', content=part.content)
             elif isinstance(part, UserPromptPart):
                 yield chat.ChatCompletionUserMessageParam(role='user', content=part.content)
             elif isinstance(part, ToolReturnPart):

pydantic_ai/models/test.py

@@ -129,6 +129,7 @@ class TestAgentModel(AgentModel):
     result: _utils.Either[str | None, Any | None]
     result_tools: list[ToolDefinition]
     seed: int
+    model_name: str = 'test'
 
     async def request(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
@@ -142,7 +143,7 @@ class TestAgentModel(AgentModel):
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
     ) -> AsyncIterator[StreamedResponse]:
         model_response = self._request(messages, model_settings)
-        yield TestStreamedResponse(model_response, messages)
+        yield TestStreamedResponse(_model_name=self.model_name, _structured_response=model_response, _messages=messages)
 
     def gen_tool_args(self, tool_def: ToolDefinition) -> Any:
         return _JsonSchemaTestData(tool_def.parameters_json_schema, self.seed).generate()
@@ -151,7 +152,8 @@ class TestAgentModel(AgentModel):
         # if there are tools, the first thing we want to do is call all of them
         if self.tool_calls and not any(isinstance(m, ModelResponse) for m in messages):
             return ModelResponse(
-                parts=[ToolCallPart.from_raw_args(name, self.gen_tool_args(args)) for name, args in self.tool_calls]
+                parts=[ToolCallPart.from_raw_args(name, self.gen_tool_args(args)) for name, args in self.tool_calls],
+                model_name=self.model_name,
             )
 
         if messages:
@@ -177,7 +179,7 @@ class TestAgentModel(AgentModel):
                             if tool.name in new_retry_names
                         ]
                     )
-                return ModelResponse(parts=retry_parts)
+                return ModelResponse(parts=retry_parts, model_name=self.model_name)
 
         if response_text := self.result.left:
             if response_text.value is None:
@@ -189,20 +191,26 @@ class TestAgentModel(AgentModel):
                             if isinstance(part, ToolReturnPart):
                                 output[part.tool_name] = part.content
                 if output:
-                    return ModelResponse.from_text(pydantic_core.to_json(output).decode())
+                    return ModelResponse(
+                        parts=[TextPart(pydantic_core.to_json(output).decode())], model_name=self.model_name
+                    )
                 else:
-                    return ModelResponse.from_text('success (no tool calls)')
+                    return ModelResponse(parts=[TextPart('success (no tool calls)')], model_name=self.model_name)
             else:
-                return ModelResponse.from_text(response_text.value)
+                return ModelResponse(parts=[TextPart(response_text.value)], model_name=self.model_name)
         else:
             assert self.result_tools, 'No result tools provided'
             custom_result_args = self.result.right
             result_tool = self.result_tools[self.seed % len(self.result_tools)]
             if custom_result_args is not None:
-                return ModelResponse(parts=[ToolCallPart.from_raw_args(result_tool.name, custom_result_args)])
+                return ModelResponse(
+                    parts=[ToolCallPart.from_raw_args(result_tool.name, custom_result_args)], model_name=self.model_name
+                )
             else:
                 response_args = self.gen_tool_args(result_tool)
-                return ModelResponse(parts=[ToolCallPart.from_raw_args(result_tool.name, response_args)])
+                return ModelResponse(
+                    parts=[ToolCallPart.from_raw_args(result_tool.name, response_args)], model_name=self.model_name
+                )
 
 
 @dataclass

pydantic_ai/models/vertexai.py

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

pydantic_ai/result.py

@@ -11,35 +11,49 @@ import logfire_api
 from typing_extensions import TypeVar
 
 from . import _result, _utils, exceptions, messages as _messages, models
-from .tools import AgentDeps, RunContext
+from .tools import AgentDepsT, RunContext
 from .usage import Usage, UsageLimits
 
-__all__ = 'ResultData', 'ResultValidatorFunc', 'RunResult', 'StreamedRunResult'
+__all__ = 'ResultDataT', 'ResultDataT_inv', 'ResultValidatorFunc', 'RunResult', 'StreamedRunResult'
 
 
-ResultData = TypeVar('ResultData', default=str)
-"""Type variable for the result data of a run."""
+T = TypeVar('T')
+"""An invariant TypeVar."""
+ResultDataT_inv = TypeVar('ResultDataT_inv', default=str)
+"""
+An invariant type variable for the result data of a model.
+
+We need to use an invariant typevar for `ResultValidator` and `ResultValidatorFunc` because the result data type is used
+in both the input and output of a `ResultValidatorFunc`. This can theoretically lead to some issues assuming that types
+possessing ResultValidator's are covariant in the result data type, but in practice this is rarely an issue, and
+changing it would have negative consequences for the ergonomics of the library.
+
+At some point, it may make sense to change the input to ResultValidatorFunc to be `Any` or `object` as doing that would
+resolve these potential variance issues.
+"""
+ResultDataT = TypeVar('ResultDataT', default=str, covariant=True)
+"""Covariant type variable for the result data type of a run."""
 
 ResultValidatorFunc = Union[
-    Callable[[RunContext[AgentDeps], ResultData], ResultData],
-    Callable[[RunContext[AgentDeps], ResultData], Awaitable[ResultData]],
-    Callable[[ResultData], ResultData],
-    Callable[[ResultData], Awaitable[ResultData]],
+    Callable[[RunContext[AgentDepsT], ResultDataT_inv], ResultDataT_inv],
+    Callable[[RunContext[AgentDepsT], ResultDataT_inv], Awaitable[ResultDataT_inv]],
+    Callable[[ResultDataT_inv], ResultDataT_inv],
+    Callable[[ResultDataT_inv], Awaitable[ResultDataT_inv]],
 ]
 """
-A function that always takes `ResultData` and returns `ResultData` and:
+A function that always takes and returns the same type of data (which is the result type of an agent run), and:
 
 * may or may not take [`RunContext`][pydantic_ai.tools.RunContext] as a first argument
 * may or may not be async
 
-Usage `ResultValidatorFunc[AgentDeps, ResultData]`.
+Usage `ResultValidatorFunc[AgentDeps, T]`.
 """
 
 _logfire = logfire_api.Logfire(otel_scope='pydantic-ai')
 
 
 @dataclass
-class _BaseRunResult(ABC, Generic[ResultData]):
+class _BaseRunResult(ABC, Generic[ResultDataT]):
     """Base type for results.
 
     You should not import or use this type directly, instead use its subclasses `RunResult` and `StreamedRunResult`.
@@ -119,10 +133,10 @@ class _BaseRunResult(ABC, Generic[ResultData]):
 
 
 @dataclass
-class RunResult(_BaseRunResult[ResultData]):
+class RunResult(_BaseRunResult[ResultDataT]):
     """Result of a non-streamed run."""
 
-    data: ResultData
+    data: ResultDataT
     """Data from the final response in the run."""
     _result_tool_name: str | None
     _usage: Usage
@@ -165,14 +179,14 @@ class RunResult(_BaseRunResult[ResultData]):
 
 
 @dataclass
-class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultData]):
+class StreamedRunResult(_BaseRunResult[ResultDataT], Generic[AgentDepsT, ResultDataT]):
     """Result of a streamed run that returns structured data via a tool call."""
 
     _usage_limits: UsageLimits | None
     _stream_response: models.StreamedResponse
-    _result_schema: _result.ResultSchema[ResultData] | None
-    _run_ctx: RunContext[AgentDeps]
-    _result_validators: list[_result.ResultValidator[AgentDeps, ResultData]]
+    _result_schema: _result.ResultSchema[ResultDataT] | None
+    _run_ctx: RunContext[AgentDepsT]
+    _result_validators: list[_result.ResultValidator[AgentDepsT, ResultDataT]]
     _result_tool_name: str | None
     _on_complete: Callable[[], Awaitable[None]]
     is_complete: bool = field(default=False, init=False)
@@ -185,7 +199,7 @@ class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultDat
     [`get_data`][pydantic_ai.result.StreamedRunResult.get_data] completes.
     """
 
-    async def stream(self, *, debounce_by: float | None = 0.1) -> AsyncIterator[ResultData]:
+    async def stream(self, *, debounce_by: float | None = 0.1) -> AsyncIterator[ResultDataT]:
         """Stream the response as an async iterable.
 
         The pydantic validator for structured data will be called in
@@ -269,7 +283,12 @@ class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultDat
                     yield combined_validated_text
 
                 lf_span.set_attribute('combined_text', combined_validated_text)
-                await self._marked_completed(_messages.ModelResponse.from_text(combined_validated_text))
+                await self._marked_completed(
+                    _messages.ModelResponse(
+                        parts=[_messages.TextPart(combined_validated_text)],
+                        model_name=self._stream_response.model_name(),
+                    )
+                )
 
     async def stream_structured(
         self, *, debounce_by: float | None = 0.1
@@ -306,7 +325,7 @@ class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultDat
                 lf_span.set_attribute('structured_response', msg)
                 await self._marked_completed(msg)
 
-    async def get_data(self) -> ResultData:
+    async def get_data(self) -> ResultDataT:
         """Stream the whole response, validate and return it."""
         usage_checking_stream = _get_usage_checking_stream_response(
             self._stream_response, self._usage_limits, self.usage
@@ -332,7 +351,7 @@ class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultDat
 
     async def validate_structured_result(
         self, message: _messages.ModelResponse, *, allow_partial: bool = False
-    ) -> ResultData:
+    ) -> ResultDataT:
         """Validate a structured result message."""
         if self._result_schema is not None and self._result_tool_name is not None:
             match = self._result_schema.find_named_tool(message.parts, self._result_tool_name)
@@ -351,17 +370,17 @@ class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultDat
             text = '\n\n'.join(x.content for x in message.parts if isinstance(x, _messages.TextPart))
             for validator in self._result_validators:
                 text = await validator.validate(
-                    text,  # pyright: ignore[reportArgumentType]
+                    text,
                     None,
                     self._run_ctx,
                 )
-            # Since there is no result tool, we can assume that str is compatible with ResultData
-            return cast(ResultData, text)
+            # Since there is no result tool, we can assume that str is compatible with ResultDataT
+            return cast(ResultDataT, text)
 
     async def _validate_text_result(self, text: str) -> str:
         for validator in self._result_validators:
-            text = await validator.validate(  # pyright: ignore[reportAssignmentType]
-                text,  # pyright: ignore[reportArgumentType]
+            text = await validator.validate(
+                text,
                 None,
                 self._run_ctx,
             )

pydantic_ai/settings.py

@@ -12,7 +12,8 @@ if TYPE_CHECKING:
 class ModelSettings(TypedDict, total=False):
     """Settings to configure an LLM.
 
-    Here we include only settings which apply to multiple models / model providers.
+    Here we include only settings which apply to multiple models / model providers,
+    though not all of these settings are supported by all models.
     """
 
     max_tokens: int
@@ -24,6 +25,8 @@ class ModelSettings(TypedDict, total=False):
     * Anthropic
     * OpenAI
     * Groq
+    * Cohere
+    * Mistral
     """
 
     temperature: float
@@ -40,6 +43,8 @@ class ModelSettings(TypedDict, total=False):
     * Anthropic
     * OpenAI
     * Groq
+    * Cohere
+    * Mistral
     """
 
     top_p: float
@@ -55,6 +60,8 @@ class ModelSettings(TypedDict, total=False):
     * Anthropic
     * OpenAI
     * Groq
+    * Cohere
+    * Mistral
     """
 
     timeout: float | Timeout
@@ -66,6 +73,16 @@ class ModelSettings(TypedDict, total=False):
     * Anthropic
     * OpenAI
     * Groq
+    * Mistral
+    """
+
+    parallel_tool_calls: bool
+    """Whether to allow parallel tool calls.
+
+    Supported by:
+    * OpenAI
+    * Groq
+    * Anthropic
     """