lionel-gemini-flash-2.5-via-openrouter-legacy-transformer.py

"""
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())