Update app/app.py

app/app.py

@@ -4,7 +4,9 @@ import asyncio
 import logging
 import uuid
 import re
-from fastapi import FastAPI, HTTPException, Request
+from typing import Dict, List, Optional
+from datetime import datetime
+from fastapi import FastAPI, HTTPException, Request, BackgroundTasks
 from pydantic import BaseModel
 from llama_cpp import Llama
 # Correctly reference the module within the 'app' package
@@ -21,6 +23,131 @@ class RequestIdAdapter(logging.LoggerAdapter):
 
 logger = logging.getLogger("app")
 
+# -----------------------------
+# ✅ Queue Management Classes
+# -----------------------------
+class QueuedRequest:
+    def __init__(self, request_id: str, question: str, timestamp: datetime):
+        self.request_id = request_id
+        self.question = question
+        self.timestamp = timestamp
+        self.status = "queued"  # queued, processing, completed, failed, timeout
+        self.result: Optional[Dict] = None
+        self.error: Optional[str] = None
+
+class RequestQueue:
+    def __init__(self, max_size: int = 15):
+        self.queue: List[QueuedRequest] = []
+        self.processing: Optional[QueuedRequest] = None
+        self.completed_requests: Dict[str, QueuedRequest] = {}
+        self.max_size = max_size
+        self.lock = asyncio.Lock()
+    
+    async def add_request(self, request_id: str, question: str) -> Dict:
+        async with self.lock:
+            if len(self.queue) >= self.max_size:
+                return {
+                    "status": "queue_full",
+                    "message": f"Queue is full (max {self.max_size} requests). Please try again later.",
+                    "queue_position": None,
+                    "estimated_wait_time": None
+                }
+            
+            queued_request = QueuedRequest(request_id, question, datetime.now())
+            
+            # If nothing is processing and queue is empty, this request can be processed immediately
+            if not self.processing and not self.queue:
+                self.processing = queued_request
+                queued_request.status = "processing"
+                return {
+                    "status": "processing",
+                    "message": "Your request is being processed now.",
+                    "queue_position": 0,
+                    "estimated_wait_time": "0-2 minutes"
+                }
+            
+            # Add to queue
+            self.queue.append(queued_request)
+            position = len(self.queue)
+            estimated_wait = f"{position * 2}-{(position + 1) * 2} minutes"
+            
+            return {
+                "status": "queued",
+                "message": f"Using free CPU tier - can only process one request at a time. Your request is #{position} in queue and will be processed after current requests are completed.",
+                "queue_position": position,
+                "estimated_wait_time": estimated_wait
+            }
+    
+    async def get_next_request(self) -> Optional[QueuedRequest]:
+        async with self.lock:
+            if self.queue:
+                next_request = self.queue.pop(0)
+                self.processing = next_request
+                next_request.status = "processing"
+                return next_request
+            return None
+    
+    async def complete_request(self, request_id: str, result: Dict = None, error: str = None):
+        async with self.lock:
+            if self.processing and self.processing.request_id == request_id:
+                if result:
+                    self.processing.result = result
+                    self.processing.status = "completed"
+                elif error:
+                    self.processing.error = error
+                    self.processing.status = "failed"
+                
+                # Store completed request for result retrieval
+                self.completed_requests[request_id] = self.processing
+                self.processing = None
+                
+                # Clean up old completed requests (keep only last 50)
+                if len(self.completed_requests) > 50:
+                    oldest_keys = list(self.completed_requests.keys())[:-50]
+                    for key in oldest_keys:
+                        del self.completed_requests[key]
+    
+    async def get_request_status(self, request_id: str) -> Optional[Dict]:
+        async with self.lock:
+            # Check if currently processing
+            if self.processing and self.processing.request_id == request_id:
+                return {
+                    "status": self.processing.status,
+                    "message": "Your request is currently being processed.",
+                    "result": self.processing.result
+                }
+            
+            # Check completed requests
+            if request_id in self.completed_requests:
+                req = self.completed_requests[request_id]
+                return {
+                    "status": req.status,
+                    "message": "Request completed." if req.status == "completed" else "Request failed.",
+                    "result": req.result,
+                    "error": req.error
+                }
+            
+            # Check queue
+            for i, req in enumerate(self.queue):
+                if req.request_id == request_id:
+                    return {
+                        "status": "queued",
+                        "message": f"Your request is #{i+1} in queue.",
+                        "queue_position": i + 1,
+                        "estimated_wait_time": f"{(i+1) * 2}-{(i+2) * 2} minutes"
+                    }
+            
+            return None
+    
+    async def get_queue_info(self) -> Dict:
+        async with self.lock:
+            return {
+                "queue_length": len(self.queue),
+                "currently_processing": self.processing.request_id if self.processing else None,
+                "max_queue_size": self.max_size,
+                "completed_requests_count": len(self.completed_requests)
+            }
+
 # -----------------------------
 # ✅ Configuration
 # -----------------------------
