src/api/schemas.py

from __future__ import annotations
"""Pydantic schemas for API requests and responses."""

from datetime import datetime
from typing import Any, List, Optional, Dict

from pydantic import BaseModel, Field


class QueryRequest(BaseModel):
    """Request schema for query endpoint."""

    query: str = Field(..., description="The user's question", min_length=1, max_length=2000)
    enable_search: bool = Field(True, description="Enable web search")
    enable_scraping: bool = Field(True, description="Enable web scraping for deep content")
    max_sources: int = Field(5, description="Maximum number of sources to use", ge=1, le=20)
    stream: bool = Field(False, description="Enable streaming response")
    conversation_id: Optional[str] = Field(None, description="Conversation ID for context")


class SourceInfo(BaseModel):
    """Source citation information."""

    title: str
    url: str
    snippet: str = ""


class QueryResponse(BaseModel):
    """Response schema for query endpoint."""

    answer: str = Field(..., description="The generated answer")
    sources: List[SourceInfo] = Field(default_factory=list, description="Sources used")
    follow_up_questions: List[str] = Field(default_factory=list, description="Suggested follow-up questions")
    confidence: float = Field(..., description="Confidence score", ge=0, le=1)
    conversation_id: Optional[str] = Field(None, description="Conversation ID for follow-up")
    processing_time_ms: int = Field(..., description="Processing time in milliseconds")
    metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")


class StreamingChunk(BaseModel):
    """Schema for streaming response chunks."""

    type: str = Field(..., description="Chunk type: 'start', 'content', 'source', 'done'")
    content: str = Field("", description="Content for this chunk")
    source: Optional[SourceInfo] = Field(None, description="Source info (for source chunks)")
    metadata: Dict[str, Any] = Field(default_factory=dict, description="Chunk metadata")


class ErrorResponse(BaseModel):
    """Error response schema."""

    error: str = Field(..., description="Error message")
    error_code: str = Field(..., description="Error code")
    details: Dict[str, Any] = Field(default_factory=dict, description="Error details")


class HealthResponse(BaseModel):
    """Health check response."""

    status: str = Field(..., description="Service status")
    version: str = Field(..., description="API version")
    timestamp: datetime = Field(default_factory=datetime.utcnow, description="Response timestamp")
    components: Dict[str, bool] = Field(default_factory=dict, description="Component health status")


class ConversationMessage(BaseModel):
    """A message in a conversation."""

    role: str = Field(..., description="Message role: 'user' or 'assistant'")
    content: str = Field(..., description="Message content")
    timestamp: datetime = Field(default_factory=datetime.utcnow, description="Message timestamp")


class ConversationHistory(BaseModel):
    """Conversation history response."""

    conversation_id: str
    messages: List[ConversationMessage]
    created_at: datetime
    updated_at: datetime