ruby2210/rag-chatbot

1"""2Chat endpoints for the RAG Chatbot application.3Implements the API endpoints for chat functionality.4"""5from fastapi import APIRouter, HTTPException, Depends, Query6from pydantic import BaseModel, Field7from typing import Optional, List, Dict, Any8from uuid import uuid49from sqlalchemy.orm import Session10from ..services.rag_service import rag_service11from ..services.agent_service import agent_service12from ..utils.database import get_db13from ..utils.config import settings14from ..utils.logging import get_logger15from ..models.chat_history import ChatHistory16from datetime import datetime17 18 19logger = get_logger(__name__)20router = APIRouter()21 22 23class ChatRequest(BaseModel):24    """25    Request model for chat endpoint.26    """27    message: str = Field(..., description="The user's message/query")28    session_id: Optional[str] = Field(None, description="Session identifier; if not provided, a new session is created")29    context_type: str = Field("full_book", description="Type of context to use: 'full_book' or 'selected_text_only'")30 31 32class ChatResponse(BaseModel):33    """34    Response model for chat endpoint.35    """36    response: str = Field(..., description="The chatbot's response to the user's query")37    session_id: str = Field(..., description="The session identifier")38    context_type: str = Field(..., description="The type of context used")39    retrieved_sources: List[Dict[str, Any]] = Field(default=[], description="List of sources used to generate the response")40 41 42class SelectedTextChatRequest(BaseModel):43    """44    Request model for selected-text chat endpoint.45    """46    message: str = Field(..., description="The user's message/query")47    selected_text: str = Field(..., description="The text selected by the user that will be the only context")48    session_id: Optional[str] = Field(None, description="Session identifier; if not provided, a new session is created")49    context_type: str = Field("selected_text_only", description="Must be 'selected_text_only' for this endpoint")50 51 52class HistoryMessage(BaseModel):53    """54    Model for individual messages in chat history.55    """56    message_id: str = Field(..., description="Unique identifier for the message")57    role: str = Field(..., description="Either 'user' or 'assistant'")58    content: str = Field(..., description="The actual message content")59    timestamp: str = Field(..., description="ISO format timestamp")60    context_type: str = Field(..., description="The type of context used for this message")61 62 63class ChatHistoryResponse(BaseModel):64    """65    Response model for chat history endpoint.66    """67    session_id: str = Field(..., description="The session identifier")68    history: List[HistoryMessage] = Field(..., description="List of messages in chronological order")69 70 71class ClearChatRequest(BaseModel):72    """73    Request model for clear chat endpoint.74    """75    session_id: str = Field(..., description="The session identifier to clear")76 77 78class ClearChatResponse(BaseModel):79    """80    Response model for clear chat endpoint.81    """82    session_id: str = Field(..., description="The session identifier that was cleared")83    status: str = Field(..., description="Confirmation status")84    message: str = Field(..., description="Human-readable confirmation message")85 86 87@router.post("/chat", response_model=ChatResponse)88async def chat_endpoint(request: ChatRequest):89    """90    Chat endpoint that handles both full-book and selected-text queries.91    """92    try:93        # Validate context type94        if request.context_type not in ["full_book", "selected_text_only"]:95            raise HTTPException(status_code=400, detail="context_type must be 'full_book' or 'selected_text_only'")96 97        # Create or use provided session ID98        session_id = request.session_id or agent_service.create_session()99 100        # Process the query based on context type101        if request.context_type == "full_book":102            result = rag_service.full_book_query(session_id, request.message)103        else:104            # For selected_text_only, we expect this to be handled by the dedicated endpoint105            # but we'll support it here for flexibility106            result = {107                "response": "Selected text queries should use the /chat/selected-text endpoint",108                "session_id": session_id,109                "context_type": request.context_type,110                "retrieved_sources": []111            }112 113        return ChatResponse(**result)114 115    except Exception as e:116        logger.error(f"Error in chat endpoint: {e}")117        raise HTTPException(status_code=500, detail=f"Error processing chat request: {str(e)}")118 119 120@router.post("/chat/selected-text", response_model=ChatResponse)121async def selected_text_chat_endpoint(request: SelectedTextChatRequest):122    """123    Chat endpoint specifically for selected-text-only queries.124    Enforces strict context isolation to only the provided selected text.125    """126    try:127        # Validate that selected_text is provided128        if not request.selected_text or not request.selected_text.strip():129            raise HTTPException(status_code=400, detail="selected_text is required and cannot be empty")130 131        # Create or use provided session ID132        session_id = request.session_id or agent_service.create_session()133 134        # Process the selected-text query135        result = rag_service.selected_text_query(session_id, request.message, request.selected_text)136 137        return ChatResponse(**result)138 139    except Exception as e:140        logger.error(f"Error in selected-text chat endpoint: {e}")141        raise HTTPException(status_code=500, detail=f"Error processing selected-text chat request: {str(e)}")142 143 144@router.get("/chat/history", response_model=ChatHistoryResponse)145async def get_chat_history(session_id: str = Query(..., description="The session identifier")):146    """147    Retrieve the chat history for a specific session.148    """149    try:150        from sqlalchemy.orm import Session151        from ..utils.database import SessionLocal152        from ..models.chat_history import ChatHistory as ChatHistoryModel153 154        db = SessionLocal()155        try:156            # Query the database for chat history for this session with all details157            history_records = db.query(ChatHistoryModel).filter(158                ChatHistoryModel.session_id == session_id159            ).order_by(ChatHistoryModel.timestamp).all()160 161            # Format the history for the response162            history_messages = []163            for record in history_records:164                history_messages.append(HistoryMessage(165                    message_id=record.message_id,166                    role=record.role,167                    content=record.content,168                    timestamp=record.timestamp.isoformat() if record.timestamp else datetime.utcnow().isoformat(),169                    context_type=record.query_context_type170                ))171 172            return ChatHistoryResponse(173                session_id=session_id,174                history=history_messages175            )176        finally:177            db.close()178 179    except Exception as e:180        logger.error(f"Error in get chat history endpoint: {e}")181        raise HTTPException(status_code=500, detail=f"Error retrieving chat history: {str(e)}")182 183 184@router.post("/chat/clear", response_model=ClearChatResponse)185async def clear_chat_endpoint(request: ClearChatRequest):186    """187    Endpoint to clear chat history for a specific session and reset the session state.188    """189    try:190        # Clear the session191        agent_service.clear_session(request.session_id)192 193        # Return confirmation194        return ClearChatResponse(195            session_id=request.session_id,196            status="cleared",197            message="Chat history has been cleared and session reset"198        )199 200    except Exception as e:201        logger.error(f"Error in clear chat endpoint: {e}")202        raise HTTPException(status_code=500, detail=f"Error clearing chat: {str(e)}")