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

pydantic_ai/models/mistral.py

@@ -1,7 +1,7 @@
 from __future__ import annotations as _annotations
 
 import os
-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
@@ -12,7 +12,7 @@ import pydantic_core
 from httpx import AsyncClient as AsyncHTTPClient, Timeout
 from typing_extensions import assert_never
 
-from .. import UnexpectedModelBehavior
+from .. import UnexpectedModelBehavior, _utils
 from .._utils import now_utc as _now_utc
 from ..messages import (
     ArgsJson,
@@ -20,6 +20,7 @@ from ..messages import (
     ModelRequest,
     ModelResponse,
     ModelResponsePart,
+    ModelResponseStreamEvent,
     RetryPromptPart,
     SystemPromptPart,
     TextPart,
@@ -32,10 +33,8 @@ from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
     AgentModel,
-    EitherStreamedResponse,
     Model,
-    StreamStructuredResponse,
-    StreamTextResponse,
+    StreamedResponse,
     cached_async_http_client,
 )
 
@@ -164,7 +163,7 @@ class MistralAgentModel(AgentModel):
     @asynccontextmanager
     async def request_stream(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> AsyncIterator[EitherStreamedResponse]:
+    ) -> AsyncIterator[StreamedResponse]:
         """Make a streaming request to the model from Pydantic AI call."""
         response = await self._stream_completions_create(messages, model_settings)
         async with response:
@@ -282,11 +281,11 @@ class MistralAgentModel(AgentModel):
 
         parts: list[ModelResponsePart] = []
         if text := _map_content(content):
-            parts.append(TextPart(text))
+            parts.append(TextPart(content=text))
 
         if isinstance(tool_calls, list):
             for tool_call in tool_calls:
-                tool = _map_mistral_to_pydantic_tool_call(tool_call)
+                tool = _map_mistral_to_pydantic_tool_call(tool_call=tool_call)
                 parts.append(tool)
 
         return ModelResponse(parts, timestamp=timestamp)
@@ -295,45 +294,19 @@ class MistralAgentModel(AgentModel):
     async def _process_streamed_response(
         result_tools: list[ToolDefinition],
         response: MistralEventStreamAsync[MistralCompletionEvent],
-    ) -> EitherStreamedResponse:
+    ) -> StreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
-        start_usage = 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')
 
-        # Iterate until we get either `tool_calls` or `content` from the first chunk.
-        while True:
-            try:
-                event = await response.__anext__()
-                chunk = event.data
-            except StopAsyncIteration as e:
-                raise UnexpectedModelBehavior('Streamed response ended without content or tool calls') from e
-
-            start_usage += _map_usage(chunk)
-
-            if chunk.created:
-                timestamp = datetime.fromtimestamp(chunk.created, tz=timezone.utc)
-            else:
-                timestamp = _now_utc()
-
-            if chunk.choices:
-                delta = chunk.choices[0].delta
-                content = _map_content(delta.content)
-
-                tool_calls: list[MistralToolCall] | None = None
-                if delta.tool_calls:
-                    tool_calls = delta.tool_calls
-
-                if tool_calls or content and result_tools:
-                    return MistralStreamStructuredResponse(
-                        {c.id if c.id else 'null': c for c in tool_calls or []},
-                        {c.name: c for c in result_tools},
-                        response,
-                        content,
-                        timestamp,
-                        start_usage,
-                    )
+        if first_chunk.data.created:
+            timestamp = datetime.fromtimestamp(first_chunk.data.created, tz=timezone.utc)
+        else:
+            timestamp = datetime.now(tz=timezone.utc)
 
-                elif content:
-                    return MistralStreamTextResponse(content, response, timestamp, start_usage)
+        return MistralStreamedResponse(peekable_response, timestamp, {c.name: c for c in result_tools})
 
     @staticmethod
     def _map_to_mistral_tool_call(t: ToolCallPart) -> MistralToolCall:
@@ -467,108 +440,73 @@ class MistralAgentModel(AgentModel):
             assert_never(message)
 
 
-@dataclass
-class MistralStreamTextResponse(StreamTextResponse):
-    """Implementation of `StreamTextResponse` for Mistral models."""
-
-    _first: str | None
-    _response: MistralEventStreamAsync[MistralCompletionEvent]
-    _timestamp: datetime
-    _usage: Usage
-    _buffer: list[str] = field(default_factory=list, init=False)
-
-    async def __anext__(self) -> None:
-        if self._first is not None and len(self._first) > 0:
-            self._buffer.append(self._first)
-            self._first = None
-            return None
-
-        chunk = await self._response.__anext__()
-        self._usage += _map_usage(chunk.data)
-
-        try:
-            choice = chunk.data.choices[0]
-        except IndexError:
-            raise StopAsyncIteration()
-
-        content = choice.delta.content
-        if choice.finish_reason is None:
-            assert content is not None, f'Expected delta with content, invalid chunk: {chunk!r}'
-
-        if text := _map_content(content):
-            self._buffer.append(text)
-
-    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
+MistralToolCallId = Union[str, None]
 
 
 @dataclass
-class MistralStreamStructuredResponse(StreamStructuredResponse):
-    """Implementation of `StreamStructuredResponse` for Mistral models."""
+class MistralStreamedResponse(StreamedResponse):
+    """Implementation of `StreamedResponse` for Mistral models."""
 
-    _function_tools: dict[str, MistralToolCall]
-    _result_tools: dict[str, ToolDefinition]
-    _response: MistralEventStreamAsync[MistralCompletionEvent]
-    _delta_content: str | None
+    _response: AsyncIterable[MistralCompletionEvent]
     _timestamp: datetime
