src/api/schemas/query.py

from pydantic import BaseModel, Field

from src.api.schemas.common import SourceCitation


class AskRequest(BaseModel):
    question: str = Field(..., min_length=3, description="The question to ask")
    model: str | None = Field(None, description="LiteLLM model string (e.g. claude-sonnet-4-20250514)")
    max_sources: int = Field(5, ge=1, le=20, description="Maximum source chunks to retrieve")
    author: str | None = Field(None, description="Filter by author name")
    era: str | None = Field(None, description="Filter by era (pre_1900, 1900_1930, 1930_1950, post_1950)")


class AskResponse(BaseModel):
    answer: str
    sources: list[SourceCitation]
    model_used: str


class SearchRequest(BaseModel):
    query: str = Field(..., min_length=2, description="Search query")
    n_results: int = Field(10, ge=1, le=50)
    author: str | None = None
    era: str | None = None