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

pydantic_ai/models/openai.py

@@ -1,6 +1,6 @@
 from __future__ import annotations as _annotations
 
-from collections.abc import AsyncIterator, Iterable
+from collections.abc import AsyncIterable, AsyncIterator, Iterable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
@@ -10,13 +10,14 @@ from typing import Literal, Union, overload
 from httpx import AsyncClient as AsyncHTTPClient
 from typing_extensions import assert_never
 
-from .. import UnexpectedModelBehavior, _utils, result
+from .. import UnexpectedModelBehavior, _utils, usage
 from .._utils import guard_tool_call_id as _guard_tool_call_id
 from ..messages import (
     ModelMessage,
     ModelRequest,
     ModelResponse,
     ModelResponsePart,
+    ModelResponseStreamEvent,
     RetryPromptPart,
     SystemPromptPart,
     TextPart,
@@ -24,15 +25,12 @@ from ..messages import (
     ToolReturnPart,
     UserPromptPart,
 )
-from ..result import Usage
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
     AgentModel,
-    EitherStreamedResponse,
     Model,
-    StreamStructuredResponse,
-    StreamTextResponse,
+    StreamedResponse,
     cached_async_http_client,
     check_allow_model_requests,
 )
@@ -41,7 +39,6 @@ try:
     from openai import NOT_GIVEN, AsyncOpenAI, AsyncStream
     from openai.types import ChatModel, chat
     from openai.types.chat import ChatCompletionChunk
-    from openai.types.chat.chat_completion_chunk import ChoiceDeltaToolCall
 except ImportError as _import_error:
     raise ImportError(
         'Please install `openai` to use the OpenAI model, '
@@ -54,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):
@@ -66,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,
@@ -75,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.
 
@@ -90,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:
@@ -101,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,
@@ -118,6 +122,7 @@ class OpenAIModel(Model):
             self.model_name,
             allow_text_result,
             tools,
+            self.system_prompt_role,
         )
 
     def name(self) -> str:
@@ -143,17 +148,18 @@ 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
-    ) -> tuple[ModelResponse, result.Usage]:
+    ) -> tuple[ModelResponse, usage.Usage]:
         response = await self._completions_create(messages, False, model_settings)
         return self._process_response(response), _map_usage(response)
 
     @asynccontextmanager
     async def request_stream(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> AsyncIterator[EitherStreamedResponse]:
+    ) -> AsyncIterator[StreamedResponse]:
         response = await self._completions_create(messages, True, model_settings)
         async with response:
             yield await self._process_streamed_response(response)
@@ -189,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,
@@ -200,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]
@@ -211,42 +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]) -> EitherStreamedResponse:
+    async def _process_streamed_response(self, response: AsyncStream[ChatCompletionChunk]) -> OpenAIStreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
-        timestamp: datetime | None = None
-        start_usage = Usage()
-        # the first chunk may contain enough information so we iterate until we get either `tool_calls` or `content`
-        while True:
-            try:
-                chunk = await response.__anext__()
-            except StopAsyncIteration as e:
-                raise UnexpectedModelBehavior('Streamed response ended without content or tool calls') from e
-
-            timestamp = timestamp or datetime.fromtimestamp(chunk.created, tz=timezone.utc)
-            start_usage += _map_usage(chunk)
-
-            if chunk.choices:
-                delta = chunk.choices[0].delta
-
-                if delta.content is not None:
-                    return OpenAIStreamTextResponse(delta.content, response, timestamp, start_usage)
-                elif delta.tool_calls is not None:
-                    return OpenAIStreamStructuredResponse(
-                        response,
-                        {c.index: c for c in delta.tool_calls},
-                        timestamp,
-                        start_usage,
-                    )
-                # else continue until we get either delta.content or delta.tool_calls
+        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(
+            _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] = []
@@ -268,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):
@@ -295,88 +287,35 @@ class OpenAIAgentModel(AgentModel):
 
 
 @dataclass
-class OpenAIStreamTextResponse(StreamTextResponse):
-    """Implementation of `StreamTextResponse` for OpenAI models."""
-
-    _first: str | None
-    _response: AsyncStream[ChatCompletionChunk]
-    _timestamp: datetime
-    _usage: result.Usage
-    _buffer: list[str] = field(default_factory=list, init=False)
-
-    async def __anext__(self) -> None:
-        if self._first is not None:
-            self._buffer.append(self._first)
-            self._first = None
-            return None
-
-        chunk = await self._response.__anext__()
-        self._usage += _map_usage(chunk)
-        try:
-            choice = chunk.choices[0]
-        except IndexError:
-            raise StopAsyncIteration()
-
-        # we don't raise StopAsyncIteration on the last chunk because usage comes after this
-        if choice.finish_reason is None:
-            assert choice.delta.content is not None, f'Expected delta with content, invalid chunk: {chunk!r}'
-        if choice.delta.content is not None:
-            self._buffer.append(choice.delta.content)
-
-    def get(self, *, final: bool = False) -> Iterable[str]:
-        yield from self._buffer
-        self._buffer.clear()
-
-    def usage(self) -> Usage:
-        return self._usage
-
-    def timestamp(self) -> datetime:
-        return self._timestamp
-
-
-@dataclass
-class OpenAIStreamStructuredResponse(StreamStructuredResponse):
-    """Implementation of `StreamStructuredResponse` for OpenAI models."""
+class OpenAIStreamedResponse(StreamedResponse):
+    """Implementation of `StreamedResponse` for OpenAI models."""
 
