src/schemas.py

"""
In LLMs, you are querying a probabilistic text engine. 
If you ask it for an age, it might give you 25, or it might give you "Twenty-five", or it might say "Based on the data, the user is 25 years old."

If your system expects 25 but gets a whole sentence, your code crashes in production.

**Schemas act as the bouncers at the door of your application. We use Pydantic to define the exact shape of the data we expect.**
"""
from pydantic import BaseModel, Field
from typing import Dict, Any, Optional

# 1. THE CONFIGURATION SCHEMA
# This dictates how our router behaves. Notice how we set smart defaults.
class RouterConfig(BaseModel):
    """Configuration for how the OmniRouter should route the request."""
    provider: str = Field(
        default="openai", 
        description="The LLM provider to use (e.g., 'openai', 'anthropic', 'local')"
    )
    model: str = Field(
        default="gpt-4-turbo", 
        description="The specific model string to use"
    )
    # Defensive Engineering: An LLM temperature cannot be less than 0 or greater than 2.
    # ge = greater than or equal to, le = less than or equal to.
    temperature: float = Field(
        default=0.7, 
        ge=0.0, 
        le=2.0, 
        description="Creativity score for the model"
    )
    max_retries: int = Field(
        default=3, 
        description="How many times to retry on API failure or rate limit"
    )

    # --- NEW CAPABILITY --- Added
    fallback_provider: Optional[str] = Field(
        default=None, 
        description="If the primary provider completely fails, switch to this one"
    )
    fallback_model: Optional[str] = Field(
        default=None,
        description="The model to use for the fallback provider"
    )

# 2. THE STANDARDIZED OUTPUT SCHEMA
# This solves the main pain point from our README. 
# Whether OpenAI or Anthropic answers, the rest of our app gets THIS exact object.
class LLMResponse(BaseModel):
    """The standardized output format returned from ANY provider."""
    content: str = Field(description="The actual text response generated by the LLM")
    provider_used: str = Field(description="Which provider actually generated this response")
    model_used: str = Field(description="The specific model used")
    
    # We track tokens heavily for cost optimization (Week 14 concept)
    prompt_tokens: int = Field(default=0, description="Tokens used in the prompt")
    completion_tokens: int = Field(default=0, description="Tokens used in the completion")
    
    # We will calculate this automatically later
    cost_estimate: float = Field(default=0.0, description="Estimated cost of this call in USD")