Update flask_Character.py

flask_Character.py

@@ -1,4 +1,3 @@
-
 # This software is licensed under a **dual-license model**
 # For individuals and businesses earning **under $1M per year**, this software is licensed under the **MIT License**
 # Businesses or organizations with **annual revenue of $1,000,000 or more** must obtain permission to use this software commercially.
@@ -20,44 +19,75 @@ from starlette.exceptions import HTTPException as StarletteHTTPException
 from langchain.prompts import PromptTemplate
 from langchain_groq import ChatGroq
 from pydantic import BaseModel, Field, BeforeValidator, model_serializer
+# Ensure you have pydantic >= 2.0.0 for core_schema and typing_extensions for Annotated
 from typing_extensions import Annotated
+from pydantic_core import core_schema # Import core_schema for direct use in __get_pydantic_json_schema__
 
 from pymongo import MongoClient
 from pymongo.errors import ConnectionFailure, OperationFailure
 from bson import ObjectId
 
 # --- MongoDB Configuration ---
-MONGO_URI = "mongodb+srv://precison9:P1LhtFknkT75yg5L@cluster0.isuwpef.mongodb.net" # Replace with your actual URI
+# IMPORTANT: Replace with your actual URI.
+# For security, consider using environment variables for your MONGO_URI in production.
+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"
 
+# Global variables for MongoDB client and collections
 client: Optional[MongoClient] = None
-db: Optional[Any] = None # Changed to Optional[Any] as Database type is not directly imported for annotation here
+db: Optional[Any] = None
 extracted_emails_collection: Optional[Any] = None
 generated_replies_collection: Optional[Any] = None
 
 # --- Pydantic ObjectId Handling ---
 class CustomObjectId(str):
+    """
+    Custom Pydantic type for handling MongoDB ObjectIds.
+    It validates that the input is a valid ObjectId string and
+    ensures it's represented as a string in JSON Schema.
+    """
     @classmethod
     def __get_validators__(cls):
         yield cls.validate
 
     @classmethod
-    def validate(cls, v, info):
+    def validate(cls, v):
+        # Allow None or empty string to pass through for optional fields if not handled by Optional[PyObjectId]
+        if v is None or v == "":
+            return None # Or raise ValueError if not allowed
+
+        # Ensure input is a string or convertible to string for ObjectId.is_valid
+        if not isinstance(v, (str, ObjectId)):
+            raise ValueError("ObjectId must be a string or ObjectId instance")
+
+        # Convert ObjectId to string if it's already an ObjectId instance
+        if isinstance(v, ObjectId):
+            return str(v)
+
+        # Validate string format
         if not ObjectId.is_valid(v):
-            raise ValueError("Invalid ObjectId")
-        return str(v)
+            raise ValueError("Invalid ObjectId format")
+        return cls(v) # Return an instance of CustomObjectId (which is a str subclass)
 
+
+    # This method is crucial for Pydantic v2 to generate correct OpenAPI schema
     @classmethod
-    def __get_pydantic_json_schema__(cls, core_schema, handler):
-        json_schema = handler(core_schema)
-        json_schema["type"] = "string"
-        json_schema["example"] = "60c728ef238b9c7b9e0f6c2a"
+    def __get_pydantic_json_schema__(
+        cls, _core_schema: core_schema.CoreSchema, handler
+    ) -> Dict[str, Any]:
+        # We tell Pydantic that this custom type should be represented as a standard string
+        # in the generated JSON Schema (OpenAPI documentation).
+        # We use handler to process a simple string schema.
+        json_schema = handler(core_schema.str_schema())
+        json_schema["example"] = "60c728ef238b9c7b9e0f6c2a" # Add an example for clarity
         return json_schema
 
+# Annotated type for convenience in models
 PyObjectId = Annotated[CustomObjectId, BeforeValidator(str)]
 
+
 # ---------------------- Models ----------------------
 class Contact(BaseModel):
     name: str
@@ -79,6 +109,7 @@ class Task(BaseModel):
     due_date: date
 
 class ExtractedData(BaseModel):
