Update flask_Character.py

flask_Character.py

@@ -11,10 +11,9 @@ from datetime import date, datetime, timedelta
 from typing import List, Optional, Literal, Dict, Any, Tuple
 import traceback
 import asyncio
-from uuid import uuid4, UUID # For unique request IDs
+from uuid import uuid4, UUID 
 
 from fastapi import FastAPI, HTTPException, Response, Query, Depends, status
-from fastapi.responses import FileResponse
 from fastapi.exception_handlers import http_exception_handler
 from starlette.exceptions import HTTPException as StarletteHTTPException
 
@@ -22,7 +21,7 @@ from langchain.prompts import PromptTemplate
 from langchain_groq import ChatGroq
 from pydantic import BaseModel, Field, BeforeValidator, model_serializer
 from typing_extensions import Annotated
-from pydantic_core import core_schema # Import core_schema for direct use in __get_pydantic_json_schema__
+from pydantic_core import core_schema 
 
 from pymongo import MongoClient
 from pymongo.errors import ConnectionFailure, OperationFailure
@@ -33,11 +32,9 @@ MAX_BATCH_SIZE = 20
 BATCH_INTERVAL_SECONDS = 1.0
 
 # --- Queues and pending request stores for batching ---
-extract_data_queue: Optional[asyncio.Queue] = None
-generate_reply_queue: Optional[asyncio.Queue] = None
+extract_data_queue: Optional[asyncio.Queue[Tuple[Any, UUID]]] = None
+generate_reply_queue: Optional[asyncio.Queue[Tuple[Any, UUID]]] = None
 
-# Dictionaries to store futures for pending requests, keyed by unique request ID
-# The value will be an asyncio.Future that the endpoint handler will await
 extract_pending_requests: Dict[UUID, asyncio.Future] = {}
 generate_pending_requests: Dict[UUID, asyncio.Future] = {}
 
@@ -46,7 +43,7 @@ shutdown_event = asyncio.Event()
 
 
 # --- MongoDB Configuration ---
-MONGO_URI = "mongodb+srv://precison9:P1LhtFknkT75yg5L@cluster0.isuwpef.mongodb.net" # Use os.getenv in prod
+MONGO_URI = "mongodb+srv://precison9:P1LhtFknkT75yg5L@cluster0.isuwpef.mongodb.net" 
 DB_NAME = "email_assistant_db"
 EXTRACTED_EMAILS_COLLECTION = "extracted_emails"
 GENERATED_REPLIES_COLLECTION = "generated_replies"
@@ -167,24 +164,23 @@ class GenerateReplyResponse(BaseModel):
     stored_id: str = Field(..., description="The MongoDB ID of the stored reply.")
     cached: bool = Field(..., description="True if the reply was retrieved from cache, False if newly generated.")
 
-# --- Query Models for GET Endpoints ---
 class ExtractedEmailQuery(BaseModel):
-    contact_name: Optional[str] = Query(None)
-    appointment_title: Optional[str] = Query(None)
-    task_title: Optional[str] = Query(None)
-    from_date: Optional[date] = Query(None)
-    to_date: Optional[date] = Query(None)
+    contact_name: Optional[str] = Query(None, description="Filter by contact name.")
+    appointment_title: Optional[str] = Query(None, description="Filter by appointment title.")
+    task_title: Optional[str] = Query(None, description="Filter by task title.")
+    from_date: Optional[date] = Query(None, description="Filter by processed date (start).")
+    to_date: Optional[date] = Query(None, description="Filter by processed date (end).")
     limit: int = Query(10, ge=1, le=100)
 
 class GeneratedReplyQuery(BaseModel):
-    language: Optional[Literal["Italian", "English"]] = Query(None)
-    style: Optional[str] = Query(None)
-    tone: Optional[str] = Query(None)
-    from_date: Optional[date] = Query(None)
-    to_date: Optional[date] = Query(None)
+    language: Optional[Literal["Italian", "English"]] = Query(None, description="Filter by language.")
+    style: Optional[str] = Query(None, description="Filter by style.")
+    tone: Optional[str] = Query(None, description="Filter by tone.")
+    from_date: Optional[date] = Query(None, description="Filter by generation date (start).")
+    to_date: Optional[date] = Query(None, description="Filter by generation date (end).")
     limit: int = Query(10, ge=1, le=100)
 