@@ -31,12 +158,16 @@ LLM_TIMEOUT_SECONDS = int(os.getenv("LLM_TIMEOUT_SECONDS", "120"))
 RELEVANCE_THRESHOLD = float(os.getenv("RELEVANCE_THRESHOLD", "0.3"))
 TOP_K_SEARCH = int(os.getenv("TOP_K_SEARCH", "4"))
 TOP_K_CONTEXT = int(os.getenv("TOP_K_CONTEXT", "2"))
+MAX_QUEUE_SIZE = int(os.getenv("MAX_QUEUE_SIZE", "15"))
 
 # -----------------------------
 # ✅ Initialize FastAPI App
 # -----------------------------
 app = FastAPI(title="NEEPCO DoP RAG Chatbot", version="2.1.0")
 
+# Initialize request queue
+request_queue = RequestQueue(max_size=MAX_QUEUE_SIZE)
+
 @app.middleware("http")
 async def add_request_id(request: Request, call_next):
     request_id = str(uuid.uuid4())
@@ -101,25 +232,40 @@ class Feedback(BaseModel):
     comment: str | None = None
 
 # -----------------------------
-# ✅ Endpoints
+# ✅ Background Processing
 # -----------------------------
-def get_logger_adapter(request: Request):
-    return RequestIdAdapter(logger, {'request_id': getattr(request.state, 'request_id', 'N/A')})
+async def process_queued_requests():
+    """Background task to process queued requests"""
+    while True:
+        try:
+            next_request = await request_queue.get_next_request()
+            if next_request:
+                logger.info(f"Processing queued request: {next_request.request_id}")
+                try:
+                    result = await process_chat_request(next_request.question, next_request.request_id)
+                    await request_queue.complete_request(next_request.request_id, result=result)
+                    logger.info(f"Completed request: {next_request.request_id}")
+                except Exception as e:
+                    error_msg = f"Error processing request: {str(e)}"
+                    logger.error(f"Failed to process request {next_request.request_id}: {e}", exc_info=True)
+                    await request_queue.complete_request(next_request.request_id, error=error_msg)
+            else:
+                # No requests in queue, wait a bit before checking again
+                await asyncio.sleep(2)
+        except Exception as e:
+            logger.error(f"Error in background processor: {e}", exc_info=True)
+            await asyncio.sleep(5)
 
-@app.get("/")
-async def root():
-    return {"status": "✅ Server is running."}
+# Start background processor
+@app.on_event("startup")
+async def startup_event():
+    asyncio.create_task(process_queued_requests())
 
-@app.get("/health")
-async def health_check():
-    status = {
-        "status": "ok",
-        "database_status": "ready" if db_ready else "error",
-        "model_status": "ready" if model_ready else "error"
-    }
-    if not db_ready or not model_ready:
-        raise HTTPException(status_code=503, detail=status)
-    return status
+# -----------------------------
+# ✅ Core Processing Function
+# -----------------------------
+def get_logger_adapter(request_id: str):
+    return RequestIdAdapter(logger, {'request_id': request_id})
 
 async def generate_llm_response(prompt: str, request_id: str):
     loop = asyncio.get_running_loop()
@@ -132,23 +278,23 @@ async def generate_llm_response(prompt: str, request_id: str):
         raise ValueError("Empty response from LLM")
     return answer
 
-@app.post("/chat")
-async def chat(query: Query, request: Request):
-    adapter = get_logger_adapter(request)
-    question_lower = query.question.strip().lower()
+async def process_chat_request(question: str, request_id: str) -> Dict:
+    """Core chat processing logic extracted for reuse"""
+    adapter = get_logger_adapter(request_id)
+    question_lower = question.strip().lower()
 
     # --- GREETING & INTRO HANDLING ---
     greeting_keywords = ["hello", "hi", "hey", "what can you do", "who are you"]
     if question_lower in greeting_keywords:
-        adapter.info(f"Handling a greeting or introductory query: '{query.question}'")
+        adapter.info(f"Handling a greeting or introductory query: '{question}'")
         intro_message = (
             "Hello! I am an AI assistant specifically trained on NEEPCO's Delegation of Powers (DoP) policy document. "
             "My purpose is to help you find accurate information and answer questions based on this specific dataset. "
             "I am currently running on a CPU-based environment. How can I assist you with the DoP policy today?"
         )
         return {
-            "request_id": getattr(request.state, 'request_id', 'N/A'),
-            "question": query.question,
+            "request_id": request_id,
+            "question": question,
             "context_used": "NA - Greeting",
             "answer": intro_message
         }
