Kelvin-programmer/rag-chatbot

1"""FastAPI application for the RAG Chatbot API."""2 3import logging4import os5import tempfile6import time7from collections import defaultdict8from contextlib import asynccontextmanager9 10from fastapi import FastAPI, File, HTTPException, Request, UploadFile11from fastapi.middleware.cors import CORSMiddleware12from fastapi.responses import FileResponse13from fastapi.staticfiles import StaticFiles14from starlette.middleware.base import BaseHTTPMiddleware15from starlette.responses import JSONResponse16 17from .config import Settings18from .rag_engine import RAGEngine19from .schemas import (20    DocumentsResponse,21    HealthResponse,22    IngestResponse,23    QueryRequest,24    QueryResponse,25)26 27settings = Settings()28 29logging.basicConfig(30    level=logging.DEBUG if settings.debug else logging.INFO,31    format="%(asctime)s | %(levelname)-8s | %(name)s | %(message)s",32)33logger = logging.getLogger(__name__)34 35 36# ---------------------------------------------------------------------------37# Rate-limiting middleware38# ---------------------------------------------------------------------------39 40class RateLimitMiddleware(BaseHTTPMiddleware):41    """Simple sliding-window rate limiter keyed by client IP."""42 43    def __init__(self, app, max_requests: int = 10, window_seconds: int = 60):44        super().__init__(app)45        self.max_requests = max_requests46        self.window_seconds = window_seconds47        self._hits: dict[str, list[float]] = defaultdict(list)48 49    async def dispatch(self, request: Request, call_next):50        client_ip = request.client.host if request.client else "unknown"51        now = time.time()52 53        # Prune timestamps outside the current window54        self._hits[client_ip] = [55            t for t in self._hits[client_ip] if now - t < self.window_seconds56        ]57 58        if len(self._hits[client_ip]) >= self.max_requests:59            return JSONResponse(60                status_code=429,61                content={"detail": "Rate limit exceeded. Try again later."},62            )63 64        self._hits[client_ip].append(now)65        return await call_next(request)66 67 68# ---------------------------------------------------------------------------69# Application lifecycle70# ---------------------------------------------------------------------------71 72engine: RAGEngine | None = None73 74 75@asynccontextmanager76async def lifespan(app: FastAPI):77    global engine78    logger.info("Starting %s ...", settings.app_name)79    engine = RAGEngine(settings)80    engine.vector_store.load()81    logger.info("Ready — %d documents loaded", engine.vector_store.count)82    yield83    logger.info("Shutting down")84 85 86# ---------------------------------------------------------------------------87# FastAPI app88# ---------------------------------------------------------------------------89 90app = FastAPI(91    title=settings.app_name,92    description="Production-grade Retrieval-Augmented Generation chatbot API",93    version="1.0.0",94    lifespan=lifespan,95)96 97# Serve the frontend98_static_dir = os.path.join(os.path.dirname(__file__), "static")99app.mount("/static", StaticFiles(directory=_static_dir), name="static")100 101 102@app.get("/", include_in_schema=False)103async def root():104    return FileResponse(os.path.join(_static_dir, "index.html"))105 106app.add_middleware(107    CORSMiddleware,108    allow_origins=settings.cors_origins,109    allow_credentials=True,110    allow_methods=["*"],111    allow_headers=["*"],112)113app.add_middleware(114    RateLimitMiddleware,115    max_requests=settings.rate_limit_requests,116    window_seconds=settings.rate_limit_window,117)118 119 120# ---------------------------------------------------------------------------121# Routes122# ---------------------------------------------------------------------------123 124 125@app.get("/api/v1/health", response_model=HealthResponse, tags=["System"])126async def health_check():127    """Return service health and basic stats."""128    return HealthResponse(129        status="healthy",130        version="1.0.0",131        documents_loaded=engine.vector_store.count if engine else 0,132    )133 134 135@app.post("/api/v1/query", response_model=QueryResponse, tags=["Chat"])136async def query_documents(body: QueryRequest):137    """Ask a question about the loaded documents."""138    if not engine or engine.vector_store.count == 0:139        raise HTTPException(140            status_code=400, detail="No documents loaded. Upload a PDF first."141        )142 143    result = engine.query(body.question)144    return result145 146 147@app.post(148    "/api/v1/documents/upload", response_model=IngestResponse, tags=["Documents"]149)150async def upload_document(file: UploadFile = File(...)):151    """Upload a PDF and ingest it into the vector store."""152    if not file.filename or not file.filename.lower().endswith(".pdf"):153        raise HTTPException(status_code=400, detail="Only PDF files are supported.")154 155    size = 0156    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")157    try:158        while chunk := await file.read(8192):159            size += len(chunk)160            if size > settings.max_upload_size_mb * 1024 * 1024:161                tmp.close()162                os.unlink(tmp.name)163                raise HTTPException(164                    status_code=413,165                    detail=f"File exceeds {settings.max_upload_size_mb} MB limit.",166                )167            tmp.write(chunk)168        tmp.close()169 170        count = engine.ingest_pdf(tmp.name)171        return IngestResponse(172            message=f"Processed {file.filename}",173            chunks_added=count,174            total_documents=engine.vector_store.count,175        )176    finally:177        if os.path.exists(tmp.name):178            os.unlink(tmp.name)179 180 181@app.get("/api/v1/documents", response_model=DocumentsResponse, tags=["Documents"])182async def get_document_stats():183    """Return the number of indexed document chunks."""184    return DocumentsResponse(185        total_documents=engine.vector_store.count if engine else 0,186        persist_dir=settings.vector_store_path,187    )188 189 190@app.delete("/api/v1/documents", tags=["Documents"])191async def clear_documents():192    """Remove all documents from the vector store."""193    if engine:194        engine.vector_store.clear()195    return {"message": "All documents cleared."}196