Muhammad-Umer-Khan/PersonalAssistant

1import logging
2from fastapi import FastAPI, HTTPException
3from fastapi.middleware.cors import CORSMiddleware
4from pydantic import BaseModel
5from rag_pipeline import CustomDocChatbot
6from logger import logging
7
8logger = logging.getLogger(__name__)
9
10# Initialize FastAPI app with descriptive title
11app = FastAPI(title="Muhammad Umer Khan's RAG Bot")
12
13# Enable CORS for React frontend compatibility
14app.add_middleware(
15    CORSMiddleware,
16    allow_origins=["http://localhost:3000", "https://portfolio-sigma-mocha-67.vercel.app"],
17    allow_credentials=True,
18    allow_methods=["*"],
19    allow_headers=["*"],
20)
21
22# Initialize chatbot instance
23try:
24    chatbot = CustomDocChatbot()
25    logger.info({"message": "🤖 Chatbot initialized successfully"})
26except Exception as e:
27    logger.critical({"message": f"❌ Failed to initialize chatbot: {str(e)}"})
28    raise
29
30# Define request model for /chat endpoint
31class QueryRequest(BaseModel):
32    """Pydantic model for validating chat query requests."""
33    query: str
34
35@app.get("/")
36async def root():
37    """Root endpoint returning a welcome message."""
38    return {"message": "Hello, I am Muhammad Umer Khan's AI Bot! 🤖"}
39
40@app.post("/chat")
41async def chat(request: QueryRequest):
42    """
43    Handle chat queries with rate limiting and caching.
44    
45    Args:
46        request (QueryRequest): JSON payload with the user's query.
47    
48    Returns:
49        dict: Response containing the chatbot's reply.
50    
51    Raises:
52        HTTPException: If the query is invalid or processing fails.
53    """
54    try:
55        response = await chatbot.query(request.query)
56        logger.info({"message": f"💬 Query processed: {request.query} | Response: {response}"})
57        return {"reply": response}
58    except Exception as e:
59        logger.error({"message": f"❌ API error: {str(e)}"})
60        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")
61
62@app.get("/health")
63async def health_check():
64    """
65    Health check endpoint to verify LLM and vector store status.
66    
67    Returns:
68        dict: Status indicating if the chatbot is operational.
69    
70    Raises:
71        HTTPException: If critical components are not initialized.
72    """
73    if hasattr(chatbot, 'qa_chain') and hasattr(chatbot, 'vector_db'):
74        logger.info({"message": "✅ Health check passed"})
75        return {"status": "healthy"}
76    logger.error({"message": "❌ Health check failed"})
77    raise HTTPException(status_code=503, detail="Service Unavailable")
78
79@app.on_event("shutdown")
80async def shutdown_event():
81    """Clean up resources on application shutdown."""
82    await chatbot.shutdown()
83    logger.info({"message": "🛑 Application shutdown gracefully"})
84
85if __name__ == "__main__":
86    import uvicorn
87    logger.info({"message": "🚀 Starting FastAPI server on port 8000"})
88    uvicorn.run(app, host="0.0.0.0", port=8000)