+    # Use PyObjectId for the _id field
     id: Optional[PyObjectId] = Field(alias="_id", default=None)
     contacts: List[Contact]
     appointments: List[Appointment]
@@ -87,14 +118,17 @@ class ExtractedData(BaseModel):
     processed_at: datetime = Field(default_factory=datetime.utcnow)
 
     class Config:
-        populate_by_name = True
-        arbitrary_types_allowed = True
+        populate_by_name = True # Allow setting 'id' or '_id'
+        arbitrary_types_allowed = True # Allow CustomObjectId and ObjectId
 
+    # Custom serializer for JSON output to ensure ObjectId is converted to string
     @model_serializer(when_used='json')
     def serialize_model(self):
         data = self.model_dump(by_alias=True, exclude_none=True)
+        # Ensure _id is a string when serializing to JSON
         if "_id" in data and isinstance(data["_id"], ObjectId):
             data["_id"] = str(data["_id"])
+        # Ensure dates are correctly serialized to ISO format if they are date objects
         if 'appointments' in data:
             for appt in data['appointments']:
                 if isinstance(appt.get('start_date'), date):
@@ -121,6 +155,7 @@ class GenerateReplyRequest(BaseModel):
     emoji: str = Field("Auto", examples=["Auto", "None", "Occasional", "Frequent"])
 
 class GeneratedReplyData(BaseModel):
+    # Use PyObjectId for the _id field
     id: Optional[PyObjectId] = Field(alias="_id", default=None)
     original_email_text: str
     generated_reply_text: str
@@ -161,30 +196,42 @@ class GeneratedReplyQuery(BaseModel):
 
 # ---------------------- Utility Functions ----------------------
 def extract_last_json_block(text: str) -> Optional[str]:
+    """
+    Extracts the last JSON block enclosed in ```json``` from a string,
+    or a standalone JSON object if no code block is found.
+    """
     pattern = r'```json\s*(.*?)\s*```'
     matches = re.findall(pattern, text, re.DOTALL)
     if matches:
         return matches[-1].strip()
+    # Fallback: try to find a standalone JSON object
     match = re.search(r'\{.*\}', text, re.DOTALL)
     if match:
         return match.group(0)
     return None
 
-def parse_date(date_str: Optional[str], current_date: date) -> Optional[date]: # Allow None to pass through
-    if not date_str: return None # If input is None (e.g. optional end_date), return None
+def parse_date(date_str: Optional[str], current_date: date) -> Optional[date]:
+    """
+    Parses a date string, handling 'today', 'tomorrow', and YYYY-MM-DD format.
+    Returns None if input is None or cannot be parsed into a date.
+    """
+    if not date_str:
+        return None
     date_str_lower = date_str.lower().strip()
-    if date_str_lower == "today": return current_date
-    if date_str_lower == "tomorrow": return current_date + timedelta(days=1)
+    if date_str_lower == "today":
+        return current_date
+    if date_str_lower == "tomorrow":
+        return current_date + timedelta(days=1)
     try:
         return datetime.strptime(date_str_lower, "%Y-%m-%d").date()
     except ValueError:
-        # For "appointments" "start_date", you assumed "today" if parsing failed or not present.
-        # For "end_date" it was optional. We need to be consistent.
-        # Given the original normalize_llm_output, start_date defaulted to today, end_date was optional.
-        # This parse_date is more general. The default handling should be in normalize_llm_output.
-        return current_date # Fallback, or raise error, or return None depending on strictness
+        return None # Return None if parsing fails, let normalize_llm_output handle defaults
 
 def normalize_llm_output(data: dict, current_date: date, original_email_text: str) -> ExtractedData:
+    """
+    Normalizes and validates LLM extracted data into ExtractedData Pydantic model.
+    Handles defaults for dates and name splitting.
+    """
     def split_name(full_name: str) -> tuple[str, str]:
         parts = full_name.strip().split()
         name = parts[0] if parts else ""
