app.py

import time
import traceback
from fastapi import FastAPI, HTTPException, status, Request
from rich import print as rprint
from rich.panel import Panel

from models import QueryRequest, QueryResponse, Question
from query_service import QueryService

app = FastAPI(
    title="High-Performance RAG API",
    description="Processes documents and answers a list of questions efficiently with intelligent caching.",
    version="3.0.0",
)

query_service = QueryService()

@app.middleware("http")
async def add_process_time_header(request: Request, call_next):
    start_time = time.time()
    response = await call_next(request)
    process_time = time.time() - start_time
    response.headers["X-Process-Time"] = f"{process_time:.4f} sec"
    rprint(f"Request '{request.method} {request.url.path}' completed in {process_time:.4f}s")
    return response

@app.post(
    "/api/v1/hackrx/run",
    response_model=QueryResponse,
    tags=["RAG Pipeline"],
    summary="Process a Document and Answer a Batch of Questions",
    status_code=status.HTTP_200_OK
)
async def run_submission(request_body: QueryRequest):
    try:
        rprint(Panel(f"New Query Request"))
        questions = [Question(question=q) for q in request_body.questions]
        results = query_service.process_queries(str(request_body.documents), questions)
        final_answers = [result.answer for result in results]
        return QueryResponse(answers=final_answers)
    except ValueError as e:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e))
    except Exception as e:
        tb_str = traceback.format_exc()
        rprint(Panel(f"[bold red]Querying failed:[/bold red]\n{tb_str}", title="[red]Server Error[/red]"))
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="An internal server error occurred.")

@app.get("/health", tags=["Monitoring"], summary="API Health Check")
def health_check():
    return {"status": "ok"}