"""Pydantic models for LLM responses."""

from typing import List, Optional

from pydantic import BaseModel, Field


# ================== Search Response Models ==================


class LLMSearchCitation(BaseModel):
    """A citation from LLM search results."""

    url: str
    title: Optional[str] = None
    text: Optional[str] = None


class LLMSearchResult(BaseModel):
    """A search result from LLM search (Perplexity only)."""

    title: str
    url: str
    snippet: Optional[str] = None


class LLMSearchUsage(BaseModel):
    """Token usage statistics for LLM search."""

    input_tokens: int
    output_tokens: int
    total_tokens: int


class LLMSearchResponse(BaseModel):
    """Response for LLM-powered web search."""

    id: str
    model: str
    content: str
    citations: List[LLMSearchCitation] = Field(default_factory=list)
    search_results: List[LLMSearchResult] = Field(default_factory=list)
    usage: LLMSearchUsage


# ================== Complete Response Models ==================


class LLMUsage(BaseModel):
    """Token usage statistics for LLM completion."""

    input_tokens: int
    output_tokens: int
    total_tokens: int


class LLMCompleteResponse(BaseModel):
    """Response for LLM text completion."""

    id: str
    model: str
    content: str
    stop_reason: str
    usage: LLMUsage
