fix coverage/linting

Author: ajac-zero

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -1,6 +1,7 @@
 from __future__ import annotations as _annotations
 
 import base64
+import itertools
 import json
 import warnings
 from collections.abc import AsyncIterable, AsyncIterator, Sequence
@@ -62,6 +63,7 @@
         ChatCompletionContentPartParam,
         ChatCompletionContentPartTextParam,
     )
+    from openai.types.chat.chat_completion_chunk import Choice
     from openai.types.chat.chat_completion_content_part_image_param import ImageURL
     from openai.types.chat.chat_completion_content_part_input_audio_param import InputAudio
     from openai.types.chat.chat_completion_content_part_param import File, FileFile
@@ -530,9 +532,17 @@ async def _completions_create(
             raise  # pragma: lax no cover
 
     def _validate_completion(self, response: chat.ChatCompletion) -> chat.ChatCompletion:
+        """Hook that validates chat completions before processing.
+
+        This method may be overridden by subclasses of `OpenAIChatModel` to apply custom completion validations.
+        """
         return chat.ChatCompletion.model_validate(response.model_dump())
 
     def _process_reasoning(self, response: chat.ChatCompletion) -> list[ThinkingPart]:
+        """Hook that maps reasoning tokens to thinking parts.
+
+        This method may be overridden by subclasses of `OpenAIChatModel` to apply custom mappings.
+        """
         message = response.choices[0].message
         items: list[ThinkingPart] = []
 
@@ -550,6 +560,10 @@ def _process_reasoning(self, response: chat.ChatCompletion) -> list[ThinkingPart
         return items
 
     def _process_provider_details(self, response: chat.ChatCompletion) -> dict[str, Any]:
+        """Hook that response content to provider details.
+
+        This method may be overridden by subclasses of `OpenAIChatModel` to apply custom mappings.
+        """
         choice = response.choices[0]
         provider_details: dict[str, Any] = {}
 
@@ -692,7 +706,7 @@ def _get_web_search_options(self, model_request_parameters: ModelRequestParamete
     def _map_model_response(self, message: ModelResponse) -> chat.ChatCompletionMessageParam:
         """Hook that determines how `ModelResponse` is mapped into `ChatCompletionMessageParam` objects before sending.
 
-        Subclasses of `OpenAIChatModel` should override this method to provide their own mapping logic.
+        Subclasses of `OpenAIChatModel` may override this method to provide their own mapping logic.
         """
         texts: list[str] = []
         tool_calls: list[ChatCompletionMessageFunctionToolCallParam] = []
@@ -1740,19 +1754,28 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
     async def _validate_response(self):
         """Hook that validates incoming chunks.
 
-        This method should be overridden by subclasses of `OpenAIStreamedResponse` to apply custom chunk validations.
+        This method may be overridden by subclasses of `OpenAIStreamedResponse` to apply custom chunk validations.
 
         By default, this is a no-op since `ChatCompletionChunk` is already validated.
         """
         async for chunk in self._response:
             yield chunk
 
     def _map_part_delta(self, chunk: ChatCompletionChunk):
-        """Hook that maps delta content to events.
+        """Hook that determines the sequence of mappings that will be called to produce events.
 
-        This method should be overridden by subclasses of `OpenAIStreamResponse` to customize the mapping.
+        This method may be overridden by subclasses of `OpenAIStreamResponse` to customize the mapping.
         """
         choice = chunk.choices[0]
+        return itertools.chain(
+            self._map_thinking_delta(choice), self._map_text_delta(choice), self._map_tool_call_delta(choice)
+        )
+
+    def _map_thinking_delta(self, choice: Choice):
+        """Hook that maps thinking delta content to events.
+
+        This method may be overridden by subclasses of `OpenAIStreamResponse` to customize the mapping.
+        """
         # 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):
@@ -1774,6 +1797,11 @@ def _map_part_delta(self, chunk: ChatCompletionChunk):
                 provider_name=self.provider_name,
             )
 
+    def _map_text_delta(self, choice: Choice):
+        """Hook that maps text delta content to events.
+
+        This method may be overridden by subclasses of `OpenAIStreamResponse` to customize the mapping.
+        """
         # Handle the text part of the response
         content = choice.delta.content
         if content:
@@ -1789,6 +1817,11 @@ def _map_part_delta(self, chunk: ChatCompletionChunk):
                     maybe_event.part.provider_name = self.provider_name
                 yield maybe_event
 
+    def _map_tool_call_delta(self, choice: Choice):
+        """Hook that maps tool call delta content to events.
+
+        This method may be overridden by subclasses of `OpenAIStreamResponse` to customize the mapping.
+        """
         for dtc in choice.delta.tool_calls or []:
             maybe_event = self._parts_manager.handle_tool_call_delta(
                 vendor_part_id=dtc.index,
@@ -1802,7 +1835,7 @@ def _map_part_delta(self, chunk: ChatCompletionChunk):
     def _map_provider_details(self, chunk: ChatCompletionChunk) -> dict[str, str] | None:
         """Hook that generates the provider details from chunk content.
 
-        This method should be overridden by subclasses of `OpenAIStreamResponse` to customize the provider details.
+        This method may be overridden by subclasses of `OpenAIStreamResponse` to customize the provider details.
         """
         choice = chunk.choices[0]
         if raw_finish_reason := choice.finish_reason:

pydantic_ai_slim/pydantic_ai/models/openrouter.py

@@ -1,11 +1,8 @@
 from dataclasses import asdict, dataclass
-from typing import Any, Literal, cast, override
+from typing import Any, Literal, cast
 
-from openai import APIError
-from openai.types import chat
-from openai.types.chat import chat_completion, chat_completion_chunk
 from pydantic import AliasChoices, BaseModel, Field, TypeAdapter
-from typing_extensions import TypedDict, assert_never
+from typing_extensions import TypedDict, assert_never, override
 
 from .. import _utils
 from ..exceptions import ModelHTTPError
@@ -15,7 +12,6 @@
     FilePart,
     FinishReason,
     ModelResponse,
-    PartStartEvent,
     TextPart,
     ThinkingPart,
     ToolCallPart,
@@ -25,7 +21,18 @@
 from ..settings import ModelSettings
 from ..usage import RequestUsage
 from . import ModelRequestParameters
-from .openai import OpenAIChatModel, OpenAIChatModelSettings, OpenAIStreamedResponse
+
+try:
+    from openai import APIError
+    from openai.types import chat
+    from openai.types.chat import chat_completion, chat_completion_chunk
+
+    from .openai import OpenAIChatModel, OpenAIChatModelSettings, OpenAIStreamedResponse
+except ImportError as _import_error:
+    raise ImportError(
+        'Please install `openai` to use the OpenRouter model, '
+        'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
+    ) from _import_error
 
 _CHAT_FINISH_REASON_MAP: dict[Literal['stop', 'length', 'tool_calls', 'content_filter', 'error'], FinishReason] = {
     'stop': 'stop',
@@ -592,10 +599,9 @@ async def _validate_response(self):
             raise ModelHTTPError(status_code=error.code, model_name=self._model_name, body=error.message)
 
     @override
-    def _map_part_delta(self, chunk: chat.ChatCompletionChunk):
+    def _map_thinking_delta(self, choice: chat_completion_chunk.Choice):
         # We can cast with confidence because chunk was validated in `_validate_response`
-        chunk = cast(OpenRouterChatCompletionChunk, chunk)
-        choice = chunk.choices[0]
+        choice = cast(OpenRouterChunkChoice, choice)
 
         if reasoning_details := choice.delta.reasoning_details:
             for detail in reasoning_details:
@@ -607,31 +613,6 @@ def _map_part_delta(self, chunk: chat.ChatCompletionChunk):
                     provider_name=self._provider_name,
                 )
 
-        # 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
-
     @override
     def _map_provider_details(self, chunk: chat.ChatCompletionChunk) -> dict[str, str] | None:
         chunk = cast(OpenRouterChatCompletionChunk, chunk)