Update flask_Character.py

flask_Character.py

@@ -60,8 +60,8 @@ class CustomObjectId(str):
         # This validator is only called if the field is not None
         # Pydantic's Optional[PyObjectId] handles the None case before this validator
         if v is None or v == "":
-            return None # Should not be reached if Optional[PyObjectId] is used correctly
-        
+            return None
+
         if not isinstance(v, (str, ObjectId)):
             raise ValueError("ObjectId must be a string or ObjectId instance")
 
@@ -180,6 +180,12 @@ class GeneratedReplyData(BaseModel):
             data["_id"] = str(data["_id"])
         return data
 
+# NEW: Response Model for /generate-reply endpoint
+class GenerateReplyResponse(BaseModel):
+    reply: str = Field(..., description="The AI-generated reply text.")
+    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, description="Filter by contact name (case-insensitive partial match).")
@@ -391,7 +397,10 @@ async def handle_single_reply_request(request_data: GenerateReplyRequest, future
         return
     try:
         if generated_replies_collection is None:
-            raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Database service not available for caching/storage.")
+            # If DB is not available, set a specific exception on the future
+            if not future.done():
+                future.set_exception(HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Database service not available for caching/storage."))
+            return
 
         cache_query = {
             "original_email_text": request_data.email_text,
@@ -469,15 +478,22 @@ async def process_reply_batches():
                     continue
 
             now = asyncio.get_event_loop().time()
-            oldest_item_timestamp = reply_request_queue[0][2]
+            # Safety check: ensure queue is not empty before accessing index 0
+            if reply_request_queue:
+                oldest_item_timestamp = reply_request_queue[0][2]
+            else:
+                # If queue became empty while waiting, loop again
+                continue
 
             # Condition to trigger batch processing: queue is full OR timeout reached for oldest item
             if len(reply_request_queue) >= MAX_BATCH_SIZE or \
                (now - oldest_item_timestamp >= BATCH_TIMEOUT):
                 num_to_take = min(len(reply_request_queue), MAX_BATCH_SIZE)
                 for _ in range(num_to_take):
-                    req, fut, _ = reply_request_queue.pop(0)
-                    batch_to_fire.append((req, fut))
+                    # Safety check: ensure queue is not empty before popping
+                    if reply_request_queue:
+                        req, fut, _ = reply_request_queue.pop(0)
+                        batch_to_fire.append((req, fut))
             else:
                 # Calculate time to wait for the next batch or timeout
                 time_to_wait = BATCH_TIMEOUT - (now - oldest_item_timestamp)
@@ -492,7 +508,6 @@ async def process_reply_batches():
             await asyncio.gather(*tasks)
         else:
             # Short sleep to prevent busy-waiting if queue is empty but not waiting
-            # (e.g., if a notify happened just before the wait, but queue was already processed)
             await asyncio.sleep(0.001)
 
 
@@ -541,7 +556,6 @@ async def startup_event():
         print(f"Successfully connected to MongoDB: {DB_NAME}")
 
         # Start the batch processor task if not already running
-        # Using asyncio.create_task is the correct way to run background tasks in FastAPI
         if batch_processor_task is None or batch_processor_task.done():
             batch_processor_task = asyncio.create_task(process_reply_batches())
             print("Batch processor task for replies started.")
@@ -550,7 +564,6 @@ async def startup_event():
 
     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
@@ -563,17 +576,14 @@ async def startup_event():
         extracted_emails_collection = None
         generated_replies_collection = None
     finally:
-        # 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 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. 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 startup sequence completed for MongoDB client & Batch Processor initialization.")
@@ -583,11 +593,9 @@ async def startup_event():
 async def shutdown_event():
     global client, batch_processor_task
     print("FastAPI app shutting down.")
-    # 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 during shutdown.")
@@ -596,7 +604,6 @@ async def shutdown_event():
             traceback.print_exc()
         batch_processor_task = None
 
-    # Close MongoDB client connection
     if client:
         client.close()
         print("MongoDB client closed.")
@@ -612,21 +619,18 @@ async def health_check():
     db_ok = False
     if client is not None and db is not None:
         try:
-            # Attempt a simple database operation to confirm connectivity
-            # For async functions, ensure you use await asyncio.to_thread for blocking MongoDB operations
             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
+            db_ok = False
 
     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."
         else:
-            # 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:
@@ -637,7 +641,6 @@ async def health_check():
     if db_ok:
         return {"status": "ok", "message": "Email Assistant API is up.", "database": db_status, "batch_processor": batch_processor_status}
     else:
-        # If DB is not OK, return a 503 Service Unavailable
         raise HTTPException(
             status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
             detail={"message": "Service unavailable due to issues.", "database": db_status, "batch_processor": batch_processor_status}
@@ -654,16 +657,13 @@ async def extract_email_data(request: ProcessEmailRequest):
         raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="MongoDB not available for extracted email storage. Check server startup logs.")
     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)
         if 'appointments' in extracted_data_dict:
             for appt in extracted_data_dict['appointments']:
-                # MongoDB stores dates as datetime.datetime, so convert
                 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:
@@ -673,11 +673,8 @@ 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 into MongoDB
         result = await asyncio.to_thread(extracted_emails_collection.insert_one, extracted_data_dict)
 
-        # Update the Pydantic model's ID with the generated MongoDB ObjectId for the response
-        # PyObjectId's __get_validators__ should handle this conversion from ObjectId to str
         extracted_data.id = result.inserted_id
         return extracted_data
     except ValueError as e:
@@ -696,7 +693,7 @@ async def extract_email_data_excel(request: ProcessEmailRequest):
     raise HTTPException(status_code=status.HTTP_501_NOT_IMPLEMENTED, detail="Excel functionality is currently disabled.")
 
 
-@app.post("/generate-reply", summary="Generate a smart reply to an email (batched & cached)")
+@app.post("/generate-reply", response_model=GenerateReplyResponse, 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).
@@ -713,18 +710,14 @@ async def generate_email_reply(request: GenerateReplyRequest):
         reply_queue_condition.notify() # Notify the batch processor that a new request is available
 
     try:
-        # Give a reasonable timeout for the client to wait for a reply
-        # This timeout should be greater than BATCH_TIMEOUT
         client_timeout = BATCH_TIMEOUT + 10.0 # e.g., 0.5s batch + 10s LLM response buffer
         result = await asyncio.wait_for(future, timeout=client_timeout)
         return result
     except asyncio.TimeoutError:
-        # If the client times out, ensure the future is cancelled if not already done
         if not future.done():
             future.cancel()
         raise HTTPException(status_code=status.HTTP_504_GATEWAY_TIMEOUT, detail=f"Request timed out after {client_timeout}s waiting for batch processing. The LLM might be busy or the request queue too long.")
     except Exception as e:
-        # Re-raise HTTPException if it's already one, otherwise wrap in 500
         if isinstance(e, HTTPException):
             raise e
         traceback.print_exc()
@@ -759,12 +752,6 @@ async def query_extracted_emails_endpoint(query_params: ExtractedEmailQuery = De
 
         results = []
         for doc_raw in extracted_docs_raw:
-            # Convert _id to string for Pydantic model if it's an ObjectId
-            # PyObjectId type hint handles this on model parsing
-            # if isinstance(doc_raw.get("_id"), ObjectId):
-            #     doc_raw["_id"] = str(doc_raw["_id"])
-
-            # 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()
@@ -793,7 +780,6 @@ async def query_generated_replies_endpoint(query_params: GeneratedReplyQuery = D
         if query_params.from_date:
             date_query["$gte"] = datetime.combine(query_params.from_date, datetime.min.time())
         if query_params.to_date:
-            # Query up to the end of the 'to_date' day
             date_query["$lt"] = datetime.combine(query_params.to_date + timedelta(days=1), datetime.min.time())
         if date_query:
             mongo_query["generated_at"] = date_query
@@ -803,9 +789,6 @@ async def query_generated_replies_endpoint(query_params: GeneratedReplyQuery = D
         generated_docs_raw = await asyncio.to_thread(list, cursor)
         results = []
         for doc_raw in generated_docs_raw:
-            # PyObjectId type hint handles this on model parsing
-            # if isinstance(doc_raw.get("_id"), ObjectId):
-            #     doc_raw["_id"] = str(doc_raw["_id"])
             results.append(GeneratedReplyData(**doc_raw))
         return results
     except Exception as e: