"""Chat API endpoints for RAG chatbot.""" from typing import Optional from fastapi import APIRouter, HTTPException, Depends from pydantic import BaseModel, Field from src.rag.engine import RAGEngine from src.rag.document_store import DocumentStore router = APIRouter(prefix="/chat", tags=["chat"]) # Initialize RAG engine rag_engine = RAGEngine() class ChatMessage(BaseModel): """Chat message request model.""" message: str = Field(..., min_length=1, max_length=2000) session_id: Optional[str] = None selected_text: Optional[str] = None conversation_history: Optional[list] = None language: Optional[str] = "en" # Response language (en, ur, ur-PK, ar, es, ...) class ChatResponse(BaseModel): """Chat response model.""" answer: str sources: list citations: list query: str confidence: float has_sources: bool source_count: int timestamp: str session_id: Optional[str] = None class IndexDocumentRequest(BaseModel): """Request to index a document.""" content: str = Field(..., min_length=1) title: Optional[str] = None source_url: Optional[str] = None file_path: Optional[str] = None section: Optional[str] = None tags: Optional[list] = None metadata: Optional[dict] = None class IndexDocumentResponse(BaseModel): """Response for document indexing.""" success: bool document_id: str message: str @router.post("/message", response_model=ChatResponse) async def send_message(request: ChatMessage) -> ChatResponse: """Send a message to the RAG chatbot. Args: request: Chat message request. Returns: ChatResponse with answer and sources. """ try: # Initialize RAG engine await rag_engine.initialize() # Query using RAG engine with conversation history result = await rag_engine.query( question=request.message, conversation_history=request.conversation_history, selected_text=request.selected_text, top_k=8, include_citations=True, language=request.language, ) return ChatResponse( answer=result["answer"], sources=result["sources"], citations=result["citations"], query=result["query"], confidence=result["confidence"], has_sources=result["has_sources"], source_count=result["source_count"], timestamp=result["timestamp"], session_id=request.session_id, ) except Exception as e: raise HTTPException( status_code=500, detail=f"Error processing chat message: {str(e)}", ) @router.post("/selected-text", response_model=ChatResponse) async def query_selected_text(request: ChatMessage) -> ChatResponse: """Query about selected text content. Args: request: Chat message with selected text. Returns: ChatResponse with contextual answer. """ if not request.selected_text: raise HTTPException( status_code=400, detail="selected_text is required for this endpoint", ) try: await rag_engine.initialize() result = await rag_engine.query( question=request.message, selected_text=request.selected_text, conversation_history=request.conversation_history, include_citations=True, language=request.language, ) return ChatResponse( answer=result["answer"], sources=result["sources"], citations=result["citations"], query=result["query"], confidence=result["confidence"], has_sources=result["has_sources"], source_count=result["source_count"], timestamp=result["timestamp"], session_id=request.session_id, ) except Exception as e: raise HTTPException( status_code=500, detail=f"Error processing selected text query: {str(e)}", ) @router.post("/index", response_model=IndexDocumentResponse) async def index_document(request: IndexDocumentRequest) -> IndexDocumentResponse: """Index a document for RAG retrieval. Args: request: Document indexing request. Returns: IndexDocumentResponse with document ID. """ try: document_store = DocumentStore() doc_id = await document_store.add_document( content=request.content, title=request.title, source_url=request.source_url, file_path=request.file_path, section=request.section, tags=request.tags, metadata=request.metadata, ) return IndexDocumentResponse( success=True, document_id=doc_id, message="Document indexed successfully", ) except Exception as e: raise HTTPException( status_code=500, detail=f"Error indexing document: {str(e)}", ) @router.get("/health") async def chat_health(): """Check RAG chatbot health.""" return { "status": "healthy", "service": "rag-chatbot", "collection": rag_engine.collection_name, }