app/api/schemas.py

from typing import Any

from pydantic import BaseModel, Field


class DocumentSearchRequest(BaseModel):
    query: str = Field(..., description="Search query")
    top_k: int = Field(
        default=10, ge=1, le=50, description="Number of documents to retrieve"
    )


class RetrievedDocument(BaseModel):
    content: str = Field(..., description="Document content")
    metadata: dict[str, Any] = Field(
        default_factory=dict, description="Document metadata"
    )
    score: float | None = Field(None, description="Relevance score")
    chunk_id: int | None = Field(None, description="Chunk identifier")
    source_file: str | None = Field(None, description="Source filename")


class DocumentSearchResponse(BaseModel):
    query: str
    documents: list[RetrievedDocument]
    total_results: int
    retrieval_method: str = "hybrid"


class AnswerGenerationRequest(BaseModel):
    query: str = Field(..., description="Question to answer")
    top_k: int = Field(
        default=10, ge=1, le=50, description="Number of documents to use for context"
    )
    include_sources: bool = Field(
        default=True, description="Include source documents in response"
    )
    temperature: float = Field(default=0.1, ge=0, le=1, description="LLM temperature")
    max_tokens: int = Field(
        default=1000, ge=100, le=4000, description="Maximum tokens in response"
    )


class AnswerGenerationResponse(BaseModel):
    query: str
    answer: str
    sources: list[RetrievedDocument] | None = None
    retrieval_method: str
    documents_used: int