Skip to content

Commit 126710a

Browse files
KludexKludex
authored
Add OpenRouter provider (#1778)
1 parent 320dd4f commit 126710a

File tree

7 files changed

+309
-5
lines changed

7 files changed

+309
-5
lines changed

‎docs/models/openai.md‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -294,14 +294,11 @@ Once you have the API key, you can use it with the `OpenAIProvider`:
294294
```python
295295
from pydantic_ai import Agent
296296
from pydantic_ai.models.openai import OpenAIModel
297-
from pydantic_ai.providers.openai import OpenAIProvider
297+
from pydantic_ai.providers.openrouter import OpenRouterProvider
298298

299299
model = OpenAIModel(
300300
'anthropic/claude-3.5-sonnet',
301-
provider=OpenAIProvider(
302-
base_url='https://openrouter.ai/api/v1',
303-
api_key='your-openrouter-api-key',
304-
),
301+
provider=OpenRouterProvider(api_key='your-openrouter-api-key'),
305302
)
306303
agent = Agent(model)
307304
...

‎pydantic_ai_slim/pydantic_ai/providers/__init__.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,10 @@ def infer_provider(provider: str) -> Provider[Any]:
5252
from .deepseek import DeepSeekProvider
5353

5454
return DeepSeekProvider()
55+
elif provider == 'openrouter':
56+
from .openrouter import OpenRouterProvider
57+
58+
return OpenRouterProvider()
5559
elif provider == 'azure':
5660
from .azure import AzureProvider
5761

‎pydantic_ai_slim/pydantic_ai/providers/openrouter.py‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
from __future__ import annotations as _annotations
2+
3+
import os
4+
from typing import overload
5+
6+
from httpx import AsyncClient as AsyncHTTPClient
7+
from openai import AsyncOpenAI
8+
9+
from pydantic_ai.exceptions import UserError
10+
from pydantic_ai.models import cached_async_http_client
11+
from pydantic_ai.providers import Provider
12+
13+
try:
14+
from openai import AsyncOpenAI
15+
except ImportError as _import_error: # pragma: no cover
16+
raise ImportError(
17+
'Please install the `openai` package to use the OpenRouter provider, '
18+
'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
19+
) from _import_error
20+
21+
22+
class OpenRouterProvider(Provider[AsyncOpenAI]):
23+
"""Provider for OpenRouter API."""
24+
25+
@property
26+
def name(self) -> str:
27+
return 'openrouter'
28+
29+
@property
30+
def base_url(self) -> str:
31+
return 'https://openrouter.ai/api/v1'
32+
33+
@property
34+
def client(self) -> AsyncOpenAI:
35+
return self._client
36+
37+
@overload
38+
def __init__(self) -> None: ...
39+
40+
@overload
41+
def __init__(self, *, api_key: str) -> None: ...
42+
43+
@overload
44+
def __init__(self, *, api_key: str, http_client: AsyncHTTPClient) -> None: ...
45+
46+
@overload
47+
def __init__(self, *, openai_client: AsyncOpenAI | None = None) -> None: ...
48+
49+
def __init__(
50+
self,
51+
*,
52+
api_key: str | None = None,
53+
openai_client: AsyncOpenAI | None = None,
54+
http_client: AsyncHTTPClient | None = None,
55+
) -> None:
56+
api_key = api_key or os.getenv('OPENROUTER_API_KEY')
57+
if not api_key and openai_client is None:
58+
raise UserError(
59+
'Set the `OPENROUTER_API_KEY` environment variable or pass it via `OpenRouterProvider(api_key=...)`'
60+
'to use the OpenRouter provider.'
61+
)
62+
63+
if openai_client is not None:
64+
self._client = openai_client
65+
elif http_client is not None:
66+
self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=http_client)
67+
else:
68+
http_client = cached_async_http_client(provider='openrouter')
69+
self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=http_client)

‎tests/conftest.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -276,6 +276,11 @@ def mistral_api_key() -> str:
276276
return os.getenv('MISTRAL_API_KEY', 'mock-api-key')
277277

278278

279+
@pytest.fixture(scope='session')
280+
def openrouter_api_key() -> str:
281+
return os.getenv('OPENROUTER_API_KEY', 'mock-api-key')
282+
283+
279284
@pytest.fixture(scope='session')
280285
def bedrock_provider():
281286
try:

‎tests/providers/cassettes/test_openrouter/test_openrouter_with_google_model.yaml‎

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
1+
interactions:
2+
- request:
3+
headers:
4+
accept:
5+
- application/json
6+
accept-encoding:
7+
- gzip, deflate
8+
connection:
9+
- keep-alive
10+
content-length:
11+
- '164'
12+
content-type:
13+
- application/json
14+
host:
15+
- openrouter.ai
16+
method: POST
17+
parsed_body:
18+
messages:
19+
- content: Be helpful.
20+
role: system
21+
- content: Tell me a joke.
22+
role: user
23+
model: google/gemini-2.0-flash-exp:free
24+
n: 1
25+
stream: false
26+
uri: https://openrouter.ai/api/v1/chat/completions
27+
response:
28+
headers:
29+
access-control-allow-origin:
30+
- '*'
31+
connection:
32+
- keep-alive
33+
content-length:
34+
- '242'
35+
content-type:
36+
- application/json
37+
vary:
38+
- Accept-Encoding
39+
parsed_body:
40+
error:
41+
code: 429
42+
message: Provider returned error
43+
metadata:
44+
provider_name: Google
45+
raw: google/gemini-2.0-flash-exp:free is temporarily rate-limited upstream; please retry shortly.
46+
user_id: user_2uRh0l3Yi3hdjBArTOSmLXWJBc4
47+
status:
48+
code: 429
49+
message: Too Many Requests
50+
- request:
51+
headers:
52+
accept:
53+
- application/json
54+
accept-encoding:
55+
- gzip, deflate
56+
connection:
57+
- keep-alive
58+
content-length:
59+
- '164'
60+
content-type:
61+
- application/json
62+
host:
63+
- openrouter.ai
64+
method: POST
65+
parsed_body:
66+
messages:
67+
- content: Be helpful.
68+
role: system
69+
- content: Tell me a joke.
70+
role: user
71+
model: google/gemini-2.0-flash-exp:free
72+
n: 1
73+
stream: false
74+
uri: https://openrouter.ai/api/v1/chat/completions
75+
response:
76+
headers:
77+
access-control-allow-origin:
78+
- '*'
79+
connection:
80+
- keep-alive
81+
content-length:
82+
- '252'
83+
content-type:
84+
- application/json
85+
vary:
86+
- Accept-Encoding
87+
parsed_body:
88+
error:
89+
code: 429
90+
message: Provider returned error
91+
metadata:
92+
provider_name: Google AI Studio
93+
raw: google/gemini-2.0-flash-exp:free is temporarily rate-limited upstream; please retry shortly.
94+
user_id: user_2uRh0l3Yi3hdjBArTOSmLXWJBc4
95+
status:
96+
code: 429
97+
message: Too Many Requests
98+
- request:
99+
headers:
100+
accept:
101+
- application/json
102+
accept-encoding:
103+
- gzip, deflate
104+
connection:
105+
- keep-alive
106+
content-length:
107+
- '164'
108+
content-type:
109+
- application/json
110+
host:
111+
- openrouter.ai
112+
method: POST
113+
parsed_body:
114+
messages:
115+
- content: Be helpful.
116+
role: system
117+
- content: Tell me a joke.
118+
role: user
119+
model: google/gemini-2.0-flash-exp:free
120+
n: 1
121+
stream: false
122+
uri: https://openrouter.ai/api/v1/chat/completions
123+
response:
124+
headers:
125+
access-control-allow-origin:
126+
- '*'
127+
connection:
128+
- keep-alive
129+
content-length:
130+
- '476'
131+
content-type:
132+
- application/json
133+
transfer-encoding:
134+
- chunked
135+
vary:
136+
- Accept-Encoding
137+
parsed_body:
138+
choices:
139+
- finish_reason: stop
140+
index: 0
141+
logprobs: null
142+
message:
143+
content: "Why don't scientists trust atoms? \n\nBecause they make up everything!\n"
144+
reasoning: null
145+
refusal: null
146+
role: assistant
147+
native_finish_reason: STOP
148+
created: 1747736401
149+
id: gen-1747736401-niWjLHssb1xvyg7Ow2Nf
150+
model: google/gemini-2.0-flash-exp:free
151+
object: chat.completion
152+
provider: Google
153+
usage:
154+
completion_tokens: 17
155+
prompt_tokens: 8
156+
total_tokens: 25
157+
status:
158+
code: 200
159+
message: OK
160+
version: 1

‎tests/providers/test_openrouter.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
import re
2+
3+
import httpx
4+
import pytest
5+
from inline_snapshot import snapshot
6+
7+
from pydantic_ai.agent import Agent
8+
from pydantic_ai.exceptions import UserError
9+
10+
from ..conftest import TestEnv, try_import
11+
12+
with try_import() as imports_successful:
13+
import openai
14+
15+
from pydantic_ai.models.openai import OpenAIModel
16+
from pydantic_ai.providers.openrouter import OpenRouterProvider
17+
18+
19+
pytestmark = [
20+
pytest.mark.skipif(not imports_successful(), reason='openai not installed'),
21+
pytest.mark.vcr,
22+
pytest.mark.anyio,
23+
]
24+
25+
26+
def test_openrouter_provider():
27+
provider = OpenRouterProvider(api_key='api-key')
28+
assert provider.name == 'openrouter'
29+
assert provider.base_url == 'https://openrouter.ai/api/v1'
30+
assert isinstance(provider.client, openai.AsyncOpenAI)
31+
assert provider.client.api_key == 'api-key'
32+
33+
34+
def test_openrouter_provider_need_api_key(env: TestEnv) -> None:
35+
env.remove('OPENROUTER_API_KEY')
36+
with pytest.raises(
37+
UserError,
38+
match=re.escape(
39+
'Set the `OPENROUTER_API_KEY` environment variable or pass it via `OpenRouterProvider(api_key=...)`'
40+
'to use the OpenRouter provider.'
41+
),
42+
):
43+
OpenRouterProvider()
44+
45+
46+
def test_openrouter_provider_pass_http_client() -> None:
47+
http_client = httpx.AsyncClient()
48+
provider = OpenRouterProvider(http_client=http_client, api_key='api-key')
49+
assert provider.client._client == http_client # type: ignore[reportPrivateUsage]
50+
51+
52+
def test_openrouter_pass_openai_client() -> None:
53+
openai_client = openai.AsyncOpenAI(api_key='api-key')
54+
provider = OpenRouterProvider(openai_client=openai_client)
55+
assert provider.client == openai_client
56+
57+
58+
async def test_openrouter_with_google_model(allow_model_requests: None, openrouter_api_key: str) -> None:
59+
provider = OpenRouterProvider(api_key=openrouter_api_key)
60+
model = OpenAIModel('google/gemini-2.0-flash-exp:free', provider=provider)
61+
agent = Agent(model, instructions='Be helpful.')
62+
response = await agent.run('Tell me a joke.')
63+
assert response.output == snapshot("""\
64+
Why don't scientists trust atoms? \n\
65+
66+
Because they make up everything!
67+
""")

‎tests/providers/test_provider_names.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,13 @@
2323
from pydantic_ai.providers.groq import GroqProvider
2424
from pydantic_ai.providers.mistral import MistralProvider
2525
from pydantic_ai.providers.openai import OpenAIProvider
26+
from pydantic_ai.providers.openrouter import OpenRouterProvider
2627

2728
test_infer_provider_params = [
2829
('anthropic', AnthropicProvider, 'ANTHROPIC_API_KEY'),
2930
('cohere', CohereProvider, 'CO_API_KEY'),
3031
('deepseek', DeepSeekProvider, 'DEEPSEEK_API_KEY'),
32+
('openrouter', OpenRouterProvider, 'OPENROUTER_API_KEY'),
3133
('openai', OpenAIProvider, 'OPENAI_API_KEY'),
3234
('azure', AzureProvider, 'AZURE_OPENAI'),
3335
('google-vertex', GoogleVertexProvider, None),

0 commit comments

Comments
 (0)