replace OpenRouterThinking with encoding in 'id'

Author: ajac-zero

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -538,27 +538,6 @@ def _validate_completion(self, response: chat.ChatCompletion) -> chat.ChatComple
         """
         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] = []
-
-        # The `reasoning_content` field is only present in DeepSeek models.
-        # https://api-docs.deepseek.com/guides/reasoning_model
-        if reasoning_content := getattr(message, 'reasoning_content', None):
-            items.append(ThinkingPart(id='reasoning_content', content=reasoning_content, provider_name=self.system))
-
-        # The `reasoning` field is only present in gpt-oss via Ollama and OpenRouter.
-        # - https://cookbook.openai.com/articles/gpt-oss/handle-raw-cot#chat-completions-api
-        # - https://openrouter.ai/docs/use-cases/reasoning-tokens#basic-usage-with-reasoning-tokens
-        if reasoning := getattr(message, 'reasoning', None):
-            items.append(ThinkingPart(id='reasoning', content=reasoning, provider_name=self.system))
-
-        return items
-
     def _process_provider_details(self, response: chat.ChatCompletion) -> dict[str, Any]:
         """Hook that response content to provider details.
 
@@ -616,7 +595,7 @@ def _process_response(self, response: chat.ChatCompletion | str) -> ModelRespons
         choice = response.choices[0]
         items: list[ModelResponsePart] = []
 
-        if thinking_parts := self._process_reasoning(response):
+        if thinking_parts := self._process_thinking(choice.message):
             items.extend(thinking_parts)
 
         if choice.message.content:
@@ -648,6 +627,26 @@ def _process_response(self, response: chat.ChatCompletion | str) -> ModelRespons
             finish_reason=self._map_finish_reason(choice.finish_reason),
         )
 
+    def _process_thinking(self, message: chat.ChatCompletionMessage) -> list[ThinkingPart] | None:
+        """Hook that maps reasoning tokens to thinking parts.
+
+        This method may be overridden by subclasses of `OpenAIChatModel` to apply custom mappings.
+        """
+        items: list[ThinkingPart] = []
+
+        # The `reasoning_content` field is only present in DeepSeek models.
+        # https://api-docs.deepseek.com/guides/reasoning_model
+        if reasoning_content := getattr(message, 'reasoning_content', None):
+            items.append(ThinkingPart(id='reasoning_content', content=reasoning_content, provider_name=self.system))
+
+        # The `reasoning` field is only present in gpt-oss via Ollama and OpenRouter.
+        # - https://cookbook.openai.com/articles/gpt-oss/handle-raw-cot#chat-completions-api
+        # - https://openrouter.ai/docs/use-cases/reasoning-tokens#basic-usage-with-reasoning-tokens
+        if reasoning := getattr(message, 'reasoning', None):
+            items.append(ThinkingPart(id='reasoning', content=reasoning, provider_name=self.system))
+
+        return items
+
     async def _process_streamed_response(
         self, response: AsyncStream[ChatCompletionChunk], model_request_parameters: ModelRequestParameters
     ) -> OpenAIStreamedResponse:
@@ -674,7 +673,11 @@ async def _process_streamed_response(
         )
 
     @property
-    def _streamed_response_cls(self):
+    def _streamed_response_cls(self) -> type[OpenAIStreamedResponse]:
+        """Returns the `StreamedResponse` type that will be used for streamed responses.
+
+        This method may be overridden by subclasses of `OpenAIChatModel` to provide their own `StreamedResponse` type.
+        """
         return OpenAIStreamedResponse
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[chat.ChatCompletionToolParam]:
@@ -743,6 +746,10 @@ def _map_model_response(self, message: ModelResponse) -> chat.ChatCompletionMess
     def _map_finish_reason(
         self, key: Literal['stop', 'length', 'tool_calls', 'content_filter', 'function_call']
     ) -> FinishReason | None:
