schemas.py

from pydantic import BaseModel, EmailStr

# Authentication
class Token(BaseModel):
    access_token: str
    token_type: str

class TokenData(BaseModel):
    username: str | None = None

class User(BaseModel):
    username: str
    email: EmailStr | None = None

class UserInDB(User):
    hashed_password: str

class LoginRequest(BaseModel):
    username: str
    password: str

# Generation
class GenerationRequest(BaseModel):
    prompt: str
    max_length: int = 100
    temperature: float = 0.7
    top_k: int = 50
    top_p: float = 0.95

class GenerationResponse(BaseModel):
    generated_text: str

# WebSocket
class WebSocketMessage(BaseModel):
    prompt: str

# Error Handling
class HTTPError(BaseModel):
    detail: str

    class Config:
        schema_extra = {
            "example": {"detail": "Error message"},
        }

# Rate Limiting
class RateLimitResponse(BaseModel):
    message: str
    retry_after: int