-# ---------------------- Utility Functions ----------------------
+# --- Utility Functions ---
 def extract_last_json_block(text: str) -> Optional[str]:
     pattern = r'```json\s*(.*?)\s*```'
     matches = re.findall(pattern, text, re.DOTALL)
@@ -212,17 +208,17 @@ def normalize_llm_output(data: dict, current_date: date, original_email_text: st
     tasks_data = [Task(task_title=t.get("task_title", "Untitled"), task_description=t.get("task_description", "No description"), due_date=parse_date(t.get("due_date"), current_date) or current_date) for t in data.get("tasks", [])]
     return ExtractedData(contacts=contacts_data, appointments=appointments_data, tasks=tasks_data, original_email_text=original_email_text)
 
-# ---------------------- Core Logic (Internal Functions) ----------------------
+# --- Core Logic (Internal Functions) ---
 def _process_email_internal(email_text: str, api_key: str, current_date: date) -> ExtractedData:
     if not email_text: raise ValueError("Email text cannot be empty for processing.")
     llm = ChatGroq(model="meta-llama/llama-4-scout-17b-16e-instruct", temperature=0, max_tokens=2000, groq_api_key=api_key)
     prompt_today_str = current_date.isoformat()
     prompt_tomorrow_str = (current_date + timedelta(days=1)).isoformat()
     prompt_template_str = f"""You are an expert email assistant tasked with extracting structured information from an Italian email.**Your response MUST be a single, complete JSON object, wrapped in a ```json``` block.****DO NOT include any conversational text, explanations, or preambles outside the JSON block.****The JSON should contain three top-level keys: "contacts", "appointments", and "tasks".**If a category has no items, its list should be empty (e.g., "contacts": []).Here is the required JSON schema for each category:- **contacts**: List of Contact objects.    Each Contact object must have:    - `name` (string, full name)    - `last_name` (string, last name) - You should infer this from the full name.    - `email` (string, optional, null if not present)    - `phone_number` (string, optional, null if not present)- **appointments**: List of Appointment objects.    Each Appointment object must have:    - `title` (string, short, meaningful title in Italian based on the meeting's purpose)    - `description` (string, summary of the meeting's goal)    - `start_date` (string, YYYY-MM-DD. If not explicitly mentioned, use "{prompt_today_str}" for "today", or "{prompt_tomorrow_str}" for "tomorrow")    - `start_time` (string, optional, e.g., "10:30 AM", null if not present)    - `end_date` (string, YYYY-MM-DD, optional, null if unknown or not applicable)    - `end_time` (string, optional, e.g., "11:00 AM", null if not present)- **tasks**: List of Task objects.    Each Task object must have:    - `task_title` (string, short summary of action item)    - `task_description` (string, more detailed explanation)    - `due_date` (string, YYYY-MM-DD. Infer from context, e.g., "entro domani" becomes "{prompt_tomorrow_str}", "today" becomes "{prompt_today_str}")---Email:{{email}}"""
-    prompt_template = PromptTemplate(input_variables=["email"], template=prompt_template_str) # Removed unused prompt_today_str, prompt_tomorrow_str from input_variables
+    prompt_template = PromptTemplate(input_variables=["email"], template=prompt_template_str)
     chain = prompt_template | llm
     try:
-        llm_output = chain.invoke({"email": email_text}) # Removed unused variables
+        llm_output = chain.invoke({"email": email_text})
         llm_output_str = llm_output.content
         json_str = extract_last_json_block(llm_output_str)
         if not json_str: raise ValueError(f"No JSON block found in LLM output. LLM response: {llm_output_str}")
@@ -243,7 +239,10 @@ def _generate_response_internal(email_text: str, api_key: str, language: Literal
     except Exception as e: traceback.print_exc(); raise
 
 # --- FastAPI Application ---
-app = FastAPI(title="Email Assistant API", description="API for extracting structured data and generating replies.", version="1.2.0", docs_url="/", redoc_url="/redoc")
+# NOTE: To make this batching system work correctly with asyncio.Queue,
+# you must run the server with a SINGLE WORKER. For example:
+# uvicorn main:app --workers 1
+app = FastAPI(title="Email Assistant API", description="API with batch processing.", version="1.3.0", docs_url="/", redoc_url="/redoc")
 
 # --- Exception Handlers ---
 @app.exception_handler(StarletteHTTPException)
@@ -255,130 +254,74 @@ async def global_exception_handler_wrapper(request, exc):
     traceback.print_exc()
     return Response(content=json.dumps({"detail": f"Internal Server Error: {str(exc)}", "type": "unhandled_exception"}), status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, media_type="application/json")
 
-# --- Batch Worker Functions ---
+# --- Batch Worker Helper Functions ---
 async def _execute_single_extract_task(request_item: ProcessEmailRequest, request_id: UUID) -> Tuple[UUID, Any]:
-    """Helper to run a single extract task and return its result or exception."""
     current_date = date.today()
     try:
-        # Run blocking LLM call in a thread pool
-        result = await asyncio.to_thread(
-            _process_email_internal, request_item.email_text, request_item.groq_api_key, current_date
-        )
+        result = await asyncio.to_thread(_process_email_internal, request_item.email_text, request_item.groq_api_key, current_date)
         return request_id, result
     except Exception as e:
         return request_id, e
 
-async def extract_data_batch_worker():
-    global extract_data_queue, extract_pending_requests, shutdown_event
-    print(f"[{datetime.now()}] Extract Data Batch Worker started.")
-    while not shutdown_event.is_set():
-        try:
-            await asyncio.wait_for(asyncio.sleep(BATCH_INTERVAL_SECONDS), timeout=BATCH_INTERVAL_SECONDS + 0.1) # Ensures it runs roughly every interval
-        except asyncio.TimeoutError: # woken up by shutdown
-            if shutdown_event.is_set(): break
-
-        batch_to_process: List[Tuple[ProcessEmailRequest, UUID]] = []
-        while len(batch_to_process) < MAX_BATCH_SIZE:
-            try:
-                request_obj, req_id = extract_data_queue.get_nowait()
-                batch_to_process.append((request_obj, req_id))
-            except asyncio.QueueEmpty:
-                break # No more items for this batch
-
-        if not batch_to_process:
-            continue
-
-        print(f"[{datetime.now()}] Extract Worker: Processing batch of {len(batch_to_process)} requests.")
-        
-        # Concurrently execute all LLM calls for the current batch
-        llm_tasks = [_execute_single_extract_task(req_obj, req_id) for req_obj, req_id in batch_to_process]
-        results = await asyncio.gather(*llm_tasks) # Results are (request_id, result_or_exception)
-
-        for req_id, result_or_exc in results:
-            future = extract_pending_requests.pop(req_id, None) # Get and remove future
-            if future and not future.done():
-                if isinstance(result_or_exc, Exception):
-                    future.set_exception(result_or_exc)
-                else:
-                    future.set_result(result_or_exc) # This is ExtractedData object (pre-DB)
-            elif future and future.done():
-                 print(f"[{datetime.now()}] Extract Worker: Future for {req_id} was already done (e.g. timed out).")
-
-    print(f"[{datetime.now()}] Extract Data Batch Worker shutting down.")
-    # Clear out any remaining requests in the queue by setting exception
-    while not extract_data_queue.empty():
-        try:
-            _, req_id = extract_data_queue.get_nowait()
-            future = extract_pending_requests.pop(req_id, None)
-            if future and not future.done():
-                future.set_exception(HTTPException(status_code=503, detail="Service shutting down, request cancelled."))
-        except asyncio.QueueEmpty:
-            break
-
-
 async def _execute_single_generate_reply_task(request_item: GenerateReplyRequest, request_id: UUID) -> Tuple[UUID, Any]:
-    """Helper to run a single generate reply task and return its result or exception."""
     try:
-        # Run blocking LLM call in a thread pool
-        result = await asyncio.to_thread(
-            _generate_response_internal,
-            request_item.email_text, request_item.groq_api_key, request_item.language,
-            request_item.length, request_item.style, request_item.tone, request_item.emoji
-        )
-        return request_id, result # This is the reply string
+        result = await asyncio.to_thread(_generate_response_internal, request_item.email_text, request_item.groq_api_key, request_item.language, request_item.length, request_item.style, request_item.tone, request_item.emoji)
+        return request_id, result
     except Exception as e:
         return request_id, e
 
-async def generate_reply_batch_worker():
-    global generate_reply_queue, generate_pending_requests, shutdown_event
-    print(f"[{datetime.now()}] Generate Reply Batch Worker started.")
+# --- Batch Worker Functions ---
+async def batch_worker(
+    queue: asyncio.Queue,
+    pending_requests: Dict[UUID, asyncio.Future],
+    execution_function: callable,
+    worker_name: str
+):
+    print(f"[{datetime.now()}] {worker_name} started.")
     while not shutdown_event.is_set():
         try:
-            await asyncio.wait_for(asyncio.sleep(BATCH_INTERVAL_SECONDS), timeout=BATCH_INTERVAL_SECONDS + 0.1)
-        except asyncio.TimeoutError:
+            await asyncio.wait_for(shutdown_event.wait(), timeout=BATCH_INTERVAL_SECONDS)
             if shutdown_event.is_set(): break
+        except asyncio.TimeoutError:
+            pass  # Interval elapsed, time to process
 
-        batch_to_process: List[Tuple[GenerateReplyRequest, UUID]] = []
-        while len(batch_to_process) < MAX_BATCH_SIZE:
+        batch_to_process = []
+        while len(batch_to_process) < MAX_BATCH_SIZE and not queue.empty():
             try:
-                request_obj, req_id = generate_reply_queue.get_nowait()
-                batch_to_process.append((request_obj, req_id))
+                batch_to_process.append(queue.get_nowait())
             except asyncio.QueueEmpty:
                 break
-
+        
         if not batch_to_process:
             continue
 
-        print(f"[{datetime.now()}] Reply Worker: Processing batch of {len(batch_to_process)} requests.")
+        print(f"[{datetime.now()}] {worker_name}: Processing batch of {len(batch_to_process)} requests.")
         
-        llm_tasks = [_execute_single_generate_reply_task(req_obj, req_id) for req_obj, req_id in batch_to_process]
-        results = await asyncio.gather(*llm_tasks)
+        tasks = [execution_function(req_obj, req_id) for req_obj, req_id in batch_to_process]
+        results = await asyncio.gather(*tasks)
 
         for req_id, result_or_exc in results:
-            future = generate_pending_requests.pop(req_id, None)
+            future = pending_requests.pop(req_id, None)
             if future and not future.done():
                 if isinstance(result_or_exc, Exception):
                     future.set_exception(result_or_exc)
                 else:
-                    future.set_result(result_or_exc) # This is reply string
-            elif future and future.done():
-                 print(f"[{datetime.now()}] Reply Worker: Future for {req_id} was already done (e.g. timed out).")
+                    future.set_result(result_or_exc)
 
-    print(f"[{datetime.now()}] Generate Reply Batch Worker shutting down.")
-    while not generate_reply_queue.empty():
-        try:
-            _, req_id = generate_reply_queue.get_nowait()
-            future = generate_pending_requests.pop(req_id, None)
-            if future and not future.done():
-                future.set_exception(HTTPException(status_code=503, detail="Service shutting down, request cancelled."))
-        except asyncio.QueueEmpty:
-            break
+    print(f"[{datetime.now()}] {worker_name} shutting down. Cancelling pending requests...")
+    for req_id, future in list(pending_requests.items()):
+        if not future.done():
+            future.set_exception(HTTPException(status_code=503, detail="Service shutting down"))
+        pending_requests.pop(req_id, None)
+    print(f"[{datetime.now()}] {worker_name} stopped.")
 
 # --- FastAPI Event Handlers ---
+worker_tasks: List[asyncio.Task] = []
+
 @app.on_event("startup")
 async def startup_event():
     global client, db, extracted_emails_collection, generated_replies_collection
-    global extract_data_queue, generate_reply_queue # Initialize queues
+    global extract_data_queue, generate_reply_queue, worker_tasks
 
     print(f"[{datetime.now()}] FastAPI app startup sequence initiated.")
     try:
@@ -389,43 +332,42 @@ async def startup_event():
         generated_replies_collection = db[GENERATED_REPLIES_COLLECTION]
         print(f"[{datetime.now()}] Successfully connected to MongoDB: {DB_NAME}")
 
-        # Initialize queues and start worker tasks
         extract_data_queue = asyncio.Queue()
         generate_reply_queue = asyncio.Queue()
-        asyncio.create_task(extract_data_batch_worker())
-        asyncio.create_task(generate_reply_batch_worker())
+        
+        task1 = asyncio.create_task(batch_worker(extract_data_queue, extract_pending_requests, _execute_single_extract_task, "Extract Worker"))
+        task2 = asyncio.create_task(batch_worker(generate_reply_queue, generate_pending_requests, _execute_single_generate_reply_task, "Reply Worker"))
+        worker_tasks.extend([task1, task2])
         print(f"[{datetime.now()}] Batch processing workers started.")
 
     except (ConnectionFailure, OperationFailure) as e:
         print(f"[{datetime.now()}] ERROR: MongoDB Connection/Operation Failure: {e}")
-        # Critical error, prevent app from fully starting or indicate non-operational state
-        # For simplicity, we'll let it run but endpoints relying on DB will fail
         client = db = extracted_emails_collection = generated_replies_collection = None
     except Exception as e:
         print(f"[{datetime.now()}] ERROR: An unexpected error during startup: {e}")
         traceback.print_exc()
         client = db = extracted_emails_collection = generated_replies_collection = None
     finally:
-         # Simplified check after connection attempt
-        if not (client and db and extracted_emails_collection and generated_replies_collection):
-            print(f"[{datetime.now()}] MongoDB or dependent services (batch queues) might not be fully initialized.")
+        # CORRECTED: Use 'is not None' for pymongo objects
+        if not (client is not None and db is not None and extracted_emails_collection is not None and generated_replies_collection is not None):
+            print(f"[{datetime.now()}] WARNING: MongoDB might not be fully initialized.")
         print(f"[{datetime.now()}] FastAPI app startup sequence completed.")
 
 
 @app.on_event("shutdown")
-async def shutdown_event_handler(): # Renamed to avoid conflict with global shutdown_event
-    global client, shutdown_event # Use the global shutdown_event
+async def shutdown_event_handler():
+    global client, shutdown_event, worker_tasks
     print(f"[{datetime.now()}] FastAPI app shutting down.")
     
-    # Signal workers to stop
-    shutdown_event.set()
+    if not shutdown_event.is_set():
+        shutdown_event.set()
     
-    # Give workers a moment to process their current items or exit
-    # This timeout should be slightly longer than BATCH_INTERVAL_SECONDS to allow a final batch cycle
-    # Or join the worker tasks if they are stored globally (more robust)
-    await asyncio.sleep(BATCH_INTERVAL_SECONDS + 0.5) 
+    if worker_tasks:
+        print(f"[{datetime.now()}] Waiting for worker tasks to complete...")
+        await asyncio.gather(*worker_tasks, return_exceptions=True)
+        print(f"[{datetime.now()}] Worker tasks completed.")
 
-    if client:
+    if client is not None:
         client.close()
         print(f"[{datetime.now()}] MongoDB client closed.")
     print(f"[{datetime.now()}] FastAPI app shutdown sequence completed.")
@@ -435,52 +377,42 @@ async def shutdown_event_handler(): # Renamed to avoid conflict with global shut
 async def health_check():
     db_status = "MongoDB not connected."
     db_ok = False
-    if client and db:
+    # CORRECTED: Use 'is not None' for pymongo objects
+    if client is not None and db is not None:
         try:
-            await asyncio.to_thread(db.list_collection_names)
+            await asyncio.to_thread(client.admin.command, 'ping')
             db_status = "MongoDB connection OK."
             db_ok = True
         except Exception as e:
             db_status = f"MongoDB connection error: {e}"
     
-    queue_status = {
-        "extract_data_queue_size": extract_data_queue.qsize() if extract_data_queue else "N/A",
-        "generate_reply_queue_size": generate_reply_queue.qsize() if generate_reply_queue else "N/A"
-    }
+    queues = {}
+    if extract_data_queue is not None and generate_reply_queue is not None:
+        queues = {
+            "extract_data_queue_size": extract_data_queue.qsize(),
+            "generate_reply_queue_size": generate_reply_queue.qsize(),
+            "extract_pending_requests": len(extract_pending_requests),
+            "generate_pending_requests": len(generate_pending_requests)
+        }
     
     if db_ok:
-        return {"status": "ok", "message": "Email Assistant API is up.", "database": db_status, "queues": queue_status}
+        return {"status": "ok", "message": "API is up.", "database": db_status, "queues": queues}
     else:
-        raise HTTPException(status_code=503, detail={"message": "Service unavailable.", "database": db_status, "queues": queue_status})
+        raise HTTPException(status_code=503, detail={"message": "Service unavailable.", "database": db_status})
 
 @app.post("/extract-data", response_model=ExtractedData, summary="Extract structured data (batched)")
 async def extract_email_data(request: ProcessEmailRequest):
-    global extract_data_queue, extract_pending_requests, extracted_emails_collection
-
-    if not extracted_emails_collection or not extract_data_queue:
+    # CORRECTED: Use 'is None' check
+    if extracted_emails_collection is None or extract_data_queue is None:
         raise HTTPException(status_code=503, detail="Service not available (DB or batch queue).")
 
     request_id = uuid4()
-    future = asyncio.get_event_loop().create_future()
+    future: asyncio.Future[ExtractedData] = asyncio.get_event_loop().create_future()
     extract_pending_requests[request_id] = future
 
     try:
         await extract_data_queue.put((request, request_id))
-        print(f"[{datetime.now()}] /extract-data: Queued request {request_id}. Queue size: {extract_data_queue.qsize()}")
-
-        # Wait for the future to be resolved by the worker, with a timeout
-        try:
-            # Timeout should be configurable, longer than batch interval + processing time
-            extracted_data_obj = await asyncio.wait_for(future, timeout=60.0)
-        except asyncio.TimeoutError:
-            print(f"[{datetime.now()}] /extract-data: Request {request_id} timed out waiting for worker.")
-            # The future might still be in extract_pending_requests if worker hasn't processed it
-            # Worker will try to pop it; if already popped here, it's fine.
-            extract_pending_requests.pop(request_id, None) # Clean up if timed out
-            raise HTTPException(status_code=504, detail="Request timed out while awaiting processing in batch.")
-        
-        # If here, extracted_data_obj is the ExtractedData model instance from the worker
-        print(f"[{datetime.now()}] /extract-data: Worker processed {request_id}. Inserting to DB.")
+        extracted_data_obj = await asyncio.wait_for(future, timeout=60.0)
         
         data_to_insert = extracted_data_obj.model_dump(by_alias=True, exclude_none=True, exclude={'id'})
         if 'appointments' in data_to_insert:
@@ -492,55 +424,35 @@ async def extract_email_data(request: ProcessEmailRequest):
                 if isinstance(task_item.get('due_date'), date): task_item['due_date'] = datetime.combine(task_item['due_date'], datetime.min.time())
         
         insert_result = await asyncio.to_thread(extracted_emails_collection.insert_one, data_to_insert)
-        extracted_data_obj.id = str(insert_result.inserted_id) # Update with DB ID
+        extracted_data_obj.id = str(insert_result.inserted_id)
         
         return extracted_data_obj
-
-    except HTTPException: # Re-raise HTTPExceptions (like timeout or from future)
-        raise
-    except ValueError as ve: # Typically from Pydantic validation or LLM output parsing
-        raise HTTPException(status_code=400, detail=str(ve))
-    except Exception as e: # Any other exception from the future or this handler
-        traceback.print_exc()
-        raise HTTPException(status_code=500, detail=f"Internal server error during extract data: {e}")
+    except asyncio.TimeoutError:
+        raise HTTPException(status_code=504, detail="Request timed out while awaiting processing in batch.")
+    except Exception as e:
+        # Re-raise exceptions set by the worker (like HTTPException or others)
+        raise e
     finally:
-        # Ensure future is removed if it wasn't already (e.g. successful completion)
         extract_pending_requests.pop(request_id, None)
 
-
 @app.post("/generate-reply", response_model=GenerateReplyResponse, summary="Generate smart reply (batched)")
 async def generate_email_reply(request: GenerateReplyRequest):
-    global generate_reply_queue, generate_pending_requests, generated_replies_collection
-
-    if not generated_replies_collection or not generate_reply_queue:
+    # CORRECTED: Use 'is None' check
+    if generated_replies_collection is None or generate_reply_queue is None:
         raise HTTPException(status_code=503, detail="Service not available (DB or batch queue).")
 
-    # --- Check cache first (remains outside batching) ---
     cache_query = {"original_email_text": request.email_text, "language": request.language, "length": request.length, "style": request.style, "tone": request.tone, "emoji": request.emoji}
     cached_reply_doc = await asyncio.to_thread(generated_replies_collection.find_one, cache_query)
     if cached_reply_doc:
-        print(f"[{datetime.now()}] /generate-reply: Reply found in cache. ID: {str(cached_reply_doc['_id'])}")
         return GenerateReplyResponse(reply=cached_reply_doc["generated_reply_text"], stored_id=str(cached_reply_doc["_id"]), cached=True)
 
-    # --- If not cached, queue for generation ---
     request_id = uuid4()
-    future = asyncio.get_event_loop().create_future()
+    future: asyncio.Future[str] = asyncio.get_event_loop().create_future()
     generate_pending_requests[request_id] = future
 
     try:
         await generate_reply_queue.put((request, request_id))
-        print(f"[{datetime.now()}] /generate-reply: Queued request {request_id}. Queue size: {generate_reply_queue.qsize()}")
-
-        try:
-            # Timeout should be configurable
-            reply_content_str = await asyncio.wait_for(future, timeout=60.0) 
-        except asyncio.TimeoutError:
-            print(f"[{datetime.now()}] /generate-reply: Request {request_id} timed out waiting for worker.")
-            generate_pending_requests.pop(request_id, None)
-            raise HTTPException(status_code=504, detail="Request timed out while awaiting reply generation in batch.")
-
-        # If here, reply_content_str is the generated string from the worker
-        print(f"[{datetime.now()}] /generate-reply: Worker generated reply for {request_id}. Storing to DB.")
+        reply_content_str = await asyncio.wait_for(future, timeout=60.0)
 
         reply_data_to_store = GeneratedReplyData(
             original_email_text=request.email_text, generated_reply_text=reply_content_str,
@@ -553,16 +465,13 @@ async def generate_email_reply(request: GenerateReplyRequest):
         stored_id = str(insert_result.inserted_id)
         
         return GenerateReplyResponse(reply=reply_content_str, stored_id=stored_id, cached=False)
-
-    except HTTPException:
-        raise
+    except asyncio.TimeoutError:
+        raise HTTPException(status_code=504, detail="Request timed out while awaiting reply generation in batch.")
     except Exception as e:
-        traceback.print_exc()
-        raise HTTPException(status_code=500, detail=f"Internal server error during generate reply: {e}")
+        raise e
     finally:
         generate_pending_requests.pop(request_id, None)
 
-
 @app.get("/query-extracted-emails", response_model=List[ExtractedData], summary="Query stored extracted email data")
 async def query_extracted_emails(query_params: ExtractedEmailQuery = Depends()):
     if extracted_emails_collection is None: raise HTTPException(status_code=503, detail="MongoDB not available.")