Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion pydantic_ai_slim/pydantic_ai/models/openrouter.py
Original file line number Diff line number Diff line change
Expand Up @@ -399,7 +399,10 @@ class _OpenRouterChoice(chat_completion.Choice):
class _OpenRouterCostDetails:
"""OpenRouter specific cost details."""

upstream_inference_cost: int | None = None
upstream_inference_cost: float | None = None

# TODO rework fields, tests/models/cassettes/test_openrouter/test_openrouter_google_nested_schema.yaml
# shows an `upstream_inference_completions_cost` field as well


class _OpenRouterPromptTokenDetails(completion_usage.PromptTokensDetails):
Expand Down
62 changes: 61 additions & 1 deletion pydantic_ai_slim/pydantic_ai/providers/openrouter.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,14 @@
from __future__ import annotations as _annotations

import os
from dataclasses import replace
from typing import overload

import httpx
from openai import AsyncOpenAI

from pydantic_ai import ModelProfile
from pydantic_ai._json_schema import JsonSchema, JsonSchemaTransformer
from pydantic_ai.exceptions import UserError
from pydantic_ai.models import cached_async_http_client
from pydantic_ai.profiles.amazon import amazon_model_profile
Expand All @@ -31,6 +33,64 @@
) from _import_error


class _OpenRouterGoogleJsonSchemaTransformer(JsonSchemaTransformer):
"""Legacy Google JSON schema transformer for OpenRouter compatibility.

OpenRouter's compatibility layer doesn't fully support modern JSON Schema features
like $defs/$ref and anyOf for nullable types. This transformer restores v1.19.0
behavior by inlining definitions and simplifying nullable unions.

See: https://github.com/pydantic/pydantic-ai/issues/3617
"""

def __init__(self, schema: JsonSchema, *, strict: bool | None = None):
super().__init__(schema, strict=strict, prefer_inlined_defs=True, simplify_nullable_unions=True)

def transform(self, schema: JsonSchema) -> JsonSchema:
# Remove properties not supported by Gemini
schema.pop('$schema', None)
schema.pop('title', None)
schema.pop('discriminator', None)
schema.pop('examples', None)
schema.pop('exclusiveMaximum', None)
schema.pop('exclusiveMinimum', None)

if (const := schema.pop('const', None)) is not None:
schema['enum'] = [const]

# Convert enums to string type (legacy Gemini requirement)
if enum := schema.get('enum'):
schema['type'] = 'string'
schema['enum'] = [str(val) for val in enum]

# Convert oneOf to anyOf for discriminated unions
if 'oneOf' in schema and 'type' not in schema:
schema['anyOf'] = schema.pop('oneOf')

# Handle string format -> description
type_ = schema.get('type')
if type_ == 'string' and (fmt := schema.pop('format', None)):
description = schema.get('description')
if description:
schema['description'] = f'{description} (format: {fmt})'
else:
schema['description'] = f'Format: {fmt}'

return schema


def _openrouter_google_model_profile(model_name: str) -> ModelProfile | None:
"""Get the model profile for a Google model accessed via OpenRouter.

Uses the legacy transformer to maintain compatibility with OpenRouter's
translation layer, which doesn't fully support modern JSON Schema features.
"""
profile = google_model_profile(model_name)
if profile is None: # pragma: no cover
return None
return replace(profile, json_schema_transformer=_OpenRouterGoogleJsonSchemaTransformer)


class OpenRouterProvider(Provider[AsyncOpenAI]):
"""Provider for OpenRouter API."""

Expand All @@ -48,7 +108,7 @@ def client(self) -> AsyncOpenAI:

def model_profile(self, model_name: str) -> ModelProfile | None:
provider_to_profile = {
'google': google_model_profile,
'google': _openrouter_google_model_profile,
'openai': openai_model_profile,
'anthropic': anthropic_model_profile,
'mistralai': mistral_model_profile,
Expand Down
Loading