@@ -198,8 +245,10 @@ def normalize_llm_output(data: dict, current_date: date, original_email_text: st
 
     appointments_data = []
     for a in data.get("appointments", []):
-        start_date_val = parse_date(a.get("start_date", "today"), current_date) or current_date # Default to current_date if parse_date returns None
-        end_date_val = parse_date(a.get("end_date"), current_date) # parse_date can return None if end_date is not provided or invalid
+        # Default start_date to current_date if not provided or invalid
+        start_date_val = parse_date(a.get("start_date"), current_date) or current_date
+        # end_date remains optional
+        end_date_val = parse_date(a.get("end_date"), current_date)
 
         appointments_data.append(Appointment(
             title=a.get("title", "Untitled"), description=a.get("description", "No description"),
@@ -209,7 +258,8 @@ def normalize_llm_output(data: dict, current_date: date, original_email_text: st
 
     tasks_data = []
     for t in data.get("tasks", []):
-        due_date_val = parse_date(t.get("due_date", "today"), current_date) or current_date # Default to current_date
+        # Default due_date to current_date if not provided or invalid
+        due_date_val = parse_date(t.get("due_date"), current_date) or current_date
         tasks_data.append(Task(
             task_title=t.get("task_title", "Untitled"), task_description=t.get("task_description", "No description"),
             due_date=due_date_val
@@ -218,11 +268,17 @@ def normalize_llm_output(data: dict, current_date: date, original_email_text: st
 
 # ---------------------- 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.")
+    """
+    Internal function to process email text using LLM and extract structured data.
+    """
+    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()
-    # Ensure your full, detailed prompt is used here
+
     prompt_template_str = f"""
 You are an expert email assistant tasked with extracting structured information from an Italian email.
 
@@ -263,24 +319,18 @@ Email:
     prompt_template = PromptTemplate(input_variables=["email", "prompt_today_str", "prompt_tomorrow_str"], template=prompt_template_str)
     chain = prompt_template | llm
     try:
-        # print(f"DEBUG: Invoking LLM with email_text length: {len(email_text)} and current_date: {current_date}")
         llm_output = chain.invoke({"email": email_text, "prompt_today_str": prompt_today_str, "prompt_tomorrow_str": prompt_tomorrow_str})
         llm_output_str = llm_output.content
-        # print(f"DEBUG: Raw LLM output:\n{llm_output_str[:500]}...")
 
         json_str = extract_last_json_block(llm_output_str)
-        # print(f"DEBUG: Extracted JSON string:\n{json_str}")
 
-        if not json_str: raise ValueError(f"No JSON block found in LLM output. LLM response: {llm_output_str}")
+        if not json_str:
+            raise ValueError(f"No JSON block found in LLM output. LLM response: {llm_output_str}")
         json_data = json.loads(json_str)
-        # print(f"DEBUG: Parsed JSON data: {json.dumps(json_data, indent=2)}")
 
         extracted_data = normalize_llm_output(json_data, current_date, email_text)
-        # print("DEBUG: Data normalized successfully.")
         return extracted_data
     except json.JSONDecodeError as e:
-        # print(f"ERROR: JSON Decode Error: {e}")
-        # print(f"ERROR: LLM response that caused error:\n{llm_output_str}")
         raise ValueError(f"Failed to parse JSON from LLM output: {e}\nLLM response was:\n{llm_output_str}")
     except Exception as e:
         traceback.print_exc()
@@ -290,9 +340,12 @@ def _generate_response_internal(
     email_text: str, api_key: str, language: Literal["Italian", "English"],
     length: str, style: str, tone: str, emoji: str
 ) -> str:
-    if not email_text: return "Cannot generate reply for empty email text."
+    """
+    Internal function to generate a reply to an email using LLM.
+    """
+    if not email_text:
+        return "Cannot generate reply for empty email text."
     llm = ChatGroq(model="meta-llama/llama-4-scout-17b-16e-instruct", temperature=0.7, max_tokens=800, groq_api_key=api_key)
-    # Ensure your full, detailed prompt is used here
     prompt_template_str="""
 You are an assistant that helps reply to emails.
 
@@ -343,6 +396,7 @@ async def handle_single_reply_request(request_data: GenerateReplyRequest, future
             "tone": request_data.tone,
             "emoji": request_data.emoji,
         }
+        # Use await asyncio.to_thread for blocking MongoDB operations
         cached_reply_doc = await asyncio.to_thread(generated_replies_collection.find_one, cache_query)
 
         if cached_reply_doc:
@@ -351,7 +405,8 @@ async def handle_single_reply_request(request_data: GenerateReplyRequest, future
                 "stored_id": str(cached_reply_doc["_id"]),
                 "cached": True
             }
-            if not future.done(): future.set_result(response)
+            if not future.done():
+                future.set_result(response)
             return
 
         reply_content = await asyncio.to_thread(
@@ -374,6 +429,7 @@ async def handle_single_reply_request(request_data: GenerateReplyRequest, future
             tone=request_data.tone,
             emoji=request_data.emoji
         )
+        # Use model_dump for Pydantic v2
         reply_data_dict = reply_data_to_store.model_dump(by_alias=True, exclude_none=True, exclude={'id'})
 
         insert_result = await asyncio.to_thread(generated_replies_collection.insert_one, reply_data_dict)
@@ -384,7 +440,8 @@ async def handle_single_reply_request(request_data: GenerateReplyRequest, future
             "stored_id": stored_id,
             "cached": False
         }
-        if not future.done(): future.set_result(final_response)
+        if not future.done():
+            future.set_result(final_response)
 
     except Exception as e:
         traceback.print_exc()
@@ -399,6 +456,7 @@ async def process_reply_batches():
         async with reply_queue_condition:
             if not reply_request_queue:
                 await reply_queue_condition.wait()
+                # After waking up, re-check if queue is still empty (e.g., if notified but then emptied by another worker)
                 if not reply_request_queue:
                     continue
 
@@ -422,28 +480,31 @@ async def process_reply_batches():
             tasks = [handle_single_reply_request(req_data, fut) for req_data, fut in batch_to_fire]
             await asyncio.gather(*tasks)
         else:
-             await asyncio.sleep(0.001)
+            # Short sleep to prevent busy-waiting if queue is empty but not waiting
+            await asyncio.sleep(0.001)
 
 
 # ---------------------- FastAPI Application ----------------------
 app = FastAPI(
     title="Email Assistant API",
     description="API for extracting structured data from emails and generating intelligent replies using Groq LLMs, with MongoDB integration, dynamic date handling, batching, and caching.",
-    version="1.1.0", # Incremented version
-    docs_url="/",
+    version="1.1.0",
+    docs_url="/", # Sets Swagger UI to be the root path
     redoc_url="/redoc"
 )
 
 # --- Global Exception Handler ---
+# Catch Starlette HTTPExceptions (FastAPI uses these internally)
 @app.exception_handler(StarletteHTTPException)
 async def custom_http_exception_handler_wrapper(request, exc):
     return await http_exception_handler(request, exc)
 
+# Catch all other unhandled exceptions
 @app.exception_handler(Exception)
 async def global_exception_handler_wrapper(request, exc):
     print(f"Unhandled exception caught by global handler for request: {request.url}")
-    traceback.print_exc()
-    # Ensure it returns a valid FastAPI response
+    traceback.print_exc() # Print traceback to console for debugging
+    # Return a JSON response for consistency, even for unhandled errors
     return Response(content=json.dumps({"detail": f"Internal Server Error: {str(exc)}"}), status_code=500, media_type="application/json")
 
 
@@ -452,20 +513,23 @@ async def global_exception_handler_wrapper(request, exc):
 async def startup_event():
     global client, db, extracted_emails_collection, generated_replies_collection, batch_processor_task
     try:
+        # Connect to MongoDB
         client = MongoClient(MONGO_URI, serverSelectionTimeoutMS=5000)
-        client.admin.command('ping')
+        client.admin.command('ping') # Test connection
         db = client[DB_NAME]
         extracted_emails_collection = db[EXTRACTED_EMAILS_COLLECTION]
         generated_replies_collection = db[GENERATED_REPLIES_COLLECTION]
         print(f"Successfully connected to MongoDB: {DB_NAME}")
 
+        # Start the batch processor task if not already running
         if batch_processor_task is None:
-            loop = asyncio.get_event_loop()
-            batch_processor_task = loop.create_task(process_reply_batches())
+            # Use asyncio.create_task for proper task management in an async app
+            batch_processor_task = asyncio.create_task(process_reply_batches())
             print("Batch processor task for replies started.")
 
     except (ConnectionFailure, OperationFailure) as e:
         print(f"ERROR: MongoDB Connection/Operation Failure: {e}")
+        # Ensure all DB related globals are reset to None if connection fails
         client = None
         db = None
         extracted_emails_collection = None
@@ -478,87 +542,117 @@ async def startup_event():
         extracted_emails_collection = None
         generated_replies_collection = None
     finally:
-        # Corrected condition for checking client and db
+        # Final check and logging for MongoDB connection status
         if client is not None and db is not None:
             try:
+                # One last ping to confirm connection before app fully starts
                 client.admin.command('ping')
-            except Exception:
-                print("MongoDB ping failed after initial connection attempt during finally block.")
+            except Exception as e:
+                print(f"MongoDB ping failed after initial connection attempt during finally block: {e}")
                 client = None; db = None; extracted_emails_collection = None; generated_replies_collection = None
         else:
-            print("MongoDB client or db object is None after connection attempt in startup.")
-            if client is None or db is None: # Ensure all are None if one is
+            print("MongoDB client or db object is None after connection attempt in startup. Database likely not connected.")
+            # Ensure all are None if one is, to avoid partial state
+            if client is None or db is None:
                 client = None; db = None; extracted_emails_collection = None; generated_replies_collection = None
-        print("FastAPI app starting up. MongoDB client & Batch Processor initialization attempted.")
+        print("FastAPI app startup sequence completed for MongoDB client & Batch Processor initialization.")
 
 
 @app.on_event("shutdown")
 async def shutdown_event():
     global client, batch_processor_task
+    # Cancel the batch processor task
     if batch_processor_task:
         batch_processor_task.cancel()
         try:
+            # Await the task to ensure it has a chance to clean up/handle cancellation
             await batch_processor_task
         except asyncio.CancelledError:
-            print("Batch processor task for replies cancelled.")
+            print("Batch processor task for replies cancelled during shutdown.")
         except Exception as e:
             print(f"Error during batch processor task shutdown: {e}")
             traceback.print_exc()
         batch_processor_task = None
 
+    # Close MongoDB client connection
     if client:
         client.close()
         print("FastAPI app shutting down. MongoDB client closed.")
 
 
+# --- API Endpoints ---
 @app.get("/health", summary="Health Check")
 async def health_check():
-    db_status = "MongoDB not connected. Check server startup logs."
+    """
+    Checks the health of the API, including MongoDB connection and batch processor status.
+    """
+    db_status = "MongoDB not connected."
     db_ok = False
-    if client is not None and db is not None: # Corrected check
+    if client is not None and db is not None:
         try:
-            db.list_collection_names()
+            # Attempt a simple database operation to confirm connectivity
+            await asyncio.to_thread(db.list_collection_names)
             db_status = "MongoDB connection OK."
             db_ok = True
         except Exception as e:
             db_status = f"MongoDB connection error: {e}"
+            db_ok = False # Explicitly set to False on error
 
-    batch_processor_status = "Batch processor not running or state unknown."
-    if batch_processor_task is not None :
+    batch_processor_status = "Batch processor not running."
+    if batch_processor_task is not None:
         if not batch_processor_task.done():
-             batch_processor_status = "Batch processor is running."
+            batch_processor_status = "Batch processor is running."
         else:
-             batch_processor_status = "Batch processor task is done (may have completed or errored)."
-    
+            # Check if it finished with an exception
+            if batch_processor_task.exception():
+                batch_processor_status = f"Batch processor task ended with exception: {batch_processor_task.exception()}"
+            else:
+                batch_processor_status = "Batch processor task is done (may have completed or cancelled)."
+    else:
+        batch_processor_status = "Batch processor task has not been initialized."
+
     if db_ok:
-         return {"status": "ok", "message": "Email Assistant API is up.", "database": db_status, "batch_processor": batch_processor_status}
+        return {"status": "ok", "message": "Email Assistant API is up.", "database": db_status, "batch_processor": batch_processor_status}
     else:
-        # Return a JSON response for HTTPException as well for consistency
+        # If DB is not OK, return a 503 Service Unavailable
         raise HTTPException(
             status_code=503,
-            detail={"message": "Service unavailable.", "database": db_status, "batch_processor": batch_processor_status}
+            detail={"message": "Service unavailable due to issues.", "database": db_status, "batch_processor": batch_processor_status}
         )
 
 
 @app.post("/extract-data", response_model=ExtractedData, summary="Extract structured data from an email and store in MongoDB")
 async def extract_email_data(request: ProcessEmailRequest):
+    """
+    Receives an email, extracts contacts, appointments, and tasks using an LLM,
+    and stores the extracted data in MongoDB.
+    """
     if extracted_emails_collection is None:
-        raise HTTPException(status_code=503, detail="MongoDB not available for extracted_emails.")
+        raise HTTPException(status_code=503, detail="MongoDB not available for extracted email storage.")
     try:
         current_date_val = date.today()
+        # Call the internal processing function in a separate thread to not block the event loop
         extracted_data = await asyncio.to_thread(
             _process_email_internal, request.email_text, request.groq_api_key, current_date_val
         )
+
+        # Prepare data for MongoDB insertion: convert date objects to datetime for storage
         extracted_data_dict = extracted_data.model_dump(by_alias=True, exclude_none=True)
-        for appt in extracted_data_dict.get('appointments', []):
-            if isinstance(appt.get('start_date'), date): appt['start_date'] = datetime.combine(appt['start_date'], datetime.min.time())
-            if isinstance(appt.get('end_date'), date) and appt.get('end_date') is not None: appt['end_date'] = datetime.combine(appt['end_date'], datetime.min.time())
-        for task_item in extracted_data_dict.get('tasks', []):
-            if isinstance(task_item.get('due_date'), date): task_item['due_date'] = datetime.combine(task_item['due_date'], datetime.min.time())
+        if 'appointments' in extracted_data_dict:
+            for appt in extracted_data_dict['appointments']:
+                if isinstance(appt.get('start_date'), date):
+                    appt['start_date'] = datetime.combine(appt['start_date'], datetime.min.time())
+                if isinstance(appt.get('end_date'), date) and appt.get('end_date') is not None:
+                    appt['end_date'] = datetime.combine(appt['end_date'], datetime.min.time())
+        if 'tasks' in extracted_data_dict:
+            for task_item in extracted_data_dict['tasks']:
+                if isinstance(task_item.get('due_date'), date):
+                    task_item['due_date'] = datetime.combine(task_item['due_date'], datetime.min.time())
 
+        # Insert into MongoDB
         result = await asyncio.to_thread(extracted_emails_collection.insert_one, extracted_data_dict)
-        # Pydantic model expects string ID, convert from ObjectId before assigning if needed
-        # However, PyObjectId should handle this if result.inserted_id is ObjectId
+
+        # Update the Pydantic model's ID with the generated MongoDB ObjectId for the response
         extracted_data.id = str(result.inserted_id) if isinstance(result.inserted_id, ObjectId) else result.inserted_id
         return extracted_data
     except ValueError as e:
@@ -570,69 +664,95 @@ async def extract_email_data(request: ProcessEmailRequest):
 
 @app.post("/extract-data-excel", summary="Extract structured data and download as Excel (also stores in MongoDB)")
 async def extract_email_data_excel(request: ProcessEmailRequest):
+    """
+    Placeholder for future functionality to extract data and provide as an Excel download.
+    Currently disabled.
+    """
     raise HTTPException(status_code=501, detail="Excel functionality is currently disabled.")
 
 
 @app.post("/generate-reply", summary="Generate a smart reply to an email (batched & cached)")
 async def generate_email_reply(request: GenerateReplyRequest):
+    """
+    Generates an intelligent email reply based on specified parameters (language, length, style, tone, emoji).
+    Uses a batch processing system with caching for efficiency.
+    """
     if generated_replies_collection is None or batch_processor_task is None or reply_queue_condition is None:
-        raise HTTPException(status_code=503, detail="Reply generation service not fully initialized. Check server logs.")
+        raise HTTPException(status_code=503, detail="Reply generation service not fully initialized. Check server logs for database or batch processor errors.")
 
     future = asyncio.Future()
     current_time = asyncio.get_event_loop().time()
 
     async with reply_queue_condition:
         reply_request_queue.append((request, future, current_time))
-        reply_queue_condition.notify()
+        reply_queue_condition.notify() # Notify the batch processor that there's a new item
 
     try:
+        # A generous timeout for the client waiting for a response
         client_timeout = BATCH_TIMEOUT + 10.0
         result = await asyncio.wait_for(future, timeout=client_timeout)
         return result
     except asyncio.TimeoutError:
+        # If the client times out, cancel the future to clean up
         if not future.done():
             future.cancel()
-        raise HTTPException(status_code=504, detail=f"Request timed out after {client_timeout}s waiting for batch processing.")
+        raise HTTPException(status_code=504, detail=f"Request timed out after {client_timeout}s waiting for batch processing. This may indicate high load or an issue with the batch processor.")
     except Exception as e:
         if isinstance(e, HTTPException):
-            raise e
+            raise e # Re-raise FastAPI HTTPExceptions directly
         traceback.print_exc()
         raise HTTPException(status_code=500, detail=f"Error processing your reply request: {str(e)}")
 
 
 @app.get("/query-extracted-emails", response_model=List[ExtractedData], summary="Query extracted emails from MongoDB")
 async def query_extracted_emails_endpoint(query_params: ExtractedEmailQuery = Depends()):
+    """
+    Queries stored extracted email data from MongoDB based on various filters.
+    """
     if extracted_emails_collection is None:
-        raise HTTPException(status_code=503, detail="MongoDB not available for extracted_emails.")
+        raise HTTPException(status_code=503, detail="MongoDB not available for querying extracted emails.")
+
     mongo_query: Dict[str, Any] = {}
-    if query_params.contact_name: mongo_query["contacts.name"] = {"$regex": query_params.contact_name, "$options": "i"}
-    if query_params.appointment_title: mongo_query["appointments.title"] = {"$regex": query_params.appointment_title, "$options": "i"}
-    if query_params.task_title: mongo_query["tasks.task_title"] = {"$regex": query_params.task_title, "$options": "i"}
+    if query_params.contact_name:
+        mongo_query["contacts.name"] = {"$regex": query_params.contact_name, "$options": "i"} # Case-insensitive regex
+    if query_params.appointment_title:
+        mongo_query["appointments.title"] = {"$regex": query_params.appointment_title, "$options": "i"}
+    if query_params.task_title:
+        mongo_query["tasks.task_title"] = {"$regex": query_params.task_title, "$options": "i"}
 
     if query_params.from_date or query_params.to_date:
         date_query: Dict[str, datetime] = {}
-        if query_params.from_date: date_query["$gte"] = datetime.combine(query_params.from_date, datetime.min.time())
-        if query_params.to_date: date_query["$lt"] = datetime.combine(query_params.to_date + timedelta(days=1), datetime.min.time())
-        if date_query : mongo_query["processed_at"] = date_query
+        if query_params.from_date:
+            # Query for documents processed on or after the start of from_date
+            date_query["$gte"] = datetime.combine(query_params.from_date, datetime.min.time())
+        if query_params.to_date:
+            # Query for documents processed before the start of the day *after* to_date
+            date_query["$lt"] = datetime.combine(query_params.to_date + timedelta(days=1), datetime.min.time())
+        if date_query:
+            mongo_query["processed_at"] = date_query
 
     try:
+        # Sort by processed_at in descending order (most recent first)
         cursor = extracted_emails_collection.find(mongo_query).sort("processed_at", -1).limit(query_params.limit)
         extracted_docs_raw = await asyncio.to_thread(list, cursor)
-        
+
         results = []
         for doc_raw in extracted_docs_raw:
             # Convert _id to string for Pydantic model if it's an ObjectId
             if isinstance(doc_raw.get("_id"), ObjectId):
                 doc_raw["_id"] = str(doc_raw["_id"])
-            
-            # Convert datetime objects back to date objects for Pydantic model fields that are `date`
+
+            # Convert datetime objects from MongoDB back to date objects for Pydantic model fields that are `date`
             if 'appointments' in doc_raw:
                 for appt in doc_raw['appointments']:
-                    if isinstance(appt.get('start_date'), datetime): appt['start_date'] = appt['start_date'].date()
-                    if isinstance(appt.get('end_date'), datetime): appt['end_date'] = appt['end_date'].date()
+                    if isinstance(appt.get('start_date'), datetime):
+                        appt['start_date'] = appt['start_date'].date()
+                    if isinstance(appt.get('end_date'), datetime) and appt.get('end_date') is not None:
+                        appt['end_date'] = appt['end_date'].date()
             if 'tasks' in doc_raw:
                 for task_item in doc_raw['tasks']:
-                    if isinstance(task_item.get('due_date'), datetime): task_item['due_date'] = task_item['due_date'].date()
+                    if isinstance(task_item.get('due_date'), datetime):
+                        task_item['due_date'] = task_item['due_date'].date()
             results.append(ExtractedData(**doc_raw))
         return results
     except Exception as e:
@@ -642,19 +762,29 @@ async def query_extracted_emails_endpoint(query_params: ExtractedEmailQuery = De
 
 @app.get("/query-generated-replies", response_model=List[GeneratedReplyData], summary="Query generated replies from MongoDB")
 async def query_generated_replies_endpoint(query_params: GeneratedReplyQuery = Depends()):
+    """
+    Queries stored generated email replies from MongoDB based on various filters.
+    """
     if generated_replies_collection is None:
-        raise HTTPException(status_code=503, detail="MongoDB not available for generated_replies.")
+        raise HTTPException(status_code=503, detail="MongoDB not available for querying generated replies.")
+
     mongo_query: Dict[str, Any] = {}
-    if query_params.language: mongo_query["language"] = query_params.language
-    if query_params.style: mongo_query["style"] = query_params.style
-    if query_params.tone: mongo_query["tone"] = query_params.tone
+    if query_params.language:
+        mongo_query["language"] = query_params.language
+    if query_params.style:
+        mongo_query["style"] = query_params.style
+    if query_params.tone:
+        mongo_query["tone"] = query_params.tone
 
     if query_params.from_date or query_params.to_date:
         date_query: Dict[str, datetime] = {}
-        if query_params.from_date: date_query["$gte"] = datetime.combine(query_params.from_date, datetime.min.time())
-        if query_params.to_date: date_query["$lt"] = datetime.combine(query_params.to_date + timedelta(days=1), datetime.min.time())
-        if date_query: mongo_query["generated_at"] = date_query
-        
+        if query_params.from_date:
+            date_query["$gte"] = datetime.combine(query_params.from_date, datetime.min.time())
+        if query_params.to_date:
+            date_query["$lt"] = datetime.combine(query_params.to_date + timedelta(days=1), datetime.min.time())
+        if date_query:
+            mongo_query["generated_at"] = date_query
+
     try:
         cursor = generated_replies_collection.find(mongo_query).sort("generated_at", -1).limit(query_params.limit)
         generated_docs_raw = await asyncio.to_thread(list, cursor)
@@ -666,5 +796,4 @@ async def query_generated_replies_endpoint(query_params: GeneratedReplyQuery = D
         return results
     except Exception as e:
         traceback.print_exc()
-        raise HTTPException(status_code=500, detail=f"Error querying generated replies: {e}")
-
+        raise HTTPException(status_code=500, detail=f"Error querying generated replies: {e}")