from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel from typing import Optional, List, Dict import os from datetime import datetime import logging import threading import requests from chatbot import RAGChatbot logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) app = FastAPI( title="RAG Chatbot API - Multi-User", description="HR Assistant Chatbot with Per-User Session Management", version="2.0.0", docs_url="/docs", redoc_url="/redoc" ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Global base chatbot instance base_chatbot = None # Per-user session storage user_sessions = {} session_lock = threading.Lock() # Configuration MAX_SESSIONS = 100 SESSION_TIMEOUT = 3600 # 1 hour class UserSession: """Isolated session for each user""" def __init__(self, user_id: str): self.user_id = user_id self.chat_history = [] self.conversation_context = { 'current_employee': None, 'last_mentioned_entities': [] } self.last_activity = datetime.now() def update_activity(self): self.last_activity = datetime.now() def cleanup_old_sessions(): """Remove inactive sessions""" with session_lock: current_time = datetime.now() to_remove = [] for user_id, session in user_sessions.items(): time_diff = (current_time - session.last_activity).total_seconds() if time_diff > SESSION_TIMEOUT: to_remove.append(user_id) for user_id in to_remove: del user_sessions[user_id] logger.info(f"Cleaned up session for user: {user_id}") def get_or_create_session(user_id: str) -> UserSession: """Get existing session or create new one""" with session_lock: if len(user_sessions) > MAX_SESSIONS: cleanup_old_sessions() if user_id not in user_sessions: user_sessions[user_id] = UserSession(user_id) logger.info(f"Created new session for user: {user_id}") session = user_sessions[user_id] session.update_activity() return session # Pydantic models class ChatRequest(BaseModel): question: str user_id: str class ChatResponse(BaseModel): question: str answer: str timestamp: str user_id: str session_info: Dict @app.on_event("startup") async def startup_event(): global base_chatbot logger.info("=== Starting RAG Chatbot Initialization ===") try: PDF_PATH = os.getenv("PDF_PATH", "./data/policies.pdf") HF_TOKEN = os.getenv("HF_TOKEN") if not HF_TOKEN: raise ValueError("HF_TOKEN environment variable not set") logger.info(f"PDF Path: {PDF_PATH}") logger.info(f"File exists: {os.path.exists(PDF_PATH)}") if not os.path.exists(PDF_PATH): raise ValueError(f"PDF file not found at {PDF_PATH}") base_chatbot = RAGChatbot(PDF_PATH, HF_TOKEN) logger.info("=== Base chatbot initialized successfully! ===") except Exception as e: logger.error(f"Failed to initialize chatbot: {e}") raise @app.get("/") async def root(): return { "service": "RAG Chatbot API", "version": "2.0.0", "status": "healthy", "active_sessions": len(user_sessions), "chatbot_loaded": base_chatbot is not None, "endpoints": { "docs": "/docs", "chat": "POST /api/chat", "history": "GET /api/history/{user_id}", "reset": "POST /api/reset?user_id=xxx", "sessions": "GET /api/sessions" } } @app.get("/api/health") async def health_check(): if base_chatbot is None: raise HTTPException(status_code=503, detail="Chatbot not initialized") return { "status": "healthy", "timestamp": datetime.now().isoformat(), "chatbot_ready": True, "active_sessions": len(user_sessions) } @app.post("/api/chat", response_model=ChatResponse) async def chat(request: ChatRequest): """Send a question to the chatbot with user session isolation""" if base_chatbot is None: raise HTTPException(status_code=503, detail="Chatbot not initialized") if not request.question.strip(): raise HTTPException(status_code=400, detail="Question cannot be empty") if not request.user_id: raise HTTPException(status_code=400, detail="user_id is required") try: logger.info(f"User {request.user_id}: {request.question[:50]}...") # Get user session session = get_or_create_session(request.user_id) # Resolve pronouns resolved_question = base_chatbot._resolve_pronouns_for_session( request.question, session.conversation_context ) # Retrieve relevant chunks retrieved_data = base_chatbot._retrieve(resolved_question, k=20) # Search user's chat history relevant_past_chats = base_chatbot._search_session_history( resolved_question, session.chat_history, k=5 ) # Build prompt prompt = base_chatbot._build_prompt_for_session( resolved_question, retrieved_data, relevant_past_chats, session.chat_history, session.conversation_context ) # ✅ NEW: Call Hugging Face Router API payload = { "model": base_chatbot.model_name, "messages": [ { "role": "user", "content": prompt } ], "max_tokens": 512, "temperature": 0.3 } response = requests.post( base_chatbot.api_url, headers=base_chatbot.headers, json=payload, timeout=60 ) response.raise_for_status() result = response.json() # Extract answer answer = result["choices"][0]["message"]["content"] # Update user's conversation context base_chatbot._update_conversation_context_for_session( request.question, answer, session.conversation_context ) # Store in user's history chat_entry = { 'timestamp': datetime.now().isoformat(), 'question': request.question, 'answer': answer, 'used_past_context': len(relevant_past_chats) > 0 } session.chat_history.append(chat_entry) response_data = ChatResponse( question=request.question, answer=answer, timestamp=datetime.now().isoformat(), user_id=request.user_id, session_info={ 'total_messages': len(session.chat_history), 'current_context': session.conversation_context.get('current_employee') } ) logger.info(f"User {request.user_id}: Question processed successfully") return response_data except Exception as e: logger.error(f"Error for user {request.user_id}: {e}") raise HTTPException(status_code=500, detail=f"Error: {str(e)}") @app.post("/api/reset") async def reset_chat(user_id: str): """Reset chat history for specific user""" if not user_id: raise HTTPException(status_code=400, detail="user_id is required") with session_lock: if user_id in user_sessions: del user_sessions[user_id] logger.info(f"Reset session for user: {user_id}") return {"message": f"Chat history reset for user {user_id}", "status": "success"} else: return {"message": f"No session found for user {user_id}", "status": "success"} @app.get("/api/history/{user_id}") async def get_history(user_id: str): """Get chat history for specific user""" session = get_or_create_session(user_id) return { "user_id": user_id, "total_conversations": len(session.chat_history), "current_context": session.conversation_context.get('current_employee'), "history": session.chat_history } @app.get("/api/sessions") async def get_active_sessions(): """Get list of active sessions""" with session_lock: return { "total_sessions": len(user_sessions), "max_sessions": MAX_SESSIONS, "session_timeout_seconds": SESSION_TIMEOUT, "sessions": [ { "user_id": user_id, "messages": len(session.chat_history), "last_activity": session.last_activity.isoformat(), "current_context": session.conversation_context.get('current_employee') } for user_id, session in user_sessions.items() ] } @app.post("/api/cleanup") async def manual_cleanup(): """Manually trigger session cleanup""" cleanup_old_sessions() return { "message": "Cleanup completed", "active_sessions": len(user_sessions) } if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=7860)