add stream hooks

Author: ajac-zero

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -7,9 +7,9 @@
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field, replace
 from datetime import datetime
+from itertools import chain
 from typing import Any, Literal, cast, overload
 
-from openai.types.chat.chat_completion_chunk import Choice
 from pydantic import ValidationError
 from pydantic_core import to_json
 from typing_extensions import assert_never, deprecated
@@ -1702,7 +1702,8 @@ class OpenAIStreamedResponse(StreamedResponse):
     _provider_name: str
     _provider_url: str
 
-    def _handle_thinking_delta(self, choice: Choice):
+    def _handle_thinking_delta(self, chunk: ChatCompletionChunk):
+        choice = chunk.choices[0]
         # The `reasoning_content` field is only present in DeepSeek models.
         # https://api-docs.deepseek.com/guides/reasoning_model
         if reasoning_content := getattr(choice.delta, 'reasoning_content', None):
@@ -1724,12 +1725,45 @@ def _handle_thinking_delta(self, choice: Choice):
                 provider_name=self.provider_name,
             )
 
-    def _handle_provider_details(self, choice: Choice) -> dict[str, str] | None:
+    def _handle_provider_details(self, chunk: ChatCompletionChunk) -> dict[str, str] | None:
+        choice = chunk.choices[0]
         if raw_finish_reason := choice.finish_reason:
             return {'finish_reason': raw_finish_reason}
 
-    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+    def _handle_text_delta(self, chunk: ChatCompletionChunk):
+        # Handle the text part of the response
+        content = chunk.choices[0].delta.content
+        if content:
+            maybe_event = self._parts_manager.handle_text_delta(
+                vendor_part_id='content',
+                content=content,
+                thinking_tags=self._model_profile.thinking_tags,
+                ignore_leading_whitespace=self._model_profile.ignore_streamed_leading_whitespace,
+            )
+            if maybe_event is not None:  # pragma: no branch
+                if isinstance(maybe_event, PartStartEvent) and isinstance(maybe_event.part, ThinkingPart):
+                    maybe_event.part.id = 'content'
+                    maybe_event.part.provider_name = self.provider_name
+                yield maybe_event
+
+    def _handle_tool_delta(self, chunk: ChatCompletionChunk):
+        choice = chunk.choices[0]
+        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
+
+    async def _validate_response(self):
         async for chunk in self._response:
+            yield chunk
+
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+        async for chunk in self._validate_response():
             self._usage += self._map_usage(chunk)
 
             if chunk.id:  # pragma: no branch
@@ -1750,36 +1784,15 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
             if raw_finish_reason := choice.finish_reason:
                 self.finish_reason = self._map_finish_reason(raw_finish_reason)
 
-            if provider_details := self._handle_provider_details(choice):
+            if provider_details := self._handle_provider_details(chunk):
                 self.provider_details = provider_details
 
-            for thinking_part in self._handle_thinking_delta(choice):
-                yield thinking_part
-
-            # Handle the text part of the response
-            content = choice.delta.content
-            if content:
-                maybe_event = self._parts_manager.handle_text_delta(
-                    vendor_part_id='content',
-                    content=content,
-                    thinking_tags=self._model_profile.thinking_tags,
-                    ignore_leading_whitespace=self._model_profile.ignore_streamed_leading_whitespace,
-                )
-                if maybe_event is not None:  # pragma: no branch
-                    if isinstance(maybe_event, PartStartEvent) and isinstance(maybe_event.part, ThinkingPart):
-                        maybe_event.part.id = 'content'
-                        maybe_event.part.provider_name = self.provider_name
-                    yield maybe_event
-
-            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
+            for event in chain(
+                self._handle_thinking_delta(chunk),
+                self._handle_text_delta(chunk),
+                self._handle_tool_delta(chunk),
+            ):
+                yield event
 
     def _map_usage(self, response: ChatCompletionChunk):
         return _map_usage(response, self._provider_name, self._provider_url, self._model_name)

pydantic_ai_slim/pydantic_ai/models/openrouter.py

