from pydantic import BaseModel, Field
from litellm import Router
from dotenv import load_dotenv
import instructor

load_dotenv() # Gemini API key


router = Router(
    model_list=[
        {
            "model_name": "gemini-2.0-flash",
            "litellm_params": {
                "model": "gemini/gemini-2.0-flash-001",
            },
        },
        {
                "model_name": "gemini-2.5-flash",
                "litellm_params": {
                    "model": "gemini/gemini-2.5-flash-preview-04-17",
                },
            },
    ],
    cache_responses=False,
    set_verbose=True,
    debug_level="DEBUG",
)

class PersonalInfo(BaseModel):
    name: str = Field(..., description="The person's name.")
    age: int = Field(..., description="The person")
    addres: str = Field("unknown", description="The person's address.") # Default fields other than None do not work for gemini
    preoccupation: str | None = Field(None, description="The person's occupation.")

messages = [
    {
        "role": "system",
        "content": "You are a helpful assistant specialized in extracting data in structured JSON formats.."
    },
    {
        "role": "user",
        "content": "John Doe is 87 years old."
    },
]

llm = instructor.from_litellm(
    router.completion,
    mode=instructor.Mode.GEMINI_JSON,
)

response: PersonalInfo = await llm.chat.completions.create(
    model="gemini-2.0-flash",
    response_model=PersonalInfo,
    messages=messages,
)

AssertionError: Gemini `model` must be set while patching the client, not passed as a parameter to the create method