src/serving/models.py

"""
Pydantic request/response schemas for the serving API.

Defines the wire format for chat, RAG, ingestion, and health
endpoints. Keeping schemas in a separate module makes them
importable for client-side validation and testing without
pulling in FastAPI.
"""

from pydantic import BaseModel, Field


class ChatMessage(BaseModel):
    role: str = Field(..., description="One of: system, user, assistant")
    content: str


class ChatRequest(BaseModel):
    messages: list[ChatMessage]
    max_tokens: int = Field(default=512, ge=1, le=4096)
    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
    top_k: int = Field(default=20, ge=0)
    top_p: float = Field(default=0.95, ge=0.0, le=1.0)
    stream: bool = Field(default=True)


class RAGChatRequest(ChatRequest):
    top_k_docs: int = Field(default=3, ge=1, le=10, description="Number of retrieved chunks to inject as context")


class IngestResponse(BaseModel):
    filename: str
    num_chunks: int
    message: str


class HealthResponse(BaseModel):
    status: str
    model: str
    device: str
    num_chunks: int