+        """Hooks that maps a finish reason key to a [FinishReason](pydantic_ai.messages.FinishReason).
+
+        This method may be overridden by subclasses of `OpenAIChatModel` to accommodate custom keys.
+        """
         return _CHAT_FINISH_REASON_MAP.get(key)
 
     async def _map_messages(self, messages: list[ModelMessage]) -> list[chat.ChatCompletionMessageParam]:
@@ -1752,7 +1759,7 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
             if provider_details := self._map_provider_details(chunk):
                 self.provider_details = provider_details
 
-            for event in self._map_part_delta(chunk):
+            for event in self._map_part_delta(choice):
                 yield event
 
     async def _validate_response(self):
@@ -1765,12 +1772,11 @@ async def _validate_response(self):
         async for chunk in self._response:
             yield chunk
 
-    def _map_part_delta(self, chunk: ChatCompletionChunk):
+    def _map_part_delta(self, choice: Choice):
         """Hook that determines the sequence of mappings that will be called to produce events.
 
         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)
         )
@@ -1851,6 +1857,10 @@ def _map_usage(self, response: ChatCompletionChunk):
     def _map_finish_reason(
         self, key: Literal['stop', 'length', 'tool_calls', 'content_filter', 'function_call']
     ) -> FinishReason | None:
+        """Hooks that maps a finish reason key to a [FinishReason](pydantic_ai.messages.FinishReason).
+
+        This method may be overridden by subclasses of `OpenAIChatModel` to accommodate custom keys.
+        """
         return _CHAT_FINISH_REASON_MAP.get(key)
 
     @property

pydantic_ai_slim/pydantic_ai/models/openrouter.py

@@ -1,11 +1,11 @@
+import json
 from dataclasses import asdict, dataclass
 from typing import Any, Literal, cast
 
 from pydantic import AliasChoices, BaseModel, Field, TypeAdapter
 from typing_extensions import TypedDict, assert_never, override
 
-from .. import _utils
-from ..exceptions import ModelHTTPError
+from ..exceptions import ModelHTTPError, UnexpectedModelBehavior
 from ..messages import (
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
@@ -263,85 +263,72 @@ class OpenRouterError(BaseModel):
     message: str
 
 
-class BaseReasoningDetail(BaseModel):
+class _BaseReasoningDetail(BaseModel):
     """Common fields shared across all reasoning detail types."""
 
     id: str | None = None
     format: Literal['unknown', 'openai-responses-v1', 'anthropic-claude-v1', 'xai-responses-v1'] | None
     index: int | None
 
 
-class ReasoningSummary(BaseReasoningDetail):
+class _ReasoningSummary(_BaseReasoningDetail):
     """Represents a high-level summary of the reasoning process."""
 
     type: Literal['reasoning.summary']
     summary: str = Field(validation_alias=AliasChoices('summary', 'content'))
 
 
-class ReasoningEncrypted(BaseReasoningDetail):
+class _ReasoningEncrypted(_BaseReasoningDetail):
     """Represents encrypted reasoning data."""
 
     type: Literal['reasoning.encrypted']
     data: str = Field(validation_alias=AliasChoices('data', 'signature'))
 
 
-class ReasoningText(BaseReasoningDetail):
+class _ReasoningText(_BaseReasoningDetail):
     """Represents raw text reasoning."""
 
     type: Literal['reasoning.text']
     text: str = Field(validation_alias=AliasChoices('text', 'content'))
     signature: str | None = None
 
 
-OpenRouterReasoningDetail = ReasoningSummary | ReasoningEncrypted | ReasoningText
-_reasoning_detail_adapter: TypeAdapter[OpenRouterReasoningDetail] = TypeAdapter(OpenRouterReasoningDetail)
-
-
-@dataclass(repr=False)
-class OpenRouterThinkingPart(ThinkingPart):
-    """A special ThinkingPart that includes reasoning attributes specific to OpenRouter."""
-
-    type: Literal['reasoning.summary', 'reasoning.encrypted', 'reasoning.text']
-    index: int | None
-    format: Literal['unknown', 'openai-responses-v1', 'anthropic-claude-v1', 'xai-responses-v1'] | None
-
-    __repr__ = _utils.dataclasses_no_defaults_repr
-
-    @classmethod
-    def from_reasoning_detail(cls, reasoning: OpenRouterReasoningDetail):
-        provider_name = 'openrouter'
-        if isinstance(reasoning, ReasoningText):
-            return cls(
-                id=reasoning.id,
-                content=reasoning.text,
-                signature=reasoning.signature,
-                provider_name=provider_name,
-                format=reasoning.format,
-                type=reasoning.type,
-                index=reasoning.index,
-            )
-        elif isinstance(reasoning, ReasoningSummary):
-            return cls(
-                id=reasoning.id,
-                content=reasoning.summary,
-                provider_name=provider_name,
-                format=reasoning.format,
-                type=reasoning.type,
-                index=reasoning.index,
-            )
-        else:
-            return cls(
-                id=reasoning.id,
-                content='',
-                signature=reasoning.data,
-                provider_name=provider_name,
-                format=reasoning.format,
-                type=reasoning.type,
-                index=reasoning.index,
-            )
-
-    def into_reasoning_detail(self):
-        return _reasoning_detail_adapter.validate_python(asdict(self)).model_dump()
+_OpenRouterReasoningDetail = _ReasoningSummary | _ReasoningEncrypted | _ReasoningText
+_reasoning_detail_adapter: TypeAdapter[_OpenRouterReasoningDetail] = TypeAdapter(_OpenRouterReasoningDetail)
+
+
+def _from_reasoning_detail(reasoning: _OpenRouterReasoningDetail) -> ThinkingPart:
+    provider_name = 'openrouter'
+    reasoning_id = reasoning.model_dump_json(include={'id', 'format', 'index', 'type'})
+    if isinstance(reasoning, _ReasoningText):
+        return ThinkingPart(
+            id=reasoning_id,
+            content=reasoning.text,
+            signature=reasoning.signature,
+            provider_name=provider_name,
+        )
+    elif isinstance(reasoning, _ReasoningSummary):
+        return ThinkingPart(
+            id=reasoning_id,
+            content=reasoning.summary,
+            provider_name=provider_name,
+        )
+    else:
+        return ThinkingPart(
+            id=reasoning_id,
+            content='',
+            signature=reasoning.data,
+            provider_name=provider_name,
+        )
+
+
+def _into_reasoning_detail(thinking_part: ThinkingPart) -> _OpenRouterReasoningDetail:
+    if thinking_part.id is None:  # pragma: lax no cover
+        raise UnexpectedModelBehavior('OpenRouter thinking part has no ID')
+
+    data = asdict(thinking_part)
+    data.update(json.loads(thinking_part.id))
+    return _reasoning_detail_adapter.validate_python(data)
 
 
 class OpenRouterCompletionMessage(chat.ChatCompletionMessage):
@@ -350,7 +337,7 @@ class OpenRouterCompletionMessage(chat.ChatCompletionMessage):
     reasoning: str | None = None
     """The reasoning text associated with the message, if any."""
 
-    reasoning_details: list[OpenRouterReasoningDetail] | None = None
+    reasoning_details: list[_OpenRouterReasoningDetail] | None = None
     """The reasoning details associated with the message, if any."""
 
 
@@ -481,22 +468,17 @@ def _validate_completion(self, response: chat.ChatCompletion) -> OpenRouterChatC
         return response
 
     @override
-    def _process_reasoning(self, response: chat.ChatCompletion) -> list[ThinkingPart]:
-        # We can cast with confidence because response was validated in `_validate_completion`
-        response = cast(OpenRouterChatCompletion, response)
-
-        message = response.choices[0].message
-        items: list[ThinkingPart] = []
+    def _process_thinking(self, message: chat.ChatCompletionMessage) -> list[ThinkingPart] | None:
+        assert isinstance(message, OpenRouterCompletionMessage)
 
         if reasoning_details := message.reasoning_details:
-            for detail in reasoning_details:
-                items.append(OpenRouterThinkingPart.from_reasoning_detail(detail))
-
-        return items
+            return [_from_reasoning_detail(detail) for detail in reasoning_details]
+        else:
+            return super()._process_thinking(message)
 
     @override
     def _process_provider_details(self, response: chat.ChatCompletion) -> dict[str, Any]:
-        response = cast(OpenRouterChatCompletion, response)
+        assert isinstance(response, OpenRouterChatCompletion)
 
         provider_details = super()._process_provider_details(response)
 
@@ -514,8 +496,8 @@ def _map_model_response(self, message: ModelResponse) -> chat.ChatCompletionMess
             if isinstance(item, TextPart):
                 texts.append(item.content)
             elif isinstance(item, ThinkingPart):
-                if item.provider_name == self.system and isinstance(item, OpenRouterThinkingPart):
-                    reasoning_details.append(item.into_reasoning_detail())
+                if item.provider_name == self.system:
+                    reasoning_details.append(_into_reasoning_detail(item).model_dump())
                 else:  # pragma: no cover
                     pass
             elif isinstance(item, ToolCallPart):
@@ -555,7 +537,7 @@ class OpenRouterChoiceDelta(chat_completion_chunk.ChoiceDelta):
     reasoning: str | None = None
     """The reasoning text associated with the message, if any."""
 
-    reasoning_details: list[OpenRouterReasoningDetail] | None = None
+    reasoning_details: list[_OpenRouterReasoningDetail] | None = None
     """The reasoning details associated with the message, if any."""
 
 
@@ -605,7 +587,7 @@ def _map_thinking_delta(self, choice: chat_completion_chunk.Choice):
 
         if reasoning_details := choice.delta.reasoning_details:
             for detail in reasoning_details:
-                thinking_part = OpenRouterThinkingPart.from_reasoning_detail(detail)
+                thinking_part = _from_reasoning_detail(detail)
                 yield self._parts_manager.handle_thinking_delta(
                     vendor_part_id='reasoning_detail',
                     id=thinking_part.id,

tests/models/test_openrouter.py

@@ -110,15 +110,19 @@ async def test_openrouter_stream_with_reasoning(allow_model_requests: None, open
         assert isinstance(thinking_event_start, PartStartEvent)
         assert thinking_event_start.part == snapshot(
             ThinkingPart(
-                content='', id='rs_0aa4f2c435e6d1dc0169082486816c8193a029b5fc4ef1764f', provider_name='openrouter'
+                content='',
+                id='{"id":"rs_0aa4f2c435e6d1dc0169082486816c8193a029b5fc4ef1764f","format":"openai-responses-v1","index":0,"type":"reasoning.encrypted"}',
+                provider_name='openrouter',
             )
         )
 
         thinking_event_end = chunks[1]
         assert isinstance(thinking_event_end, PartEndEvent)
         assert thinking_event_end.part == snapshot(
             ThinkingPart(
-                content='', id='rs_0aa4f2c435e6d1dc0169082486816c8193a029b5fc4ef1764f', provider_name='openrouter'
+                content='',
+                id='{"id":"rs_0aa4f2c435e6d1dc0169082486816c8193a029b5fc4ef1764f","format":"openai-responses-v1","index":0,"type":"reasoning.encrypted"}',
+                provider_name='openrouter',
             )
         )
 
@@ -200,7 +204,7 @@ async def test_openrouter_with_reasoning(allow_model_requests: None, openrouter_
 
     thinking_part = response.parts[0]
     assert isinstance(thinking_part, ThinkingPart)
-    assert thinking_part.id == snapshot(None)
+    assert thinking_part.id == snapshot('{"id":null,"format":"unknown","index":0,"type":"reasoning.text"}')
     assert thinking_part.content is not None
     assert thinking_part.signature is None