Vizz17/context-aware-rag

1"""Chat endpoint — query the RAG pipeline."""2 3from __future__ import annotations4 5import logging6import time7import uuid8 9from fastapi import APIRouter, HTTPException, Query, Depends10 11from app.core.config import settings12from app.api.auth import get_current_user13from app.models.schemas import ChatRequest, ChatResponse, SourceCitation, ChatHistoryListResponse, ChatSessionResponse14from app.services.llm_chain import generate_answer15from app.services.reranker import rerank16from app.services.retriever import retrieve17from app.services.vector_store import get_vector_store18 19logger = logging.getLogger(__name__)20router = APIRouter()21 22 23@router.post("/chat", response_model=ChatResponse)24def chat(req: ChatRequest, user_id: str = Depends(get_current_user)):25    """Process a user query through the full RAG pipeline."""26    start = time.perf_counter()27 28    try:29        # 1. Retrieve candidates30        t0 = time.perf_counter()31        retrieval_query = req.query32        if req.history and len(req.query.split()) < 20:33            last_user_msgs = [m.get("content", "") for m in req.history if m.get("role") == "user"]34            if last_user_msgs:35                retrieval_query = f"{last_user_msgs[-1]} {req.query}"36                logger.info(f"Expanded retrieval query: '{retrieval_query}'")37                38        candidates = retrieve(39            query=retrieval_query,40            filters=req.filters,41            user_id=user_id,42        )43        logger.info("Retrieve: %.1fs (%d candidates)", time.perf_counter() - t0, len(candidates))44 45        # 2. Rerank46        t1 = time.perf_counter()47        top_chunks = rerank(query=retrieval_query, candidates=candidates)48        logger.info("Rerank: %.1fs (%d chunks)", time.perf_counter() - t1, len(top_chunks))49 50        # 3. Generate grounded answer51        t2 = time.perf_counter()52        answer = generate_answer(query=req.query, chunks=top_chunks, history=req.history)53        logger.info("LLM generate: %.1fs", time.perf_counter() - t2)54 55        # 4. Build source citations56        sources = [57            SourceCitation(58                document=c.get("metadata", {}).get("filename", "unknown"),59                page=c.get("metadata", {}).get("page_num", 0),60                chunk_index=c.get("metadata", {}).get("chunk_index", 0),61                text=c["text"][:300],  # truncate for response62                score=round(c.get("rerank_score", c.get("score", 0.0)), 4),63            )64            for c in top_chunks65        ]66        67        # If the LLM explicitly refused to answer due to lack of context, hide the citations68        if "I don't see the answer" in answer or "I don't have enough information" in answer:69            sources = []70 71        elapsed = (time.perf_counter() - start) * 100072        session_id = req.session_id or str(uuid.uuid4())73        model_name = f"{settings.llm_provider}/{settings.llm_model}"74 75        # 5. Save chat history to vector store76        title = req.query[:50]77        if req.history and len(req.history) > 0:78            title = req.history[0].get("content", req.query)[:50]79            80        updated_history = list(req.history) if req.history else []81        updated_history.append({"role": "user", "content": req.query})82        updated_history.append({83            "role": "assistant",84            "content": answer,85            "sources": [s.model_dump() for s in sources],86            "latency": round(elapsed, 1),87            "model": model_name88        })89        90        doc_scope_id = None91        if req.filters and isinstance(req.filters, dict):92            doc_scope_id = req.filters.get("doc_id")93 94        try:95            store = get_vector_store()96            store.save_chat(session_id, title, updated_history, user_id=user_id, doc_id=doc_scope_id)97        except Exception as e:98            logger.error(f"Failed to save chat history: {e}")99 100        return ChatResponse(101            answer=answer,102            sources=sources,103            session_id=session_id,104            model=model_name,105            latency_ms=round(elapsed, 1),106            retrieval_mode="optimized" if settings.enable_optimizer else "top_k",107        )108 109    except Exception as e:110        error_str = str(e)111        if "429" in error_str or "RESOURCE_EXHAUSTED" in error_str:112            raise HTTPException(113                status_code=429,114                detail=(115                    "Rate limit reached. Please wait 1-2 minutes and try again."116                ),117            )118        119        logger.error(f"Chat error: {e}", exc_info=True)120        raise HTTPException(status_code=500, detail=str(e))121 122@router.get("/chats", response_model=ChatHistoryListResponse)123def list_chats(124    doc_id: str | None = Query(default=None, description="Filter chat sessions by selected document scope"),125    user_id: str = Depends(get_current_user)126):127    """Return all past chat sessions."""128    try:129        store = get_vector_store()130        sessions = store.list_chats(user_id=user_id, doc_id=doc_id)131        return ChatHistoryListResponse(sessions=sessions)132    except Exception as e:133        raise HTTPException(status_code=500, detail=f"Failed to fetch chats: {e}")134 135@router.get("/chats/{session_id}", response_model=ChatSessionResponse)136def get_chat(session_id: str, user_id: str = Depends(get_current_user)):137    """Retrieve a specific chat history."""138    try:139        store = get_vector_store()140        chat = store.get_chat(session_id, user_id=user_id)141        if not chat:142            raise HTTPException(status_code=404, detail="Chat session not found")143        return ChatSessionResponse(**chat)144    except HTTPException:145        raise146    except Exception as e:147        raise HTTPException(status_code=500, detail=f"Failed to fetch chat: {e}")148 149@router.delete("/chats/{session_id}")150def delete_chat(session_id: str, user_id: str = Depends(get_current_user)):151    """Delete a specific chat history."""152    try:153        store = get_vector_store()154        store.delete_chat(session_id, user_id=user_id)155        return {"status": "deleted"}156    except Exception as e:157        raise HTTPException(status_code=500, detail=f"Failed to delete chat: {e}")158