@@ -3,7 +3,7 @@
 
 from openai import AsyncStream
 from openai.types import chat
-from openai.types.chat.chat_completion import Choice
+from openai.types.chat import chat_completion, chat_completion_chunk
 from pydantic import AliasChoices, BaseModel, Field, TypeAdapter
 from typing_extensions import TypedDict, assert_never
 
@@ -346,7 +346,7 @@ class OpenRouterCompletionMessage(chat.ChatCompletionMessage):
     """The reasoning details associated with the message, if any."""
 
 
-class OpenRouterChoice(Choice):
+class OpenRouterChoice(chat_completion.Choice):
     """Wraps OpenAI chat completion choice with OpenRouter specific attributes."""
 
     native_finish_reason: str
@@ -375,14 +375,40 @@ class OpenRouterChatCompletion(chat.ChatCompletion):
     """OpenRouter specific error attribute."""
 
 
+class OpenRouterChoiceDelta(chat_completion_chunk.ChoiceDelta):
+    """Wrapped chat completion message with OpenRouter specific attributes."""
+
+    reasoning: str | None = None
+    """The reasoning text associated with the message, if any."""
+
+    reasoning_details: list[OpenRouterReasoningDetail] | None = None
+    """The reasoning details associated with the message, if any."""
+
+
+class OpenRouterChunkChoice(chat_completion_chunk.Choice):
+    """Wraps OpenAI chat completion chunk choice with OpenRouter specific attributes."""
+
+    native_finish_reason: str | None
+    """The provided finish reason by the downstream provider from OpenRouter."""
+
+    finish_reason: Literal['stop', 'length', 'tool_calls', 'content_filter', 'error'] | None  # type: ignore[reportIncompatibleVariableOverride]
+    """OpenRouter specific finish reasons for streaming chunks.
+
+    Notably, removes 'function_call' and adds 'error' finish reasons.
+    """
+
+    delta: OpenRouterChoiceDelta  # type: ignore[reportIncompatibleVariableOverride]
+    """A wrapped chat completion delta with OpenRouter specific attributes."""
+
+
 class OpenRouterChatCompletionChunk(chat.ChatCompletionChunk):
     """Wraps OpenAI chat completion with OpenRouter specific attributes."""
 
     provider: str
     """The downstream provider that was used by OpenRouter."""
 
-    choices: list[OpenRouterChoice]  # type: ignore[reportIncompatibleVariableOverride]
-    """A list of chat completion choices modified with OpenRouter specific attributes."""
+    choices: list[OpenRouterChunkChoice]  # type: ignore[reportIncompatibleVariableOverride]
+    """A list of chat completion chunk choices modified with OpenRouter specific attributes."""
 
     error: OpenRouterError | None = None
     """OpenRouter specific error attribute."""
@@ -428,6 +454,48 @@ class OpenRouterStreamedResponse(OpenAIStreamedResponse):
     def _map_usage(self, response: chat.ChatCompletionChunk):
         return _map_usage(response, self._provider_name, self._provider_url, self._model_name)
 
