diff --git a/pydantic_ai_slim/pydantic_ai/models/openrouter.py b/pydantic_ai_slim/pydantic_ai/models/openrouter.py index 07550b9a03..a219344fe4 100644 --- a/pydantic_ai_slim/pydantic_ai/models/openrouter.py +++ b/pydantic_ai_slim/pydantic_ai/models/openrouter.py @@ -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): diff --git a/pydantic_ai_slim/pydantic_ai/providers/openrouter.py b/pydantic_ai_slim/pydantic_ai/providers/openrouter.py index 70f962e047..63a73eb05d 100644 --- a/pydantic_ai_slim/pydantic_ai/providers/openrouter.py +++ b/pydantic_ai_slim/pydantic_ai/providers/openrouter.py @@ -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 @@ -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.""" @@ -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, diff --git a/tests/models/cassettes/test_openrouter/test_openrouter_google_nested_schema.yaml b/tests/models/cassettes/test_openrouter/test_openrouter_google_nested_schema.yaml new file mode 100644 index 0000000000..a30ec9d550 --- /dev/null +++ b/tests/models/cassettes/test_openrouter/test_openrouter_google_nested_schema.yaml @@ -0,0 +1,325 @@ +interactions: +- request: + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '1241' + content-type: + - application/json + host: + - openrouter.ai + method: POST + parsed_body: + messages: + - content: It's a house with a ground floor that has an entryway, a living room and a garage. + role: user + model: google/gemini-2.5-flash + stream: false + tool_choice: required + tools: + - function: + description: Insert a level with its spaces. + name: insert_level_with_spaces + parameters: + additionalProperties: false + properties: + level: + nullable: true + properties: + level_name: + type: string + level_type: + enum: + - ground + - basement + - floor + - attic + type: string + required: + - level_name + - level_type + type: object + spaces: + items: + properties: + space_name: + type: string + space_type: + enum: + - entryway + - living-room + - kitchen + - bedroom + - bathroom + - garage + type: string + required: + - space_name + - space_type + type: object + type: array + required: + - level + - spaces + type: object + strict: true + type: function + - function: + description: Result of inserting a level. + name: final_result + parameters: + properties: + level_name: + type: string + level_type: + enum: + - ground + - basement + - floor + - attic + type: string + space_count: + type: integer + required: + - level_name + - level_type + - space_count + type: object + strict: true + type: function + uri: https://openrouter.ai/api/v1/chat/completions + response: + headers: + access-control-allow-origin: + - '*' + connection: + - keep-alive + content-length: + - '1164' + content-type: + - application/json + permissions-policy: + - payment=(self "https://checkout.stripe.com" "https://connect-js.stripe.com" "https://js.stripe.com" "https://*.js.stripe.com" + "https://hooks.stripe.com") + referrer-policy: + - no-referrer, strict-origin-when-cross-origin + transfer-encoding: + - chunked + vary: + - Accept-Encoding + parsed_body: + choices: + - finish_reason: tool_calls + index: 0 + logprobs: null + message: + content: '' + reasoning: null + refusal: null + role: assistant + tool_calls: + - function: + arguments: '{"spaces":[{"space_type":"entryway","space_name":"entryway"},{"space_name":"living_room","space_type":"living-room"},{"space_name":"garage","space_type":"garage"}],"level":{"level_type":"ground","level_name":"ground_floor"}}' + name: insert_level_with_spaces + id: tool_insert_level_with_spaces_3ZiChYzj8xER8HixJe7W + index: 0 + type: function + native_finish_reason: STOP + created: 1764791728 + id: gen-1764791728-YEVpGoInRszZx8oZ508T + model: google/gemini-2.5-flash + object: chat.completion + provider: Google AI Studio + usage: + completion_tokens: 91 + completion_tokens_details: + image_tokens: 0 + reasoning_tokens: 0 + cost: 0 + cost_details: + upstream_inference_completions_cost: 0.0002275 + upstream_inference_cost: 0.0003253 + upstream_inference_prompt_cost: 9.78e-05 + is_byok: true + prompt_tokens: 326 + prompt_tokens_details: + audio_tokens: 0 + cached_tokens: 0 + video_tokens: 0 + total_tokens: 417 + status: + code: 200 + message: OK +- request: + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '1883' + content-type: + - application/json + host: + - openrouter.ai + method: POST + parsed_body: + messages: + - content: It's a house with a ground floor that has an entryway, a living room and a garage. + role: user + - content: null + role: assistant + tool_calls: + - function: + arguments: '{"spaces":[{"space_type":"entryway","space_name":"entryway"},{"space_name":"living_room","space_type":"living-room"},{"space_name":"garage","space_type":"garage"}],"level":{"level_type":"ground","level_name":"ground_floor"}}' + name: insert_level_with_spaces + id: tool_insert_level_with_spaces_3ZiChYzj8xER8HixJe7W + type: function + - content: 'Inserted level level_name=''ground_floor'' level_type= with 3 spaces' + role: tool + tool_call_id: tool_insert_level_with_spaces_3ZiChYzj8xER8HixJe7W + model: google/gemini-2.5-flash + stream: false + tool_choice: required + tools: + - function: + description: Insert a level with its spaces. + name: insert_level_with_spaces + parameters: + additionalProperties: false + properties: + level: + nullable: true + properties: + level_name: + type: string + level_type: + enum: + - ground + - basement + - floor + - attic + type: string + required: + - level_name + - level_type + type: object + spaces: + items: + properties: + space_name: + type: string + space_type: + enum: + - entryway + - living-room + - kitchen + - bedroom + - bathroom + - garage + type: string + required: + - space_name + - space_type + type: object + type: array + required: + - level + - spaces + type: object + strict: true + type: function + - function: + description: Result of inserting a level. + name: final_result + parameters: + properties: + level_name: + type: string + level_type: + enum: + - ground + - basement + - floor + - attic + type: string + space_count: + type: integer + required: + - level_name + - level_type + - space_count + type: object + strict: true + type: function + uri: https://openrouter.ai/api/v1/chat/completions + response: + headers: + access-control-allow-origin: + - '*' + connection: + - keep-alive + content-length: + - '945' + content-type: + - application/json + permissions-policy: + - payment=(self "https://checkout.stripe.com" "https://connect-js.stripe.com" "https://js.stripe.com" "https://*.js.stripe.com" + "https://hooks.stripe.com") + referrer-policy: + - no-referrer, strict-origin-when-cross-origin + transfer-encoding: + - chunked + vary: + - Accept-Encoding + parsed_body: + choices: + - finish_reason: tool_calls + index: 0 + logprobs: null + message: + content: '' + reasoning: null + refusal: null + role: assistant + tool_calls: + - function: + arguments: '{"level_type":"ground","level_name":"ground_floor","space_count":3}' + name: final_result + id: tool_final_result_HesCvwqQXZaVlFW3buU8 + index: 0 + type: function + native_finish_reason: STOP + created: 1764791730 + id: gen-1764791730-nym6GYDkejhRdrWtv8l3 + model: google/gemini-2.5-flash + object: chat.completion + provider: Google AI Studio + usage: + completion_tokens: 33 + completion_tokens_details: + image_tokens: 0 + reasoning_tokens: 0 + cost: 0 + cost_details: + upstream_inference_completions_cost: 8.25e-05 + upstream_inference_cost: 0.0002265 + upstream_inference_prompt_cost: 0.000144 + is_byok: true + prompt_tokens: 480 + prompt_tokens_details: + audio_tokens: 0 + cached_tokens: 0 + video_tokens: 0 + total_tokens: 513 + status: + code: 200 + message: OK +version: 1 diff --git a/tests/models/test_openrouter.py b/tests/models/test_openrouter.py index c52ec6e546..d439fee026 100644 --- a/tests/models/test_openrouter.py +++ b/tests/models/test_openrouter.py @@ -406,3 +406,67 @@ class FindEducationContentFilters(BaseModel): } ] ) + + +async def test_openrouter_google_nested_schema(allow_model_requests: None, openrouter_api_key: str) -> None: + """Test that nested schemas with $defs/$ref work correctly with OpenRouter + Gemini. + + This verifies the fix for https://github.com/pydantic/pydantic-ai/issues/3617 + where OpenRouter's translation layer didn't support modern JSON Schema features. + """ + from enum import Enum + + provider = OpenRouterProvider(api_key=openrouter_api_key) + + class LevelType(str, Enum): + ground = 'ground' + basement = 'basement' + floor = 'floor' + attic = 'attic' + + class SpaceType(str, Enum): + entryway = 'entryway' + living_room = 'living-room' + kitchen = 'kitchen' + bedroom = 'bedroom' + bathroom = 'bathroom' + garage = 'garage' + + class InsertLevelArg(BaseModel): + level_name: str + level_type: LevelType + + class SpaceArg(BaseModel): + space_name: str + space_type: SpaceType + + class InsertedLevel(BaseModel): + """Result of inserting a level.""" + + level_name: str + level_type: LevelType + space_count: int + + model = OpenRouterModel('google/gemini-2.5-flash', provider=provider) + agent: Agent[None, InsertedLevel] = Agent(model, output_type=InsertedLevel) + + @agent.tool_plain + def insert_level_with_spaces(level: InsertLevelArg | None, spaces: list[SpaceArg]) -> str: + """Insert a level with its spaces.""" + return f'Inserted level {level} with {len(spaces)} spaces' + + result = await agent.run("It's a house with a ground floor that has an entryway, a living room and a garage.") + + tool_call_message = result.all_messages()[1] + assert tool_call_message.parts == snapshot( + [ + ToolCallPart( + tool_name='insert_level_with_spaces', + args='{"spaces":[{"space_type":"entryway","space_name":"entryway"},{"space_name":"living_room","space_type":"living-room"},{"space_name":"garage","space_type":"garage"}],"level":{"level_type":"ground","level_name":"ground_floor"}}', + tool_call_id='tool_insert_level_with_spaces_3ZiChYzj8xER8HixJe7W', + ) + ] + ) + + assert result.output.level_type == LevelType.ground + assert result.output.space_count == 3 diff --git a/tests/providers/test_openrouter.py b/tests/providers/test_openrouter.py index ffd54bdab5..8ee94bce7f 100644 --- a/tests/providers/test_openrouter.py +++ b/tests/providers/test_openrouter.py @@ -12,7 +12,7 @@ from pydantic_ai.profiles.anthropic import AnthropicJsonSchemaTransformer, anthropic_model_profile from pydantic_ai.profiles.cohere import cohere_model_profile from pydantic_ai.profiles.deepseek import deepseek_model_profile -from pydantic_ai.profiles.google import GoogleJsonSchemaTransformer, google_model_profile +from pydantic_ai.profiles.google import google_model_profile from pydantic_ai.profiles.grok import grok_model_profile from pydantic_ai.profiles.meta import meta_model_profile from pydantic_ai.profiles.mistral import mistral_model_profile @@ -26,7 +26,10 @@ import openai from pydantic_ai.models.openrouter import OpenRouterModel - from pydantic_ai.providers.openrouter import OpenRouterProvider + from pydantic_ai.providers.openrouter import ( + OpenRouterProvider, + _OpenRouterGoogleJsonSchemaTransformer, # pyright: ignore[reportPrivateUsage] + ) pytestmark = [ @@ -109,12 +112,12 @@ def test_openrouter_provider_model_profile(mocker: MockerFixture): google_profile = provider.model_profile('google/gemini-2.5-pro-preview') google_model_profile_mock.assert_called_with('gemini-2.5-pro-preview') assert google_profile is not None - assert google_profile.json_schema_transformer == GoogleJsonSchemaTransformer + assert google_profile.json_schema_transformer == _OpenRouterGoogleJsonSchemaTransformer google_profile = provider.model_profile('google/gemma-3n-e4b-it:free') google_model_profile_mock.assert_called_with('gemma-3n-e4b-it') assert google_profile is not None - assert google_profile.json_schema_transformer == GoogleJsonSchemaTransformer + assert google_profile.json_schema_transformer == _OpenRouterGoogleJsonSchemaTransformer openai_profile = provider.model_profile('openai/o1-mini') openai_model_profile_mock.assert_called_with('o1-mini') @@ -170,3 +173,40 @@ def test_openrouter_provider_model_profile(mocker: MockerFixture): unknown_profile = provider.model_profile('unknown/model') assert unknown_profile is not None assert unknown_profile.json_schema_transformer == OpenAIJsonSchemaTransformer + + +def test_openrouter_google_json_schema_transformer(): + """Test _OpenRouterGoogleJsonSchemaTransformer covers all transformation cases.""" + schema = { + '$schema': 'http://json-schema.org/draft-07/schema#', + 'title': 'TestSchema', + 'type': 'object', + 'properties': { + 'status': {'const': 'active'}, + 'category': {'oneOf': [{'type': 'string'}, {'type': 'integer'}]}, + 'email': {'type': 'string', 'format': 'email', 'description': 'User email'}, + 'date': {'type': 'string', 'format': 'date'}, + }, + } + + transformer = _OpenRouterGoogleJsonSchemaTransformer(schema) + result = transformer.walk() + + # const -> enum conversion + assert result['properties']['status'] == {'enum': ['active'], 'type': 'string'} + + # oneOf -> anyOf conversion + assert 'anyOf' in result['properties']['category'] + assert 'oneOf' not in result['properties']['category'] + + # format -> description with existing description + assert result['properties']['email']['description'] == 'User email (format: email)' + assert 'format' not in result['properties']['email'] + + # format -> description without existing description + assert result['properties']['date']['description'] == 'Format: date' + assert 'format' not in result['properties']['date'] + + # Removed fields + assert '$schema' not in result + assert 'title' not in result