pydantic-ai-slim 0.0.17__py3-none-any.whl → 0.0.19__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, '
@@ -146,14 +143,14 @@ class OpenAIAgentModel(AgentModel):
 
     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)
@@ -214,33 +211,14 @@ class OpenAIAgentModel(AgentModel):
         return ModelResponse(items, timestamp=timestamp)
 
     @staticmethod
-    async def _process_streamed_response(response: AsyncStream[ChatCompletionChunk]) -> EitherStreamedResponse:
+    async def _process_streamed_response(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(peekable_response, datetime.fromtimestamp(first_chunk.created, tz=timezone.utc))
 
     @classmethod
     def _map_message(cls, message: ModelMessage) -> Iterable[chat.ChatCompletionMessageParam]:
@@ -295,88 +273,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 +315,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
@@ -141,25 +140,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_response, messages)
 
     def gen_tool_args(self, tool_def: ToolDefinition) -> Any:
         return _JsonSchemaTestData(tool_def.parameters_json_schema, self.seed).generate()
@@ -223,58 +206,37 @@ class TestAgentModel(AgentModel):
 
 
 @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 +396,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

@@ -164,7 +164,7 @@ class VertexAIModel(Model):
         return url, auth
 
     def name(self) -> str:
-        return f'vertexai:{self.model_name}'
+        return f'google-vertex:{self.model_name}'
 
 
 # pyright: reportUnknownMemberType=false