-    _response: AsyncStream[ChatCompletionChunk]
-    _delta_tool_calls: dict[int, ChoiceDeltaToolCall]
+    _response: AsyncIterable[ChatCompletionChunk]
     _timestamp: datetime
-    _usage: result.Usage
-
-    async def __anext__(self) -> None:
-        chunk = await self._response.__anext__()
-        self._usage += _map_usage(chunk)
-        try:
-            choice = chunk.choices[0]
-        except IndexError:
-            raise StopAsyncIteration()
-
-        if choice.finish_reason is not None:
-            raise StopAsyncIteration()
-
-        assert choice.delta.content is None, f'Expected tool calls, got content instead, invalid chunk: {chunk!r}'
-
-        for new in choice.delta.tool_calls or []:
-            if current := self._delta_tool_calls.get(new.index):
-                if current.function is None:
-                    current.function = new.function
-                elif new.function is not None:
-                    current.function.name = _utils.add_optional(current.function.name, new.function.name)
-                    current.function.arguments = _utils.add_optional(current.function.arguments, new.function.arguments)
-            else:
-                self._delta_tool_calls[new.index] = new
-
-    def get(self, *, final: bool = False) -> ModelResponse:
-        items: list[ModelResponsePart] = []
-        for c in self._delta_tool_calls.values():
-            if f := c.function:
-                if f.name is not None and f.arguments is not None:
-                    items.append(ToolCallPart.from_raw_args(f.name, f.arguments, c.id))
 
-        return ModelResponse(items, timestamp=self._timestamp)
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+        async for chunk in self._response:
+            self._usage += _map_usage(chunk)
 
-    def usage(self) -> Usage:
-        return self._usage
+            try:
+                choice = chunk.choices[0]
+            except IndexError:
+                continue
+
+            # Handle the text part of the response
+            content = choice.delta.content
+            if content is not None:
+                yield self._parts_manager.handle_text_delta(vendor_part_id='content', content=content)
+
+            for dtc in choice.delta.tool_calls or []:
+                maybe_event = self._parts_manager.handle_tool_call_delta(
+                    vendor_part_id=dtc.index,
+                    tool_name=dtc.function and dtc.function.name,
+                    args=dtc.function and dtc.function.arguments,
+                    tool_call_id=dtc.id,
+                )
+                if maybe_event is not None:
+                    yield maybe_event
 
     def timestamp(self) -> datetime:
         return self._timestamp
@@ -390,19 +329,19 @@ def _map_tool_call(t: ToolCallPart) -> chat.ChatCompletionMessageToolCallParam:
     )
 
 
-def _map_usage(response: chat.ChatCompletion | ChatCompletionChunk) -> result.Usage:
-    usage = response.usage
-    if usage is None:
-        return result.Usage()
+def _map_usage(response: chat.ChatCompletion | ChatCompletionChunk) -> usage.Usage:
+    response_usage = response.usage
+    if response_usage is None:
+        return usage.Usage()
     else:
         details: dict[str, int] = {}
