Update app/app.py

app/app.py

@@ -1,3 +1,4 @@
+# app.py - Updated with better request handling
 import os
 import json
 import asyncio
@@ -5,7 +6,7 @@ import logging
 import uuid
 import re
 from typing import Dict, List, Optional
-from datetime import datetime
+from datetime import datetime, timedelta
 from fastapi import FastAPI, HTTPException, Request, BackgroundTasks
 from pydantic import BaseModel
 from llama_cpp import Llama
@@ -31,9 +32,11 @@ class QueuedRequest:
         self.request_id = request_id
         self.question = question
         self.timestamp = timestamp
-        self.status = "queued"  # queued, processing, completed, failed, timeout
+        self.status = "queued"  # queued, processing, completed, failed, timeout, cancelled
         self.result: Optional[Dict] = None
         self.error: Optional[str] = None
+        self.cancelled = False  # Track if request was cancelled
+        self.last_accessed = datetime.now()  # Track when status was last checked
 
 class RequestQueue:
     def __init__(self, max_size: int = 15):
@@ -42,9 +45,14 @@ class RequestQueue:
         self.completed_requests: Dict[str, QueuedRequest] = {}
         self.max_size = max_size
         self.lock = asyncio.Lock()
+        self.cleanup_interval = 300  # 5 minutes
+        self.max_completed_age = 600  # 10 minutes
     
     async def add_request(self, request_id: str, question: str) -> Dict:
         async with self.lock:
+            # Clean up old requests periodically
+            await self._cleanup_old_requests()
+            
             if len(self.queue) >= self.max_size:
                 return {
                     "status": "queue_full",
@@ -90,27 +98,57 @@ class RequestQueue:
     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:
+                if self.processing.cancelled:
+                    # Don't store results for cancelled requests
+                    self.processing.status = "cancelled"
+                    logger.info(f"Request {request_id} was cancelled, not storing result")
+                elif 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
+                # Store completed request for result retrieval (even cancelled ones briefly)
                 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 cancel_request(self, request_id: str) -> bool:
+        """Cancel a request if it exists in queue or is processing"""
+        async with self.lock:
+            # Check if it's currently processing
+            if self.processing and self.processing.request_id == request_id:
+                self.processing.cancelled = True
+                logger.info(f"Marked processing request {request_id} as cancelled")
+                return True
+            
+            # Check if it's in queue
+            for i, req in enumerate(self.queue):
+                if req.request_id == request_id:
+                    cancelled_req = self.queue.pop(i)
+                    cancelled_req.status = "cancelled"
+                    cancelled_req.cancelled = True
+                    self.completed_requests[request_id] = cancelled_req
+                    logger.info(f"Cancelled queued request {request_id}")
+                    return True
+            
+            return False
     
     async def get_request_status(self, request_id: str) -> Optional[Dict]:
         async with self.lock:
+            # Update last accessed time for any request we're checking
+            current_time = datetime.now()
+            
             # Check if currently processing
             if self.processing and self.processing.request_id == request_id:
+                self.processing.last_accessed = current_time
+                if self.processing.cancelled:
+                    return {
+                        "status": "cancelled",
+                        "message": "Request was cancelled.",
+                        "result": None,
+                        "error": "Request cancelled by user"
+                    }
                 return {
                     "status": self.processing.status,
                     "message": "Your request is currently being processed.",
@@ -120,9 +158,18 @@ class RequestQueue:
             # Check completed requests
             if request_id in self.completed_requests:
                 req = self.completed_requests[request_id]
+                req.last_accessed = current_time
+                
+                status_messages = {
+                    "completed": "Request completed.",
+                    "failed": "Request failed.",
+                    "cancelled": "Request was cancelled.",
+                    "timeout": "Request timed out."
+                }
+                
                 return {
                     "status": req.status,
-                    "message": "Request completed." if req.status == "completed" else "Request failed.",
+                    "message": status_messages.get(req.status, "Request processed."),
                     "result": req.result,
                     "error": req.error
                 }
@@ -130,6 +177,7 @@ class RequestQueue:
             # Check queue
             for i, req in enumerate(self.queue):
                 if req.request_id == request_id:
+                    req.last_accessed = current_time
                     return {
                         "status": "queued",
                         "message": f"Your request is #{i+1} in queue.",
@@ -139,6 +187,21 @@ class RequestQueue:
             
             return None
     
+    async def _cleanup_old_requests(self):
+        """Clean up old completed requests and abandoned requests"""
+        current_time = datetime.now()
+        cutoff_time = current_time - timedelta(seconds=self.max_completed_age)
+        
+        # Clean up old completed requests
+        to_remove = []
+        for request_id, req in self.completed_requests.items():
+            if req.last_accessed < cutoff_time:
+                to_remove.append(request_id)
+        
+        for request_id in to_remove:
+            del self.completed_requests[request_id]
+            logger.info(f"Cleaned up old request: {request_id}")
+    
     async def get_queue_info(self) -> Dict:
         async with self.lock:
             return {
@@ -242,9 +305,28 @@ async def process_queued_requests():
             if next_request:
                 logger.info(f"Processing queued request: {next_request.request_id}")
                 try:
+                    # Check if request was cancelled before processing
+                    if next_request.cancelled:
+                        logger.info(f"Skipping cancelled request: {next_request.request_id}")
+                        await request_queue.complete_request(
+                            next_request.request_id, 
+                            error="Request was cancelled"
+                        )
+                        continue
+                    
                     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}")
+                    
+                    # Check again if request was cancelled during processing
+                    if next_request.cancelled:
+                        logger.info(f"Request was cancelled during processing: {next_request.request_id}")
+                        await request_queue.complete_request(
+                            next_request.request_id, 
+                            error="Request was cancelled during processing"
+                        )
+                    else:
+                        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)
@@ -425,14 +507,37 @@ async def chat(query: Query, request: Request):
 @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
-    }
+    try:
+        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
+        }
+    except Exception as e:
+        logger.error(f"Error checking status for {request_id}: {e}")
+        raise HTTPException(status_code=500, detail="Error checking request status")
+
+@app.delete("/cancel/{request_id}")
+async def cancel_request(request_id: str):
+    """Cancel a specific request"""
+    try:
+        cancelled = await request_queue.cancel_request(request_id)
+        if not cancelled:
+            raise HTTPException(status_code=404, detail="Request not found or cannot be cancelled")
+        
+        return {
+            "status": "cancelled",
+            "message": f"Request {request_id} has been cancelled",
+            "request_id": request_id
+        }
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Error cancelling request {request_id}: {e}")
+        raise HTTPException(status_code=500, detail="Error cancelling request")
 
 @app.get("/queue")
 async def get_queue_status():
@@ -441,7 +546,7 @@ async def get_queue_status():
 
 @app.post("/feedback")
 async def collect_feedback(feedback: Feedback, request: Request):
-    adapter = get_logger_adapter(request)
+    adapter = get_logger_adapter(request.state.request_id)
     feedback_log = {
         "type": "USER_FEEDBACK",
         "request_id": feedback.request_id,