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

pydantic_ai/models/gemini.py

@@ -2,24 +2,25 @@ from __future__ import annotations as _annotations
 
 import os
 import re
-from collections.abc import AsyncIterator, Iterable, Sequence
+from collections.abc import AsyncIterator, Sequence
 from contextlib import asynccontextmanager
 from copy import deepcopy
 from dataclasses import dataclass, field
 from datetime import datetime
 from typing import Annotated, Any, Literal, Protocol, Union
+from uuid import uuid4
 
 import pydantic
-import pydantic_core
 from httpx import USE_CLIENT_DEFAULT, AsyncClient as AsyncHTTPClient, Response as HTTPResponse
-from typing_extensions import NotRequired, TypedDict, TypeGuard, assert_never
+from typing_extensions import NotRequired, TypedDict, assert_never
 
-from .. import UnexpectedModelBehavior, _utils, exceptions, result
+from .. import UnexpectedModelBehavior, _utils, exceptions, usage
 from ..messages import (
     ModelMessage,
     ModelRequest,
     ModelResponse,
     ModelResponsePart,
+    ModelResponseStreamEvent,
     RetryPromptPart,
     SystemPromptPart,
     TextPart,
@@ -31,10 +32,8 @@ from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
     AgentModel,
-    EitherStreamedResponse,
     Model,
-    StreamStructuredResponse,
-    StreamTextResponse,
+    StreamedResponse,
     cached_async_http_client,
     check_allow_model_requests,
     get_user_agent,
@@ -171,7 +170,7 @@ class GeminiAgentModel(AgentModel):
 
     async def request(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> tuple[ModelResponse, result.Usage]:
+    ) -> tuple[ModelResponse, usage.Usage]:
         async with self._make_request(messages, False, model_settings) as http_response:
             response = _gemini_response_ta.validate_json(await http_response.aread())
         return self._process_response(response), _metadata_as_usage(response)
@@ -179,7 +178,7 @@ class GeminiAgentModel(AgentModel):
     @asynccontextmanager
     async def request_stream(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> AsyncIterator[EitherStreamedResponse]:
+    ) -> AsyncIterator[StreamedResponse]:
         async with self._make_request(messages, True, model_settings) as http_response:
             yield await self._process_streamed_response(http_response)
 
@@ -238,7 +237,7 @@ class GeminiAgentModel(AgentModel):
         return _process_response_from_parts(parts)
 
     @staticmethod
-    async def _process_streamed_response(http_response: HTTPResponse) -> EitherStreamedResponse:
+    async def _process_streamed_response(http_response: HTTPResponse) -> StreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
         aiter_bytes = http_response.aiter_bytes()
         start_response: _GeminiResponse | None = None
@@ -259,11 +258,7 @@ class GeminiAgentModel(AgentModel):
         if start_response is None:
             raise UnexpectedModelBehavior('Streamed response ended without content or tool calls')
 
-        # TODO: Update this once we rework stream responses to be more flexible
-        if _extract_response_parts(start_response).is_left():
-            return GeminiStreamStructuredResponse(_content=content, _stream=aiter_bytes)
-        else:
-            return GeminiStreamTextResponse(_json_content=content, _stream=aiter_bytes)
+        return GeminiStreamedResponse(_content=content, _stream=aiter_bytes)
 
     @classmethod
     def _message_to_gemini_content(
@@ -302,86 +297,69 @@ class GeminiAgentModel(AgentModel):
 
 
 @dataclass
-class GeminiStreamTextResponse(StreamTextResponse):
-    """Implementation of `StreamTextResponse` for the Gemini model."""
-
-    _json_content: bytearray
-    _stream: AsyncIterator[bytes]
-    _position: int = 0
-    _timestamp: datetime = field(default_factory=_utils.now_utc, init=False)
-    _usage: result.Usage = field(default_factory=result.Usage, init=False)
-
-    async def __anext__(self) -> None:
-        chunk = await self._stream.__anext__()
-        self._json_content.extend(chunk)
-
-    def get(self, *, final: bool = False) -> Iterable[str]:
-        if final:
-            all_items = pydantic_core.from_json(self._json_content)
-            new_items = all_items[self._position :]
-            self._position = len(all_items)
-            new_responses = _gemini_streamed_response_ta.validate_python(new_items)
-        else:
-            all_items = pydantic_core.from_json(self._json_content, allow_partial=True)
-            new_items = all_items[self._position : -1]
-            self._position = len(all_items) - 1
-            new_responses = _gemini_streamed_response_ta.validate_python(
-                new_items, experimental_allow_partial='trailing-strings'
-            )
-        for r in new_responses:
-            self._usage += _metadata_as_usage(r)
-            parts = r['candidates'][0]['content']['parts']
-            if _all_text_parts(parts):
-                for part in parts:
-                    yield part['text']
-            else:
-                raise UnexpectedModelBehavior(
-                    'Streamed response with unexpected content, expected all parts to be text'
-                )
-
-    def usage(self) -> result.Usage:
-        return self._usage
-
-    def timestamp(self) -> datetime:
-        return self._timestamp
-
-
-@dataclass
-class GeminiStreamStructuredResponse(StreamStructuredResponse):
-    """Implementation of `StreamStructuredResponse` for the Gemini model."""
+class GeminiStreamedResponse(StreamedResponse):
+    """Implementation of `StreamedResponse` for the Gemini model."""
 
     _content: bytearray
     _stream: AsyncIterator[bytes]
     _timestamp: datetime = field(default_factory=_utils.now_utc, init=False)
-    _usage: result.Usage = field(default_factory=result.Usage, init=False)
-
-    async def __anext__(self) -> None:
-        chunk = await self._stream.__anext__()
-        self._content.extend(chunk)
-
-    def get(self, *, final: bool = False) -> ModelResponse:
-        """Get the `ModelResponse` at this point.
 
-        NOTE: It's not clear how the stream of responses should be combined because Gemini seems to always
-        reply with a single response, when returning a structured data.
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+        async for gemini_response in self._get_gemini_responses():
+            candidate = gemini_response['candidates'][0]
+            gemini_part: _GeminiPartUnion
+            for gemini_part in candidate['content']['parts']:
+                if 'text' in gemini_part:
+                    # Using vendor_part_id=None means we can produce multiple text parts if their deltas are sprinkled
+                    # amongst the tool call deltas
+                    yield self._parts_manager.handle_text_delta(vendor_part_id=None, content=gemini_part['text'])
+
+                elif 'function_call' in gemini_part:
+                    # Here, we assume all function_call parts are complete and don't have deltas.
+                    # We do this by assigning a unique randomly generated "vendor_part_id".
+                    # We need to confirm whether this is actually true, but if it isn't, we can still handle it properly
+                    # it would just be a bit more complicated. And we'd need to confirm the intended semantics.
+                    maybe_event = self._parts_manager.handle_tool_call_delta(
+                        vendor_part_id=uuid4(),
+                        tool_name=gemini_part['function_call']['name'],
+                        args=gemini_part['function_call']['args'],
+                        tool_call_id=None,
+                    )
+                    if maybe_event is not None:
+                        yield maybe_event
+                else:
+                    assert 'function_response' in gemini_part, f'Unexpected part: {gemini_part}'
+
+    async def _get_gemini_responses(self) -> AsyncIterator[_GeminiResponse]:
+        # This method exists to ensure we only yield completed items, so we don't need to worry about
+        # partial gemini responses, which would make everything more complicated
+
+        gemini_responses: list[_GeminiResponse] = []
+        current_gemini_response_index = 0
+        # Right now, there are some circumstances where we will have information that could be yielded sooner than it is
+        # But changing that would make things a lot more complicated.
+        async for chunk in self._stream:
+            self._content.extend(chunk)
+
+            gemini_responses = _gemini_streamed_response_ta.validate_json(
+                self._content,
+                experimental_allow_partial='trailing-strings',
+            )
 
-        I'm therefore assuming that each part contains a complete tool call, and not trying to combine data from
-        separate parts.
-        """
-        responses = _gemini_streamed_response_ta.validate_json(
-            self._content,
-            experimental_allow_partial='off' if final else 'trailing-strings',
-        )
-        combined_parts: list[_GeminiPartUnion] = []
-        self._usage = result.Usage()
-        for r in responses:
+            # The idea: yield only up to the latest response, which might still be partial.
+            # Note that if the latest response is complete, we could yield it immediately, but there's not a good
+            # allow_partial API to determine if the last item in the list is complete.
+            responses_to_yield = gemini_responses[:-1]
+            for r in responses_to_yield[current_gemini_response_index:]:
+                current_gemini_response_index += 1
+                self._usage += _metadata_as_usage(r)
+                yield r
+
+        # Now yield the final response, which should be complete
+        if gemini_responses:
+            r = gemini_responses[-1]
             self._usage += _metadata_as_usage(r)
-            candidate = r['candidates'][0]
-            combined_parts.extend(candidate['content']['parts'])
-        return _process_response_from_parts(combined_parts, timestamp=self._timestamp)
-
-    def usage(self) -> result.Usage:
-        return self._usage
+            yield r
 
     def timestamp(self) -> datetime:
         return self._timestamp
@@ -458,9 +436,14 @@ def _process_response_from_parts(parts: Sequence[_GeminiPartUnion], timestamp: d
     items: list[ModelResponsePart] = []
     for part in parts:
         if 'text' in part:
-            items.append(TextPart(part['text']))
+            items.append(TextPart(content=part['text']))
         elif 'function_call' in part:
-            items.append(ToolCallPart.from_raw_args(part['function_call']['name'], part['function_call']['args']))
+            items.append(
+                ToolCallPart.from_raw_args(
+                    tool_name=part['function_call']['name'],
+                    args=part['function_call']['args'],
+                )
+            )
         elif 'function_response' in part:
             raise exceptions.UnexpectedModelBehavior(
                 f'Unsupported response from Gemini, expected all parts to be function calls or text, got: {part!r}'
@@ -575,35 +558,6 @@ class _GeminiResponse(TypedDict):
     prompt_feedback: NotRequired[Annotated[_GeminiPromptFeedback, pydantic.Field(alias='promptFeedback')]]
 
 
-# TODO: Delete the next three functions once we've reworked streams to be more flexible
-def _extract_response_parts(
-    response: _GeminiResponse,
-) -> _utils.Either[list[_GeminiFunctionCallPart], list[_GeminiTextPart]]:
-    """Extract the parts of the response from the Gemini API.
-
-    Returns Either a list of function calls (Either.left) or a list of text parts (Either.right).
-    """
-    if len(response['candidates']) != 1:
-        raise UnexpectedModelBehavior('Expected exactly one candidate in Gemini response')
-    parts = response['candidates'][0]['content']['parts']
-    if _all_function_call_parts(parts):
-        return _utils.Either(left=parts)
-    elif _all_text_parts(parts):
-        return _utils.Either(right=parts)
-    else:
-        raise exceptions.UnexpectedModelBehavior(
-            f'Unsupported response from Gemini, expected all parts to be function calls or text, got: {parts!r}'
-        )
-
-
-def _all_function_call_parts(parts: list[_GeminiPartUnion]) -> TypeGuard[list[_GeminiFunctionCallPart]]:
-    return all('function_call' in part for part in parts)
-
-
-def _all_text_parts(parts: list[_GeminiPartUnion]) -> TypeGuard[list[_GeminiTextPart]]:
-    return all('text' in part for part in parts)
-
-
 class _GeminiCandidates(TypedDict):
     """See <https://ai.google.dev/api/generate-content#v1beta.Candidate>."""
 
@@ -630,14 +584,14 @@ class _GeminiUsageMetaData(TypedDict, total=False):
     cached_content_token_count: NotRequired[Annotated[int, pydantic.Field(alias='cachedContentTokenCount')]]
 
 
-def _metadata_as_usage(response: _GeminiResponse) -> result.Usage:
+def _metadata_as_usage(response: _GeminiResponse) -> usage.Usage:
     metadata = response.get('usage_metadata')
     if metadata is None:
-        return result.Usage()
+        return usage.Usage()
     details: dict[str, int] = {}
     if cached_content_token_count := metadata.get('cached_content_token_count'):
         details['cached_content_token_count'] = cached_content_token_count
-    return result.Usage(
+    return usage.Usage(
         request_tokens=metadata.get('prompt_token_count', 0),
         response_tokens=metadata.get('candidates_token_count', 0),
         total_tokens=metadata.get('total_token_count', 0),

pydantic_ai/models/groq.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, 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 groq import NOT_GIVEN, AsyncGroq, AsyncStream
     from groq.types import chat
     from groq.types.chat import ChatCompletion, ChatCompletionChunk
-    from groq.types.chat.chat_completion_chunk import ChoiceDeltaToolCall
 except ImportError as _import_error:
     raise ImportError(
         'Please install `groq` to use the Groq model, '
@@ -157,14 +154,14 @@ class GroqAgentModel(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)
@@ -217,38 +214,23 @@ class GroqAgentModel(AgentModel):
         choice = response.choices[0]
         items: list[ModelResponsePart] = []
         if choice.message.content is not None:
-            items.append(TextPart(choice.message.content))
+            items.append(TextPart(content=choice.message.content))
         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))
+                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)
 
     @staticmethod
-    async def _process_streamed_response(response: AsyncStream[ChatCompletionChunk]) -> EitherStreamedResponse:
+    async def _process_streamed_response(response: AsyncStream[ChatCompletionChunk]) -> GroqStreamedResponse:
         """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 GroqStreamTextResponse(delta.content, response, timestamp, start_usage)
-                elif delta.tool_calls is not None:
-                    return GroqStreamStructuredResponse(
-                        response,
-                        {c.index: c for c in delta.tool_calls},
-                        timestamp,
-                        start_usage,
-                    )
+        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))
 
     @classmethod
     def _map_message(cls, message: ModelMessage) -> Iterable[chat.ChatCompletionMessageParam]:
@@ -301,90 +283,36 @@ class GroqAgentModel(AgentModel):
 
 
 @dataclass
-class GroqStreamTextResponse(StreamTextResponse):
-    """Implementation of `StreamTextResponse` for Groq models."""
+class GroqStreamedResponse(StreamedResponse):
+    """Implementation of `StreamedResponse` for Groq models."""
 
-    _first: str | None
-    _response: AsyncStream[ChatCompletionChunk]
+    _response: AsyncIterable[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)
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+        async for chunk in self._response:
+            self._usage += _map_usage(chunk)
 
-    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 GroqStreamStructuredResponse(StreamStructuredResponse):
-    """Implementation of `StreamStructuredResponse` for Groq models."""
-
-    _response: AsyncStream[ChatCompletionChunk]
-    _delta_tool_calls: dict[int, ChoiceDeltaToolCall]
-    _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)
-
-    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)
+
+            # Handle the tool calls
+            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
@@ -398,18 +326,18 @@ def _map_tool_call(t: ToolCallPart) -> chat.ChatCompletionMessageToolCallParam:
     )
 
 
-def _map_usage(completion: ChatCompletionChunk | ChatCompletion) -> result.Usage:
-    usage = None
+def _map_usage(completion: ChatCompletionChunk | ChatCompletion) -> usage.Usage:
+    response_usage = None
     if isinstance(completion, ChatCompletion):
-        usage = completion.usage
+        response_usage = completion.usage
     elif completion.x_groq is not None:
-        usage = completion.x_groq.usage
+        response_usage = completion.x_groq.usage
 
-    if usage is None:
-        return result.Usage()
+    if response_usage is None:
+        return usage.Usage()
 
-    return result.Usage(
-        request_tokens=usage.prompt_tokens,
-        response_tokens=usage.completion_tokens,
-        total_tokens=usage.total_tokens,
+    return usage.Usage(
+        request_tokens=response_usage.prompt_tokens,
+        response_tokens=response_usage.completion_tokens,
+        total_tokens=response_usage.total_tokens,
     )