Created
December 2, 2025 00:50
-
-
Save dsfaccini/4f5eb50b3ea8c618fc2b8623e4f09da5 to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| """ | |
| Fixed script with custom transformer restoring v1.19.0 behavior. | |
| This uses the OLD behavior where: | |
| - $defs ARE inlined (prefer_inlined_defs=True) | |
| - Nullable unions are simplified to nullable types (simplify_nullable_unions=True) | |
| This produces better output from Gemini 2.5-flash. | |
| """ | |
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| from enum import Enum | |
| from typing import Literal | |
| import logfire | |
| from pydantic import BaseModel, Field | |
| from pydantic_ai import Agent, RunContext | |
| from pydantic_ai._json_schema import JsonSchema, JsonSchemaTransformer | |
| from pydantic_ai.models.openai import OpenAIChatModel | |
| from pydantic_ai.profiles import ModelProfile | |
| from pydantic_ai.profiles.google import GoogleModelProfile | |
| logfire.configure() | |
| logfire.instrument_pydantic_ai() | |
| class LegacyGoogleJsonSchemaTransformer(JsonSchemaTransformer): | |
| """Restores v1.19.0 behavior for Google JSON schemas. | |
| Key differences from current transformer: | |
| - prefer_inlined_defs=True: Inlines $defs instead of using $ref | |
| - simplify_nullable_unions=True: Converts anyOf[{$ref}, {type:null}] to nullable type | |
| - Converts enums to string type (old Gemini requirement) | |
| - Converts oneOf to anyOf | |
| """ | |
| 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] | |
| # OLD BEHAVIOR: Convert enums to string type | |
| if enum := schema.get("enum"): | |
| schema["type"] = "string" | |
| schema["enum"] = [str(val) for val in enum] | |
| # OLD BEHAVIOR: 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 legacy_google_model_profile(model_name: str) -> ModelProfile | None: | |
| """Profile that uses the legacy (v1.19.0) transformer behavior.""" | |
| is_image_model = "image" in model_name | |
| is_3_or_newer = "gemini-3" in model_name | |
| return GoogleModelProfile( | |
| json_schema_transformer=LegacyGoogleJsonSchemaTransformer, # <-- The fix | |
| supports_image_output=is_image_model, | |
| supports_json_schema_output=is_3_or_newer or not is_image_model, | |
| supports_json_object_output=is_3_or_newer or not is_image_model, | |
| supports_tools=not is_image_model, | |
| google_supports_native_output_with_builtin_tools=is_3_or_newer, | |
| ) | |
| class LevelType(str, Enum): | |
| ground = "ground" | |
| basement = "basement" | |
| floor = "floor" | |
| attic = "attic" | |
| mezzanine = "mezzanine" | |
| class SpaceType(str, Enum): | |
| entryway = "entryway" | |
| living_room = "living-room" | |
| kitchen = "kitchen" | |
| bedroom = "bedroom" | |
| bathroom = "bathroom" | |
| garage = "garage" | |
| class InsertLevelArg(BaseModel): | |
| level_name: str = Field(description="Nom du niveau, tel que fourni explicitement par l'utilisateur.") | |
| level_type: LevelType = Field(description="Type de niveau parmi les choix possibles.") | |
| class SpaceArg(BaseModel): | |
| space_name: str = Field(description="Nom de l'espace.") | |
| space_type: SpaceType = Field(description="Type d'espace parmi les choix possibles.") | |
| class InsertLevelWithSpacesArgs(BaseModel): | |
| level: InsertLevelArg | None = Field( | |
| default=None, | |
| description="Définition du niveau à créer (si None, les espaces sont annexes)", | |
| ) | |
| spaces: list[SpaceArg] = Field(description="Liste des espaces à créer dans ce niveau.") | |
| class DwellingUpdate(BaseModel): | |
| type: Literal["house", "apartment"] | |
| class Deps: | |
| pass | |
| model = OpenAIChatModel( | |
| "google/gemini-2.5-flash", | |
| provider="openrouter", | |
| profile=legacy_google_model_profile, # <-- Apply the fix | |
| ) | |
| agent = Agent( | |
| model, | |
| deps_type=Deps, | |
| system_prompt="Tu es un assistant qui aide à configurer des logements.", | |
| ) | |
| @agent.tool | |
| def update_dwelling_properties(ctx: RunContext[Deps], update: DwellingUpdate) -> str: | |
| """Met à jour les propriétés du logement.""" | |
| return f"Logement mis à jour: {update.type}" | |
| @agent.tool | |
| def insert_level_with_spaces(ctx: RunContext[Deps], args: InsertLevelWithSpacesArgs) -> str: | |
| """Insère un niveau avec ses espaces.""" | |
| return f"Niveau créé: {args.level}, espaces: {args.spaces}" | |
| async def main(): | |
| with logfire.span("google_schema_fixed"): | |
| result = await agent.run( | |
| "C'est une maison avec un rez-de-chaussée qui a une entrée, un salon et un garage.", | |
| deps=Deps(), | |
| ) | |
| print(f"Result: {result.output}") | |
| if __name__ == "__main__": | |
| import asyncio | |
| asyncio.run(main()) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment