src/api/models/api_models.py

from pydantic import BaseModel, Field


# -----------------------
# Core search result model
# -----------------------
class SearchResult(BaseModel):
    title: str = Field(default="", description="Title of the article")
    feed_author: str | None = Field(default=None, description="Author of the article")
    feed_name: str | None = Field(default=None, description="Name of the feed/newsletter")
    article_author: list[str] | None = Field(default=None, description="List of article authors")
    url: str | None = Field(default=None, description="URL of the article")
    chunk_text: str | None = Field(default=None, description="Text content of the article chunk")
    score: float = Field(default=0.0, description="Relevance score of the article")


# -----------------------
# Unique titles request/response
# -----------------------
class UniqueTitleRequest(BaseModel):
    query_text: str = Field(default="", description="The user query text")
    feed_author: str | None = Field(default=None, description="Filter by author name")
    feed_name: str | None = Field(default=None, description="Filter by feed/newsletter name")
    article_author: list[str] | None = Field(default=None, description="List of article authors")
    title_keywords: str | None = Field(
        default=None, description="Keywords or phrase to match in title"
    )
    limit: int = Field(default=5, description="Number of results to return")


class UniqueTitleResponse(BaseModel):
    results: list[SearchResult] = Field(
        default_factory=list, description="List of unique title search results"
    )


# -----------------------
# Ask request model
# -----------------------
class AskRequest(BaseModel):
    query_text: str = Field(default="", description="The user query text")
    feed_author: str | None = Field(default=None, description="Filter by author name")
    feed_name: str | None = Field(default=None, description="Filter by feed/newsletter name")
    article_author: list[str] | None = Field(default=None, description="List of article authors")
    title_keywords: str | None = Field(
        default=None, description="Keywords or phrase to match in title"
    )
    limit: int = Field(default=5, description="Number of results to return")
    provider: str = Field(default="OpenRouter", description="The provider to use for the query")
    model: str | None = Field(
        default=None, description="The specific model to use for the provider, if applicable"
    )


# -----------------------
# Ask response model
# -----------------------
class AskResponse(BaseModel):
    query: str = Field(default="", description="The original query text")
    provider: str = Field(default="", description="The LLM provider used for generation")
    answer: str = Field(default="", description="Generated answer from the LLM")
    sources: list[SearchResult] = Field(
        default_factory=list, description="List of source documents used in generation"
    )
    model: str | None = Field(
        default=None, description="The specific model used by the provider, if available"
    )
    finish_reason: str | None = Field(
        default=None, description="The reason why the generation finished, if available"
    )


# -----------------------
# Streaming "response" documentation
# -----------------------
class AskStreamingChunk(BaseModel):
    delta: str = Field(default="", description="Partial text generated by the LLM")


class AskStreamingResponse(BaseModel):
    query: str = Field(default="", description="The original query text")
    provider: str = Field(default="", description="The LLM provider used for generation")
    chunks: list[AskStreamingChunk] = Field(
        default_factory=list, description="Streamed chunks of generated text"
    )