app/models.py

from typing import List, Optional, Literal
from pydantic import BaseModel, Field


class ChatMessage(BaseModel):
    """Represents a single chat message."""
    role: Literal["system", "user", "assistant"] = Field(..., description="Role of the message sender")
    content: str = Field(..., description="Content of the message")
    
    class Config:
        json_schema_extra = {
            "example": {
                "role": "user",
                "content": "Hello, how are you today?"
            }
        }


class ChatRequest(BaseModel):
    """Request model for chat completion."""
    messages: List[ChatMessage] = Field(..., description="List of chat messages")
    model: str = Field(default="llama-2-7b-chat", description="Model to use for generation")
    max_tokens: int = Field(default=2048, ge=1, le=4096, description="Maximum tokens to generate")
    temperature: float = Field(default=0.7, ge=0.0, le=2.0, description="Sampling temperature")
    top_p: float = Field(default=0.9, ge=0.0, le=1.0, description="Top-p sampling parameter")
    stream: bool = Field(default=True, description="Whether to stream the response")
    
    class Config:
        json_schema_extra = {
            "example": {
                "messages": [
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": "Hello, how are you today?"}
                ],
                "model": "llama-2-7b-chat",
                "max_tokens": 100,
                "temperature": 0.7,
                "stream": True
            }
        }


class ChatResponse(BaseModel):
    """Response model for chat completion."""
    id: str = Field(..., description="Unique response ID")
    object: str = Field(default="chat.completion", description="Object type")
    created: int = Field(..., description="Unix timestamp of creation")
    model: str = Field(..., description="Model used for generation")
    choices: List[dict] = Field(..., description="Generated choices")
    usage: Optional[dict] = Field(None, description="Token usage statistics")


class ModelInfo(BaseModel):
    """Model information response."""
    id: str = Field(..., description="Model ID")
    object: str = Field(default="model", description="Object type")
    created: int = Field(..., description="Unix timestamp of creation")
    owned_by: str = Field(default="huggingface", description="Model owner")


class ErrorResponse(BaseModel):
    """Error response model."""
    error: dict = Field(..., description="Error details")
    
    class Config:
        json_schema_extra = {
            "example": {
                "error": {
                    "message": "Invalid request parameters",
                    "type": "invalid_request_error",
                    "code": 400
                }
            }
        }