@@ -157,15 +303,16 @@ async def chat(query: Query, request: Request):
         adapter.error("Service unavailable due to initialization failure.")
         raise HTTPException(status_code=503, detail="Service is not ready. Please check logs.")
 
-    adapter.info(f"Received query: '{query.question}'")
+    adapter.info(f"Received query: '{question}'")
 
     # 1. Search Vector DB
-    search_results = db.search(query.question, top_k=TOP_K_SEARCH)
+    search_results = db.search(question, top_k=TOP_K_SEARCH)
 
     if not search_results:
         adapter.warning("No relevant context found in vector DB.")
         return {
-            "question": query.question,
+            "request_id": request_id,
+            "question": question,
             "context_used": "No relevant context found.",
             "answer": "Sorry, I could not find a relevant policy to answer that question. Please try rephrasing."
         }
@@ -181,7 +328,6 @@ async def chat(query: Query, request: Request):
     prompt = f"""<|system|>
 You are a precise and factual assistant for NEEPCO's Delegation of Powers (DoP) policy.
 Your task is to answer the user's question based ONLY on the provided context.
-
 - **Formatting Rule:** If the answer contains a list of items or steps, you **MUST** separate each item with a pipe symbol (`|`). For example: `First item|Second item|Third item`.
 - **Content Rule:** If the information is not in the provided context, you **MUST** reply with the exact phrase: "The provided policy context does not contain information on this topic."
 </s>
@@ -190,9 +336,8 @@ Your task is to answer the user's question based ONLY on the provided context.
 ```
 {context}
 ```
-
 ### Question:
-{query.question}
+{question}
 </s>
 <|assistant|>
 ### Detailed Answer:
@@ -203,7 +348,7 @@ Your task is to answer the user's question based ONLY on the provided context.
     try:
         adapter.info("Sending prompt to LLM for generation...")
         raw_answer = await asyncio.wait_for(
-            generate_llm_response(prompt, request.state.request_id),
+            generate_llm_response(prompt, request_id),
             timeout=LLM_TIMEOUT_SECONDS
         )
         adapter.info(f"LLM generation successful. Raw response: {raw_answer[:250]}...")
@@ -229,14 +374,71 @@ Your task is to answer the user's question based ONLY on the provided context.
         adapter.error(f"An unexpected error occurred during LLM generation: {e}", exc_info=True)
         answer = "Sorry, an unexpected error occurred while generating a response."
 
-    adapter.info(f"Final answer prepared. Returning to client.")
+    adapter.info(f"Final answer prepared. Returning result.")
     return {
-        "request_id": request.state.request_id,
-        "question": query.question,
+        "request_id": request_id,
+        "question": question,
         "context_used": context,
         "answer": answer
     }
 
+# -----------------------------
+# ✅ Endpoints
+# -----------------------------
+@app.get("/")
+async def root():
+    return {"status": "✅ Server is running."}
+
+@app.get("/health")
+async def health_check():
+    queue_info = await request_queue.get_queue_info()
+    status = {
+        "status": "ok",
+        "database_status": "ready" if db_ready else "error",
+        "model_status": "ready" if model_ready else "error",
+        "queue_info": queue_info
+    }
+    if not db_ready or not model_ready:
+        raise HTTPException(status_code=503, detail=status)
+    return status
+
+@app.post("/chat")
+async def chat(query: Query, request: Request):
+    """Add request to queue and return queue status"""
+    if not db_ready or not model_ready:
+        raise HTTPException(status_code=503, detail="Service is not ready. Please check logs.")
+    
+    request_id = request.state.request_id
+    adapter = get_logger_adapter(request_id)
+    
+    adapter.info(f"Received chat request: '{query.question}'")
+    
+    # Add request to queue
+    queue_status = await request_queue.add_request(request_id, query.question)
+    
+    return {
+        "request_id": request_id,
+        "question": query.question,
+        **queue_status
+    }
+
+@app.get("/status/{request_id}")
+async def get_request_status(request_id: str):
+    """Check the status of a specific request"""
+    status = await request_queue.get_request_status(request_id)
+    if not status:
+        raise HTTPException(status_code=404, detail="Request not found")
+    
+    return {
+        "request_id": request_id,
+        **status
+    }
+
+@app.get("/queue")
+async def get_queue_status():
+    """Get current queue information"""
+    return await request_queue.get_queue_info()
+
 @app.post("/feedback")
 async def collect_feedback(feedback: Feedback, request: Request):
     adapter = get_logger_adapter(request)
@@ -250,4 +452,4 @@ async def collect_feedback(feedback: Feedback, request: Request):
         "comment": feedback.comment
     }
     adapter.info(json.dumps(feedback_log))
-    return {"status": "✅ Feedback recorded. Thank you!"}
+    return {"status": "✅ Feedback recorded. Thank you!"}