-        if usage.completion_tokens_details is not None:
-            details.update(usage.completion_tokens_details.model_dump(exclude_none=True))
-        if usage.prompt_tokens_details is not None:
-            details.update(usage.prompt_tokens_details.model_dump(exclude_none=True))
-        return result.Usage(
-            request_tokens=usage.prompt_tokens,
-            response_tokens=usage.completion_tokens,
-            total_tokens=usage.total_tokens,
+        if response_usage.completion_tokens_details is not None:
+            details.update(response_usage.completion_tokens_details.model_dump(exclude_none=True))
+        if response_usage.prompt_tokens_details is not None:
+            details.update(response_usage.prompt_tokens_details.model_dump(exclude_none=True))
+        return usage.Usage(
+            request_tokens=response_usage.prompt_tokens,
+            response_tokens=response_usage.completion_tokens,
+            total_tokens=response_usage.total_tokens,
             details=details,
         )

pydantic_ai/models/test.py

@@ -2,21 +2,22 @@ from __future__ import annotations as _annotations
 
 import re
 import string
-from collections.abc import AsyncIterator, Iterable, Iterator
+from collections.abc import AsyncIterator, Iterable
 from contextlib import asynccontextmanager
-from dataclasses import dataclass, field
+from dataclasses import InitVar, dataclass, field
 from datetime import date, datetime, timedelta
 from typing import Any, Literal
 
 import pydantic_core
-from typing_extensions import assert_never
 
 from .. import _utils
 from ..messages import (
+    ArgsJson,
     ModelMessage,
     ModelRequest,
     ModelResponse,
     ModelResponsePart,
+    ModelResponseStreamEvent,
     RetryPromptPart,
     TextPart,
     ToolCallPart,
@@ -27,12 +28,10 @@ from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
     AgentModel,
-    EitherStreamedResponse,
     Model,
-    StreamStructuredResponse,
-    StreamTextResponse,
+    StreamedResponse,
 )
-from .function import _estimate_string_usage, _estimate_usage  # pyright: ignore[reportPrivateUsage]
+from .function import _estimate_string_tokens, _estimate_usage  # pyright: ignore[reportPrivateUsage]
 
 
 @dataclass
@@ -130,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
@@ -141,25 +141,9 @@ class TestAgentModel(AgentModel):
     @asynccontextmanager
     async def request_stream(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> AsyncIterator[EitherStreamedResponse]:
-        msg = self._request(messages, model_settings)
-        usage = _estimate_usage(messages)
-
-        # TODO: Rework this once we make StreamTextResponse more general
-        texts: list[str] = []
-        tool_calls: list[ToolCallPart] = []
-        for item in msg.parts:
-            if isinstance(item, TextPart):
-                texts.append(item.content)
-            elif isinstance(item, ToolCallPart):
-                tool_calls.append(item)
-            else:
-                assert_never(item)
-
-        if texts:
-            yield TestStreamTextResponse('\n\n'.join(texts), usage)
-        else:
-            yield TestStreamStructuredResponse(msg, usage)
+    ) -> 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()
@@ -168,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:
@@ -194,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:
@@ -206,75 +191,60 @@ 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
-class TestStreamTextResponse(StreamTextResponse):
-    """A text response that streams test data."""
-
-    _text: str
-    _usage: Usage
-    _iter: Iterator[str] = field(init=False)
-    _timestamp: datetime = field(default_factory=_utils.now_utc)
-    _buffer: list[str] = field(default_factory=list, init=False)
-
-    def __post_init__(self):
-        *words, last_word = self._text.split(' ')
-        words = [f'{word} ' for word in words]
-        words.append(last_word)
-        if len(words) == 1 and len(self._text) > 2:
-            mid = len(self._text) // 2
-            words = [self._text[:mid], self._text[mid:]]
-        self._iter = iter(words)
-
-    async def __anext__(self) -> None:
-        next_str = _utils.sync_anext(self._iter)
-        response_tokens = _estimate_string_usage(next_str)
-        self._usage += Usage(response_tokens=response_tokens, total_tokens=response_tokens)
-        self._buffer.append(next_str)
-
-    def get(self, *, final: bool = False) -> Iterable[str]:
-        yield from self._buffer
-        self._buffer.clear()
-
-    def usage(self) -> Usage:
-        return self._usage
-
-    def timestamp(self) -> datetime:
-        return self._timestamp
-
-
-@dataclass
-class TestStreamStructuredResponse(StreamStructuredResponse):
+class TestStreamedResponse(StreamedResponse):
     """A structured response that streams test data."""
 
     _structured_response: ModelResponse
-    _usage: Usage
-    _iter: Iterator[None] = field(default_factory=lambda: iter([None]))
-    _timestamp: datetime = field(default_factory=_utils.now_utc, init=False)
-
-    async def __anext__(self) -> None:
-        return _utils.sync_anext(self._iter)
+    _messages: InitVar[Iterable[ModelMessage]]
 
-    def get(self, *, final: bool = False) -> ModelResponse:
-        return self._structured_response
+    _timestamp: datetime = field(default_factory=_utils.now_utc, init=False)
 
-    def usage(self) -> Usage:
-        return self._usage
+    def __post_init__(self, _messages: Iterable[ModelMessage]):
+        self._usage = _estimate_usage(_messages)
+
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+        for i, part in enumerate(self._structured_response.parts):
+            if isinstance(part, TextPart):
+                text = part.content
+                *words, last_word = text.split(' ')
+                words = [f'{word} ' for word in words]
+                words.append(last_word)
+                if len(words) == 1 and len(text) > 2:
+                    mid = len(text) // 2
+                    words = [text[:mid], text[mid:]]
+                self._usage += _get_string_usage('')
+                yield self._parts_manager.handle_text_delta(vendor_part_id=i, content='')
+                for word in words:
+                    self._usage += _get_string_usage(word)
+                    yield self._parts_manager.handle_text_delta(vendor_part_id=i, content=word)
+            else:
+                args = part.args.args_json if isinstance(part.args, ArgsJson) else part.args.args_dict
+                yield self._parts_manager.handle_tool_call_part(
+                    vendor_part_id=i, tool_name=part.tool_name, args=args, tool_call_id=part.tool_call_id
+                )
 
     def timestamp(self) -> datetime:
         return self._timestamp
@@ -434,3 +404,8 @@ class _JsonSchemaTestData:
             rem //= chars
         s += _chars[self.seed % chars]
         return s
+
+
+def _get_string_usage(text: str) -> Usage:
+    response_tokens = _estimate_string_tokens(text)
+    return Usage(response_tokens=response_tokens, total_tokens=response_tokens)

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,