+    @override
+    def _map_finish_reason(  # type: ignore[reportIncompatibleMethodOverride]
+        self, key: Literal['stop', 'length', 'tool_calls', 'content_filter', 'error']
+    ) -> FinishReason | None:
+        return _CHAT_FINISH_REASON_MAP.get(key)
+
+    @override
+    def _handle_thinking_delta(self, chunk: OpenRouterChatCompletionChunk):  # type: ignore[reportIncompatibleMethodOverride]
+        delta = chunk.choices[0].delta
+        if reasoning_details := delta.reasoning_details:
+            for detail in reasoning_details:
+                thinking_part = OpenRouterThinkingPart.from_reasoning_detail(detail)
+                yield self._parts_manager.handle_thinking_delta(
+                    vendor_part_id='reasoning_detail',
+                    id=thinking_part.id,
+                    content=thinking_part.content,
+                    provider_name=self._provider_name,
+                )
+
+    @override
+    def _handle_provider_details(self, chunk: chat.ChatCompletionChunk) -> dict[str, str] | None:
+        native_chunk = OpenRouterChatCompletionChunk.model_validate(chunk.model_dump())
+
+        if provider_details := super()._handle_provider_details(chunk):
+            if provider := native_chunk.provider:
+                provider_details['downstream_provider'] = provider
+
+            if native_finish_reason := native_chunk.choices[0].native_finish_reason:
+                provider_details['native_finish_reason'] = native_finish_reason
+
+            return provider_details
+
+    @override
+    async def _validate_response(self):
+        async for chunk in self._response:
+            chunk = OpenRouterChatCompletionChunk.model_validate(chunk.model_dump())
+
+            if error := chunk.error:
+                raise ModelHTTPError(status_code=error.code, model_name=chunk.model, body=error.message)
+
+            yield chunk
+
 
 def _openrouter_settings_to_openai_settings(model_settings: OpenRouterModelSettings) -> OpenAIChatModelSettings:
     """Transforms a 'OpenRouterModelSettings' object into an 'OpenAIChatModelSettings' object.
@@ -475,6 +543,7 @@ def __init__(
         """
         super().__init__(model_name, provider=provider or OpenRouterProvider(), profile=profile, settings=settings)
 
+    @override
     def prepare_request(
         self,
         model_settings: ModelSettings | None,
@@ -485,13 +554,13 @@ def prepare_request(
         return new_settings, customized_parameters
 
     @override
-    def _map_finish_reason(
+    def _map_finish_reason(  # type: ignore[reportIncompatibleMethodOverride]
         self, key: Literal['stop', 'length', 'tool_calls', 'content_filter', 'error']
-    ) -> FinishReason | None:  # type: ignore[reportIncompatibleMethodOverride]
+    ) -> FinishReason | None:
         return _CHAT_FINISH_REASON_MAP.get(key)
 
     @override
-    def _process_reasoning(self, response: OpenRouterChatCompletion) -> list[ThinkingPart]:
+    def _process_reasoning(self, response: OpenRouterChatCompletion) -> list[ThinkingPart]:  # type: ignore[reportIncompatibleMethodOverride]
         message = response.choices[0].message
         items: list[ThinkingPart] = []
 
@@ -502,10 +571,7 @@ def _process_reasoning(self, response: OpenRouterChatCompletion) -> list[Thinkin
         return items
 
     @override
-    def _process_provider_details(self, response: OpenRouterChatCompletion) -> dict[str, Any]:
-        if error := response.error:
-            raise ModelHTTPError(status_code=error.code, model_name=response.model, body=error.message)
-
+    def _process_provider_details(self, response: OpenRouterChatCompletion) -> dict[str, Any]:  # type: ignore[reportIncompatibleMethodOverride]
         provider_details = super()._process_provider_details(response)
 
         provider_details['downstream_provider'] = response.provider
@@ -515,8 +581,14 @@ def _process_provider_details(self, response: OpenRouterChatCompletion) -> dict[
 
     @override
     def _validate_completion(self, response: chat.ChatCompletion) -> chat.ChatCompletion:
-        return OpenRouterChatCompletion.model_validate(response.model_dump())
+        response = OpenRouterChatCompletion.model_validate(response.model_dump())
 
+        if error := response.error:
+            raise ModelHTTPError(status_code=error.code, model_name=response.model, body=error.message)
+
+        return response
+
+    @override
     async def _process_streamed_response(
         self, response: AsyncStream[chat.ChatCompletionChunk], model_request_parameters: ModelRequestParameters
     ) -> OpenRouterStreamedResponse:
@@ -538,6 +610,7 @@ async def _process_streamed_response(
             _provider_url=self._provider.base_url,
         )
 
+    @override
     def _map_model_response(self, message: ModelResponse) -> chat.ChatCompletionMessageParam:
         texts: list[str] = []
         tool_calls: list[chat.ChatCompletionMessageFunctionToolCallParam] = []