""" FastAPI application for FinBot RAG system. Exposes HTTP endpoints for chat, user management, and system diagnostics. """ import logging import os from contextlib import asynccontextmanager from fastapi import FastAPI, HTTPException, Request from typing import Optional from fastapi.responses import JSONResponse from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel from pipeline.rag_pipeline import get_rag_pipeline from retrieval.user_auth import get_user_manager from vector_store import get_vector_store from ingestion.document_ingester import DocumentIngester from config import DocumentCollection # Setup logging logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", ) logger = logging.getLogger(__name__) # Logging setup remains as is # ==================== # REQUEST/RESPONSE MODELS # ==================== class ChatRequest(BaseModel): """Request model for chat endpoint.""" user_role: str query: str user_id: str = None class ChatResponse(BaseModel): """Response model for chat endpoint.""" answer: str sources: list route: str user_role: str accessible_collections: list guardrail_flags: list = [] guardrail_warnings: list = [] rbac_denied: bool = False rbac_reason: Optional[str] = None class UserInfo(BaseModel): """User information model.""" username: str name: str role: str department: str accessible_collections: list[str] = [] class CollectionInfo(BaseModel): """Collection information model.""" name: str description: str accessible_roles: list # ==================== # INITIALIZATION # ==================== async def startup_event(): """Initialize application on startup.""" logger.info("="*60) logger.info("FinBot RAG System Starting Up") logger.info("="*60) # Check for API key if not os.getenv("GROQ_API_KEY"): logger.warning("GROQ_API_KEY not set! Chat functionality will fail.") # Initialize vector store and check collections vector_store = get_vector_store() collections = vector_store.list_collections() logger.info(f"Available collections: {collections if collections else 'None (ingestion pending)'}") logger.info("FinBot RAG System Ready") logger.info("="*60) async def shutdown_event(): """Cleanup on application shutdown.""" logger.info("FinBot RAG System Shutting Down") @asynccontextmanager async def lifespan(app: FastAPI): """Manage application lifecycle.""" await startup_event() yield await shutdown_event() # ==================== # CREATE FASTAPI APP # ==================== app = FastAPI( title="FinBot RAG API", description="Advanced RAG system with RBAC, hierarchical chunking, and guardrails", version="1.0.0", lifespan=lifespan, ) # Add CORS middleware app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # ==================== # CHAT ENDPOINT # ==================== @app.post("/api/chat", response_model=ChatResponse) async def chat(request: ChatRequest): """ Process a user query through the RAG pipeline. Args: request: ChatRequest with user_role, query, and optional user_id Returns: ChatResponse with answer, sources, and metadata """ try: # Validate user role valid_roles = ["employee", "finance", "engineering", "marketing", "c_level"] if request.user_role not in valid_roles: raise HTTPException( status_code=400, detail=f"Invalid user role. Must be one of: {valid_roles}" ) # Get RAG pipeline pipeline = get_rag_pipeline() # Process query rag_response = pipeline.answer_query( user_role=request.user_role, query_text=request.query, user_id=request.user_id, ) print(rag_response) # Convert to response model return ChatResponse( answer=rag_response.answer, sources=rag_response.sources, route=rag_response.route, user_role=rag_response.user_role, accessible_collections=rag_response.accessible_collections, guardrail_flags=rag_response.guardrail_flags, guardrail_warnings=rag_response.guardrail_warnings, rbac_denied=rag_response.rbac_denied, rbac_reason=rag_response.rbac_reason, ) except HTTPException: raise except Exception as e: logger.error(f"Error processing chat request: {str(e)}", exc_info=True) # Return a proper ChatResponse with error info instead of a 500, # so the frontend always has something to display. return ChatResponse( answer="I'm sorry, I encountered an unexpected error while processing your question. Please try again in a moment.", sources=[], route="error", user_role=request.user_role, accessible_collections=[], guardrail_flags=["server_error"], guardrail_warnings=[f"Internal error: {str(e)}"], ) # ==================== # USER MANAGEMENT ENDPOINTS # ==================== @app.get("/api/users", response_model=list[UserInfo]) async def list_users(): """Get list of demo users for login screen.""" try: user_manager = get_user_manager() users = user_manager.list_users() return [ UserInfo( username=u.username, name=u.name, role=u.role.value, # Use .value to get "finance" not "UserRole.FINANCE" department=u.department, accessible_collections=user_manager.get_user_accessible_collections( u.role.value ), ) for u in users ] except Exception as e: logger.error(f"Error listing users: {str(e)}") raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") @app.get("/api/users/{username}") async def get_user(username: str): """Get specific user information.""" try: user_manager = get_user_manager() user = user_manager.get_user(username) if not user: raise HTTPException(status_code=404, detail=f"User not found: {username}") return { "username": user.username, "name": user.name, "role": user.role, "department": user.department, "accessible_collections": user_manager.get_user_accessible_collections(user.role), } except HTTPException: raise except Exception as e: logger.error(f"Error getting user: {str(e)}") raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") # ==================== # COLLECTIONS ENDPOINTS # ==================== @app.get("/api/collections", response_model=list[CollectionInfo]) async def list_collections(): """Get list of document collections.""" try: from config import COLLECTION_CONFIGS collections = [] for coll_enum in DocumentCollection: config = COLLECTION_CONFIGS.get(coll_enum) if config: collections.append( CollectionInfo( name=coll_enum.value, description=config.get("description", ""), accessible_roles=config.get("access_roles", []), ) ) return collections except Exception as e: logger.error(f"Error listing collections: {str(e)}") raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") @app.get("/api/collections/{collection_name}") async def get_collection_info(collection_name: str): """Get information about a specific collection.""" try: from config import COLLECTION_CONFIGS # Find collection coll = None for c in DocumentCollection: if c.value == collection_name: coll = c break if not coll: raise HTTPException(status_code=404, detail=f"Collection not found: {collection_name}") config = COLLECTION_CONFIGS.get(coll) # Get vector store stats vector_store = get_vector_store() stats = vector_store.get_collection_stats(collection_name) return { "name": collection_name, "description": config.get("description", ""), "accessible_roles": config.get("access_roles", []), "chunks_count": stats.get("points_count", 0) if stats else 0, } except HTTPException: raise except Exception as e: logger.error(f"Error getting collection info: {str(e)}") raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") # ==================== # INGESTION ENDPOINT (Admin) # ==================== @app.post("/api/admin/ingest") async def ingest_documents(): """ Ingest all document collections. WARNING: Only use for demo/testing! """ try: logger.info("Starting document ingestion...") ingester = DocumentIngester() results = ingester.ingest_all_collections() stats = ingester.verify_ingestion() return { "status": "success", "ingestion_results": results, "collection_stats": stats, } except Exception as e: logger.error(f"Error ingesting documents: {str(e)}") raise HTTPException(status_code=500, detail=f"Ingestion failed: {str(e)}") # ==================== # SYSTEM ENDPOINTS # ==================== @app.get("/api/health") async def health_check(): """Health check endpoint.""" try: vector_store = get_vector_store() collections = vector_store.list_collections() return { "status": "healthy", "collections_available": len(collections) > 0, "collections": collections, } except Exception as e: logger.error(f"Health check failed: {str(e)}") return JSONResponse( status_code=503, content={ "status": "unhealthy", "error": str(e), }, ) @app.get("/api/info") async def system_info(): """Get system information.""" try: return { "name": "FinBot RAG System", "version": "1.0.0", "features": [ "Role-Based Access Control (RBAC)", "Hierarchical Document Chunking", "Semantic Query Routing", "Input/Output Guardrails", "RAGAs Evaluation Support", ], "available_roles": ["employee", "finance", "engineering", "marketing", "c_level"], } except Exception as e: logger.error(f"Error getting system info: {str(e)}") raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") # ==================== # ERROR HANDLERS # ==================== @app.exception_handler(HTTPException) async def http_exception_handler(request: Request, exc: HTTPException): """Handle HTTP exceptions.""" return JSONResponse( status_code=exc.status_code, content={ "error": exc.detail, "status_code": exc.status_code, }, ) @app.exception_handler(Exception) async def general_exception_handler(request: Request, exc: Exception): """Handle general exceptions.""" logger.error(f"Unhandled exception: {str(exc)}") return JSONResponse( status_code=500, content={ "error": "Internal server error", "detail": str(exc), }, ) # ==================== # ROOT ENDPOINT # ==================== @app.get("/") async def root(): """Root endpoint with API documentation.""" return { "name": "FinBot RAG API", "version": "1.0.0", "description": "Advanced RAG system with RBAC, hierarchical chunking, and guardrails", "endpoints": { "chat": "POST /api/chat - Process a user query", "users": "GET /api/users - List demo users", "collections": "GET /api/collections - List document collections", "health": "GET /api/health - Health check", "info": "GET /api/info - System information", "ingest": "POST /api/admin/ingest - Ingest documents (admin only)", }, "documentation": "/docs", } if __name__ == "__main__": import uvicorn logger.info("Starting FinBot RAG API server...") uvicorn.run( app, host="0.0.0.0", port=8000, log_level="info", )