UnknownPixel/askPESU

1"""FastAPI application for AskPESU backend APIs."""2 3import argparse4import datetime5import logging6import time7from collections.abc import AsyncIterator8from contextlib import asynccontextmanager9 10import pytz11import torch12import uvicorn13from fastapi import FastAPI, Request14from fastapi.responses import FileResponse, JSONResponse15from fastapi.staticfiles import StaticFiles16from google.api_core.exceptions import ResourceExhausted17 18from app.docs import ask_docs, health_docs, index_docs, quota_docs19from app.models import AskRequestModel, AskResponseModel, HealthResponseModel, QuotaResponseModel20from app.quota import QuotaState21from app.rag import RetrievalAugmentedGenerator22 23 24@asynccontextmanager25async def lifespan(app: FastAPI) -> AsyncIterator[None]:26    """Lifespan event handler for startup and shutdown events."""27    # Startup28    logging.info("AskPESU API startup")29 30    # Initialize the RAG engine31    global rag32    config_path = getattr(app.state, "config_path", "conf/config.yaml")33    rag = RetrievalAugmentedGenerator(config_path)34    logging.info("RAG pipeline initialized...")35 36    yield37    # Shutdown38    logging.info("AskPESU API shutdown.")39 40 41app = FastAPI(42    title="askPESU API",43    description="Backend APIs for AskPESU, a question-answering chatbot for PES University.",44    version="0.1.0",45    docs_url="/docs",46    lifespan=lifespan,47    openapi_tags=[48        {49            "name": "Generation",50            "description": "Operations related to generating responses from the chatbot.",51        },52        {53            "name": "Monitoring",54            "description": "Health checks and other monitoring endpoints.",55        },56    ],57)58 59# Initialize globals60DIST_DIR = "frontend/out"  # Directory for static files (built from frontend)61IST = pytz.timezone("Asia/Kolkata")  # Indian Standard Time timezone62rag: RetrievalAugmentedGenerator | None = None  # Global variable to hold the RAG instance63 64# Global state to track if 'thinking' mode is enabled65THINKING_STATE = QuotaState(name="thinking", cooldown_hours=24)66# Global state to track if primary LLM is enabled67PRIMARY_STATE = QuotaState(name="primary", cooldown_hours=24)68 69# Mount static files70app.mount("/static", StaticFiles(directory=DIST_DIR), name="static")71 72 73def get_quota_status() -> dict:74    """Return quota availability for both LLMs."""75    THINKING_STATE.refresh()76    PRIMARY_STATE.refresh()77    return {78        "thinking": THINKING_STATE.status(),79        "primary": PRIMARY_STATE.status(),80    }81 82 83@app.exception_handler(ResourceExhausted)84async def resource_exhausted_exception_handler(_request: Request, exc: ResourceExhausted) -> JSONResponse:85    """Handler for resource exhausted exceptions."""86    logging.warning(f"Quota exceeded: {exc}")87    return JSONResponse(88        status_code=429,89        content={90            "status": False,91            "message": str(exc),92            "quota": get_quota_status(),93            "timestamp": datetime.datetime.now(IST).isoformat(),94        },95    )96 97 98@app.exception_handler(Exception)99async def unhandled_exception_handler(_request: Request, _exc: Exception) -> JSONResponse:100    """Handler for unhandled exceptions."""101    logging.exception("Unhandled exception occurred.")102    return JSONResponse(103        status_code=500,104        content={105            "status": False,106            "message": "Internal Server Error. Please try again later.",107            "timestamp": datetime.datetime.now(IST).isoformat(),108        },109    )110 111 112@app.get(113    "/",114    response_class=FileResponse,115    tags=["Generation"],116    responses=index_docs.response_examples,117)118async def index() -> FileResponse:119    """Serve the main entrypoint (index.html) from the built static files."""120    return FileResponse(f"{DIST_DIR}/index.html")121 122 123@app.post(124    "/ask",125    response_model=AskResponseModel,126    response_class=JSONResponse,127    openapi_extra=ask_docs.request_examples,128    responses=ask_docs.response_examples,129    tags=["Generation"],130)131async def ask(payload: AskRequestModel) -> JSONResponse:132    """Endpoint to handle question-answering requests.133 134    Automatically manages LLM quota with cooldowns.135    May raise 429 if 'thinking' or 'primary' mode is temporarily unavailable.136    """137    global THINKING_STATE, PRIMARY_STATE138    logging.debug(f"Received /ask question: {payload.query}")139    logging.debug(f"Thinking mode: {payload.thinking}")140    current_time = datetime.datetime.now(IST)141 142    # Re-enable thinking mode and primary LLM if cooldown period has expired143    THINKING_STATE.refresh()144    PRIMARY_STATE.refresh()145 146    # Check if thinking mode is requested and enabled147    if payload.thinking and not THINKING_STATE.enabled:148        logging.warning("Thinking mode was requested but currently unavailable due to quota limits.")149        raise ResourceExhausted(150            "Thinking mode is temporarily unavailable due to quota limits. "151            "Please try again later, or disable 'thinking' mode if enabled."152        )153 154    # Check if primary LLM is requested and enabled155    if not payload.thinking and not PRIMARY_STATE.enabled:156        logging.warning("Primary LLM is currently unavailable due to quota limits.")157        raise ResourceExhausted("Primary LLM is temporarily unavailable due to quota limits. Please try again later.")158 159    # Attempt to generate the answer160    start_time = time.perf_counter()161    try:162        answer = await rag.generate(query=payload.query, thinking=payload.thinking, history=payload.history)163    except ResourceExhausted:164        llm_state = THINKING_STATE if payload.thinking else PRIMARY_STATE165        llm_state.disable()166        raise167 168    latency = round(time.perf_counter() - start_time, 3)169    response = AskResponseModel(170        status=True,171        message="Answer generated successfully.",172        answer=answer,173        timestamp=current_time,174        latency=latency,175    )176    return JSONResponse(status_code=200, content=response.model_dump(mode="json", exclude_none=True))177 178 179@app.get(180    "/health",181    response_model=HealthResponseModel,182    response_class=JSONResponse,183    openapi_extra=health_docs.request_examples,184    responses=health_docs.response_examples,185    tags=["Monitoring"],186)187async def health() -> JSONResponse:188    """Health check endpoint."""189    logging.debug("Health check requested.")190    response = HealthResponseModel(191        status=True,192        message="ok",193        timestamp=datetime.datetime.now(IST),194    )195    return JSONResponse(status_code=200, content=response.model_dump(mode="json", exclude_none=True))196 197 198@app.get(199    "/quota",200    response_model=QuotaResponseModel,201    response_class=JSONResponse,202    openapi_extra=quota_docs.request_examples,203    responses=quota_docs.response_examples,204    tags=["Monitoring"],205)206async def quota() -> JSONResponse:207    """Quota status endpoint."""208    logging.debug("Quota status requested.")209    response = QuotaResponseModel(210        status=True,211        quota=get_quota_status(),212        timestamp=datetime.datetime.now(IST),213    )214    return JSONResponse(status_code=200, content=response.model_dump(mode="json", exclude_none=True))215 216 217def main() -> None:218    """Main function to run the FastAPI application with command line arguments."""219    # Set up argument parser for command line arguments220    parser = argparse.ArgumentParser(221        description="Run the FastAPI application for askPESU backend.",222    )223    parser.add_argument(224        "--host",225        type=str,226        default="0.0.0.0",227        help="Host to run the FastAPI application on. Default is 0.0.0.0",228    )229    parser.add_argument(230        "--port",231        type=int,232        default=7860,233        help="Port to run the FastAPI application on. Default is 7860",234    )235    parser.add_argument(236        "--config",237        type=str,238        default="conf/config.yaml",239        help="Path to the configuration YAML file. Default is conf/config.yaml",240    )241    parser.add_argument(242        "--debug",243        action="store_true",244        help="Run the application in debug mode with detailed logging.",245    )246    args = parser.parse_args()247 248    # Store config path in app state for lifespan handler249    app.state.config_path = args.config250 251    # Set up logging configuration252    logging_level = logging.DEBUG if args.debug else logging.INFO253    logging.basicConfig(254        level=logging_level,255        format="%(asctime)s - %(levelname)s - %(filename)s:%(funcName)s:%(lineno)d - %(message)s",256        filemode="w",257    )258 259    # Run the app260    uvicorn.run("app.app:app", host=args.host, port=args.port, reload=args.debug)261 262 263if __name__ == "__main__":264    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")265    logging.info(f"Using device: {device}")266    if device.type == "cuda":267        logging.info(f"CUDA version: {torch.version.cuda}")268        logging.info(f"Number of GPUs: {torch.cuda.device_count()}")269        for i in range(torch.cuda.device_count()):270            logging.info(f"GPU {i} name: {torch.cuda.get_device_name(i)}")271            logging.info(f"\tGPU {i} memory: {torch.cuda.get_device_properties(i).total_memory / 1024**3:.2f} GB")272            logging.info(f"\tGPU {i} memory allocated: {torch.cuda.memory_allocated(i) / 1024**3:.2f} GB")273            logging.info(f"\tGPU {i} memory reserved: {torch.cuda.memory_reserved(i) / 1024**3:.2f} GB")274        torch.set_float32_matmul_precision("high")275    else:276        logging.info("Running without GPU acceleration")277    main()278