"""
RAG-related models
"""

from pydantic import BaseModel, Field
from typing import List, Dict, Any, Optional


class RAGChatRequest(BaseModel):
    """Request model for RAG-based chat"""
    message: str = Field(..., description="User's chat message")
    conversation_id: Optional[str] = Field(None, description="Conversation ID for continuing a conversation")


class RAGChatResponse(BaseModel):
    """Response model for RAG-based chat"""
    reply: str = Field(..., description="Assistant's reply to the user")
    task_id: str = Field(..., description="ID of the task processing the request")
    sources: List[Dict[str, Any]] = Field([], description="Sources used for the response")
    conversation_id: Optional[str] = Field(None, description="Conversation ID for continuing the conversation")


class RAGChatStatusResponse(BaseModel):
    """Response model for RAG chat status check"""
    status: str = Field(..., description="Status of the RAG chat processing")
    progress: float = Field(..., description="Progress percentage of the processing")
    result: Optional[Dict[str, Any]] = Field(None, description="Result of the processing if complete")
    agent_statuses: Optional[List[Dict[str, Any]]] = Field([], description="Status of each agent in the process")