Add OpenRouterModel

Author: DanKing1903

pydantic_ai_slim/pydantic_ai/models/openrouter.py

@@ -0,0 +1,34 @@
+from typing import Any, cast
+
+from openai.types import chat
+
+from ..messages import ModelResponse
+from .openai import OpenAIModel
+
+
+class OpenRouterChatCompletion(chat.ChatCompletion):
+    """Extends ChatCompletion with OpenRouter-specific attributes.
+
+    This class extends the base ChatCompletion model to include additional
+    fields returned specifically by the OpenRouter API.
+
+    Attributes:
+        provider: The name of the upstream LLM provider (e.g., "Anthropic",
+            "OpenAI", etc.) that processed the request through OpenRouter.
+    """
+
+    provider: str
+
+
+class OpenRouterModel(OpenAIModel):
+    """Extends OpenAIModel to capture extra metadata for Openrouter."""
+
+    def _process_response(self, response: chat.ChatCompletion) -> ModelResponse:
+        response = cast(OpenRouterChatCompletion, response)
+        model_response = super()._process_response(response=response)
+        openrouter_provider: str | None = response.provider if hasattr(response, 'provider') else None
+        if openrouter_provider:
+            vendor_details: dict[str, Any] = getattr(model_response, 'vendor_details') or {}
+            vendor_details['provider'] = openrouter_provider
+            model_response.vendor_details = vendor_details
+        return model_response

tests/providers/test_openrouter.py

@@ -6,13 +6,15 @@
 
 from pydantic_ai.agent import Agent
 from pydantic_ai.exceptions import UserError
+from pydantic_ai.messages import ModelRequest, ModelResponse, TextPart, UserPromptPart
+from pydantic_ai.usage import Usage
 
-from ..conftest import TestEnv, try_import
+from ..conftest import IsDatetime, IsStr, TestEnv, try_import
 
 with try_import() as imports_successful:
     import openai
 
-    from pydantic_ai.models.openai import OpenAIModel
+    from pydantic_ai.models.openrouter import OpenRouterModel
     from pydantic_ai.providers.openrouter import OpenRouterProvider
 
 
@@ -57,11 +59,41 @@ def test_openrouter_pass_openai_client() -> None:
 
 async def test_openrouter_with_google_model(allow_model_requests: None, openrouter_api_key: str) -> None:
     provider = OpenRouterProvider(api_key=openrouter_api_key)
-    model = OpenAIModel('google/gemini-2.0-flash-exp:free', provider=provider)
-    agent = Agent(model, instructions='Be helpful.')
+    model = OpenRouterModel('google/gemini-2.0-flash-exp:free', provider=provider)
+    agent = Agent(model, instructions='Be helpful.', retries=1)
     response = await agent.run('Tell me a joke.')
     assert response.output == snapshot("""\
 Why don't scientists trust atoms? \n\
 
 Because they make up everything!
 """)
+
+    assert response.all_messages() == snapshot(
+        [
+            ModelRequest(
+                parts=[
+                    UserPromptPart(
+                        content='Tell me a joke.',
+                        timestamp=IsDatetime(iso_string=True),
+                    )
+                ],
+                instructions='Be helpful.',
+            ),
+            ModelResponse(
+                parts=[
+                    TextPart(
+                        content="""\
+Why don't scientists trust atoms? \n\
+
+Because they make up everything!
+"""
+                    )
+                ],
+                usage=Usage(requests=1, request_tokens=8, response_tokens=17, total_tokens=25, details={}),
+                model_name='google/gemini-2.0-flash-exp:free',
+                timestamp=IsDatetime(iso_string=True),
+                vendor_details={'provider': 'Google'},
+                vendor_id=IsStr(),
+            ),
+        ]
+    )