-    _usage: Usage
-
-    async def __anext__(self) -> None:
-        chunk = await self._response.__anext__()
-        self._usage += _map_usage(chunk.data)
-
-        try:
-            choice = chunk.data.choices[0]
-
-        except IndexError:
-            raise StopAsyncIteration()
-
-        if choice.finish_reason is not None:
-            raise StopAsyncIteration()
-
-        content = choice.delta.content
-        if self._result_tools:
-            if text := _map_content(content):
-                self._delta_content = (self._delta_content or '') + text
-
-    def get(self, *, final: bool = False) -> ModelResponse:
-        calls: list[ModelResponsePart] = []
-        if self._function_tools and self._result_tools or self._function_tools:
-            for tool_call in self._function_tools.values():
-                tool = _map_mistral_to_pydantic_tool_call(tool_call)
-                calls.append(tool)
-
-        elif self._delta_content and self._result_tools:
-            output_json: dict[str, Any] | None = pydantic_core.from_json(
-                self._delta_content, allow_partial='trailing-strings'
-            )
+    _result_tools: dict[str, ToolDefinition]
 
-            if output_json:
-                for result_tool in self._result_tools.values():
-                    # NOTE: Additional verification to prevent JSON validation to crash in `_result.py`
-                    # Ensures required parameters in the JSON schema are respected, especially for stream-based return types.
-                    # Example with BaseModel and required fields.
-                    if not self._validate_required_json_schema(output_json, result_tool.parameters_json_schema):
-                        continue
+    _delta_content: str = field(default='', init=False)
 
-                    tool = ToolCallPart.from_raw_args(result_tool.name, output_json)
-                    calls.append(tool)
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+        chunk: MistralCompletionEvent
+        async for chunk in self._response:
+            self._usage += _map_usage(chunk.data)
 
-        return ModelResponse(calls, timestamp=self._timestamp)
+            try:
+                choice = chunk.data.choices[0]
+            except IndexError:
+                continue
+
+            # Handle the text part of the response
+            content = choice.delta.content
+            text = _map_content(content)
+            if text:
+                # Attempt to produce a result tool call from the received text
+                if self._result_tools:
+                    self._delta_content += text
+                    maybe_tool_call_part = self._try_get_result_tool_from_text(self._delta_content, self._result_tools)
+                    if maybe_tool_call_part:
+                        yield self._parts_manager.handle_tool_call_part(
+                            vendor_part_id='result',
+                            tool_name=maybe_tool_call_part.tool_name,
+                            args=maybe_tool_call_part.args_as_dict(),
+                            tool_call_id=maybe_tool_call_part.tool_call_id,
+                        )
+                else:
+                    yield self._parts_manager.handle_text_delta(vendor_part_id='content', content=text)
 
-    def usage(self) -> Usage:
-        return self._usage
+            # Handle the explicit tool calls
+            for index, dtc in enumerate(choice.delta.tool_calls or []):
+                # It seems that mistral just sends full tool calls, so we just use them directly, rather than building
+                yield self._parts_manager.handle_tool_call_part(
+                    vendor_part_id=index, tool_name=dtc.function.name, args=dtc.function.arguments, tool_call_id=dtc.id
+                )
 
     def timestamp(self) -> datetime:
         return self._timestamp
 
+    @staticmethod
+    def _try_get_result_tool_from_text(text: str, result_tools: dict[str, ToolDefinition]) -> ToolCallPart | None:
+        output_json: dict[str, Any] | None = pydantic_core.from_json(text, allow_partial='trailing-strings')
+        if output_json:
+            for result_tool in result_tools.values():
+                # NOTE: Additional verification to prevent JSON validation to crash in `_result.py`
+                # Ensures required parameters in the JSON schema are respected, especially for stream-based return types.
+                # Example with BaseModel and required fields.
+                if not MistralStreamedResponse._validate_required_json_schema(
+                    output_json, result_tool.parameters_json_schema
+                ):
+                    continue
+
+                # The following part_id will be thrown away
+                return ToolCallPart.from_raw_args(tool_name=result_tool.name, args=output_json)
+
     @staticmethod
     def _validate_required_json_schema(json_dict: dict[str, Any], json_schema: dict[str, Any]) -> bool:
         """Validate that all required parameters in the JSON schema are present in the JSON dictionary."""
@@ -587,20 +525,20 @@ class MistralStreamStructuredResponse(StreamStructuredResponse):
                 if not isinstance(json_dict[param], list):
                     return False
                 for item in json_dict[param]:
-                    if not isinstance(item, VALIDE_JSON_TYPE_MAPPING[param_items_type]):
+                    if not isinstance(item, VALID_JSON_TYPE_MAPPING[param_items_type]):
                         return False
-            elif param_type and not isinstance(json_dict[param], VALIDE_JSON_TYPE_MAPPING[param_type]):
+            elif param_type and not isinstance(json_dict[param], VALID_JSON_TYPE_MAPPING[param_type]):
                 return False
 
             if isinstance(json_dict[param], dict) and 'properties' in param_schema:
                 nested_schema = param_schema
-                if not MistralStreamStructuredResponse._validate_required_json_schema(json_dict[param], nested_schema):
+                if not MistralStreamedResponse._validate_required_json_schema(json_dict[param], nested_schema):
                     return False
 
         return True
 
 
-VALIDE_JSON_TYPE_MAPPING: dict[str, Any] = {
+VALID_JSON_TYPE_MAPPING: dict[str, Any] = {
     'string': str,
     'integer': int,
     'number': float,

pydantic_ai/models/ollama.py

@@ -36,6 +36,7 @@ CommonOllamaModelNames = Literal[
     'mistral-nemo',
     'mixtral',
     'phi3',
+    'phi4',
     'qwq',
     'qwen',
     'qwen2',

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,
         )