Update flask_Character.py

flask_Character.py

@@ -1,21 +1,12 @@
-# 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.
-import os
-# NUMBA_CACHE_DIR and NUMBA_DISABLE_CACHE are often set for specific environments,
-# e.g., if you're experiencing issues with Numba's caching behavior or in containerized environments.
-# Keep them if they serve a specific purpose in your deployment environment.
-os.environ["NUMBA_CACHE_DIR"] = "/tmp/numba_cache"
-os.environ["NUMBA_DISABLE_CACHE"] = "1"
-
 import json
 import re
 from datetime import date, datetime, timedelta
 from typing import List, Optional, Literal, Dict, Any, Tuple
+import os
 import traceback
 import asyncio
 
-from fastapi import FastAPI, HTTPException, Response, Query, Depends, status
+from fastapi import FastAPI, HTTPException, Response, Query, Depends
 from fastapi.responses import FileResponse
 from fastapi.exception_handlers import http_exception_handler
 from starlette.exceptions import HTTPException as StarletteHTTPException
@@ -23,74 +14,45 @@ 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__
+import uvicorn
 
 from pymongo import MongoClient
 from pymongo.errors import ConnectionFailure, OperationFailure
 from bson import ObjectId
 
 # --- MongoDB Configuration ---
-# IMPORTANT: Use environment variables for your MONGO_URI in production for security.
-# Example: MONGO_URI = os.getenv("MONGO_URI", "mongodb://localhost:27017")
-MONGO_URI = "mongodb+srv://precison9:P1LhtFknkT75yg5L@cluster0.isuwpef.mongodb.net"
+# Load from environment variables for secure deployment on Hugging Face
+MONGO_URI = os.getenv("MONGO_URI", "mongodb+srv://user:pass@cluster.mongodb.net/dbname")
 DB_NAME = "email_assistant_db"
 EXTRACTED_EMAILS_COLLECTION = "extracted_emails"
-GENERATED_REPLIES_COLLECTION = "generated_replies" # Still defined, but not used by generate-reply logic
+GENERATED_REPLIES_COLLECTION = "generated_replies"
 
-# Global variables for MongoDB client and collections
 client: Optional[MongoClient] = None
 db: Optional[Any] = None
 extracted_emails_collection: Optional[Any] = None
-# generated_replies_collection is no longer needed for /generate-reply logic,
-# but kept for /query-generated-replies endpoint if that's still desired.
 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):
-        # Allow None or empty string to pass through for optional fields
-        # 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
-
-        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
+    def validate(cls, v, info):
         if not ObjectId.is_valid(v):
-            raise ValueError("Invalid ObjectId format")
-        return cls(v) # Return an instance of CustomObjectId (which is a str subclass)
+            raise ValueError("Invalid ObjectId")
+        return str(v)
 
-    # This method is crucial for Pydantic v2 to generate correct OpenAPI schema
     @classmethod
-    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).
-        json_schema = handler(core_schema.str_schema())
-        json_schema["example"] = "60c728ef238b9c7b9e0f6c2a" # Add an example for clarity
+    def __get_pydantic_json_schema__(cls, core_schema, handler):
+        json_schema = handler(core_schema)
+        json_schema["type"] = "string"
+        json_schema["example"] = "60c728ef238b9c7b9e0f6c2a"
         return json_schema
 
-# Annotated type for convenience in models
 PyObjectId = Annotated[CustomObjectId, BeforeValidator(str)]
 
-
 # ---------------------- Models ----------------------
 class Contact(BaseModel):
     name: str
@@ -112,7 +74,6 @@ 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]
@@ -121,19 +82,14 @@ class ExtractedData(BaseModel):
     processed_at: datetime = Field(default_factory=datetime.utcnow)
 
     class Config:
-        populate_by_name = True # Allow setting 'id' or '_id'
-        arbitrary_types_allowed = True # Allow CustomObjectId and ObjectId
+        populate_by_name = True
+        arbitrary_types_allowed = True
 
-    # 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
-        # Pydantic v2 usually handles this automatically for `date` types,
-        # but explicit conversion can be useful if direct manipulation is expected or for specific formats.
         if 'appointments' in data:
             for appt in data['appointments']:
                 if isinstance(appt.get('start_date'), date):
@@ -148,11 +104,11 @@ class ExtractedData(BaseModel):
 
 class ProcessEmailRequest(BaseModel):
     email_text: str = Field(..., example="Oggetto: Follow-up progetto “Delta”...")
-    groq_api_key: str = Field(..., example="YOUR_GROQ_API_KEY")
+    groq_api_key: str = Field(..., example="YOUR_GROQ_API_KEY") # Should be handled securely, see notes below
 
 class GenerateReplyRequest(BaseModel):
     email_text: str = Field(..., example="Oggetto: Follow-up progetto “Delta”...")
-    groq_api_key: str = Field(..., example="YOUR_GROQ_API_KEY")
+    groq_api_key: str = Field(..., example="YOUR_GROQ_API_KEY") # Should be handled securely, see notes below
     language: Literal["Italian", "English"] = Field("Italian", examples=["Italian", "English"])
     length: str = Field("Auto", examples=["Short", "Medium", "Long", "Auto"])
     style: str = Field("Professional", examples=["Professional", "Casual", "Formal", "Informal"])
@@ -160,7 +116,6 @@ class GenerateReplyRequest(BaseModel):
     emoji: str = Field("Auto", examples=["Auto", "None", "Occasional", "Frequent"])
 
 class GeneratedReplyData(BaseModel):
-    # Use PyObjectId for the _id field (This model is now only used for the query endpoint)
     id: Optional[PyObjectId] = Field(alias="_id", default=None)
     original_email_text: str
     generated_reply_text: str
@@ -182,12 +137,6 @@ class GeneratedReplyData(BaseModel):
             data["_id"] = str(data["_id"])
         return data
 
-# Response Model for /generate-reply endpoint (simplified)
-class GenerateReplyResponse(BaseModel):
-    reply: str = Field(..., description="The AI-generated reply text.")
-    # 'stored_id' and 'cached' are removed as caching/storage is removed
-    # from the main generate-reply logic.
-
 # --- Query Models for GET Endpoints ---
 class ExtractedEmailQuery(BaseModel):
     contact_name: Optional[str] = Query(None, description="Filter by contact name (case-insensitive partial match).")
@@ -207,44 +156,26 @@ 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]:
-    """
-    Parses a date string, handling 'today', 'tomorrow', and APAC-MM-DD format.
-    Returns None if input is None or cannot be parsed into a valid date.
-    """
-    if not date_str:
-        return None
+    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:
-        # If parsing fails, return None. The calling function (normalize_llm_output)
-        # will then decide the default (e.g., current_date).
-        return None
+        return current_date
 
 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 ""
@@ -258,9 +189,7 @@ def normalize_llm_output(data: dict, current_date: date, original_email_text: st
 
     appointments_data = []
     for a in data.get("appointments", []):
-        # 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
+        start_date_val = parse_date(a.get("start_date", "today"), current_date) or current_date
         end_date_val = parse_date(a.get("end_date"), current_date)
 
         appointments_data.append(Appointment(
@@ -271,8 +200,7 @@ def normalize_llm_output(data: dict, current_date: date, original_email_text: st
 
     tasks_data = []
     for t in data.get("tasks", []):
-        # 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
+        due_date_val = parse_date(t.get("due_date", "today"), 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
@@ -281,17 +209,10 @@ 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:
-    """
-    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.")
-
+    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.
 
@@ -313,34 +234,30 @@ Here is the required JSON schema for each category:
   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, APAC-MM-DD. If not explicitly mentioned, use "{prompt_today_str}" for "today", or "{prompt_tomorrow_str}" for "tomorrow")
+    - `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, APAC-MM-DD, optional, null if unknown or not applicable)
+    - `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, APAC-MM-DD. Infer from context, e.g., "entro domani" becomes "{prompt_tomorrow_str}", "today" becomes "{prompt_today_str}")
+    - `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", "prompt_today_str", "prompt_tomorrow_str"], template=prompt_template_str)
+    prompt_template = PromptTemplate(input_variables=["email"], template=prompt_template_str) # Removed prompt_today_str and prompt_tomorrow_str as they are in the template string
     chain = prompt_template | llm
     try:
-        llm_output = chain.invoke({"email": email_text, "prompt_today_str": prompt_today_str, "prompt_tomorrow_str": prompt_tomorrow_str})
+        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}")
+        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)
-
         extracted_data = normalize_llm_output(json_data, current_date, email_text)
         return extracted_data
     except json.JSONDecodeError as e:
@@ -353,48 +270,34 @@ def _generate_response_internal(
     email_text: str, api_key: str, language: Literal["Italian", "English"],
     length: str, style: str, tone: str, emoji: str
 ) -> str:
-    """
-    Internal function to generate a reply to an email using LLM.
-    """
-    print(f"[{datetime.now()}] _generate_response_internal: Starting LLM call. API Key starts with: {api_key[:5]}...") # Debug log
-    if not email_text:
-        print(f"[{datetime.now()}] _generate_response_internal: Email text is empty.")
-        return "Cannot generate reply for empty email text."
+    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)
+    prompt_template_str="""
+You are an assistant that helps reply to emails.
+
+Create a response to the following email with the following parameters:
+- Language: {language}
+- Length: {length}
+- Style: {style}
+- Tone: {tone}
+- Emoji usage: {emoji}
 
-    try:
-        llm = ChatGroq(model="meta-llama/llama-4-scout-17b-16e-instruct", temperature=0.7, max_tokens=800, groq_api_key=api_key)
-        prompt_template_str="""
-    You are an assistant that helps reply to emails.
-
-    Create a response to the following email with the following parameters:
-    - Language: {language}
-    - Length: {length}
-    - Style: {style}
-    - Tone: {tone}
-    - Emoji usage: {emoji}
-
-    Email:
-    {email}
-
-    Write only the reply body. Do not repeat the email or mention any instruction.
-    """
-        prompt = PromptTemplate(
-            input_variables=["email", "language", "length", "style", "tone", "emoji"],
-            template=prompt_template_str
-        )
-        chain = prompt | llm
-        print(f"[{datetime.now()}] _generate_response_internal: Invoking LLM chain...") # Debug log
-        output = chain.invoke({"email": email_text, "language": language, "length": length, "style": style, "tone": tone, "emoji": emoji})
-        print(f"[{datetime.now()}] _generate_response_internal: LLM chain returned. Content length: {len(output.content)}.") # Debug log
-        return output.content.strip()
-    except Exception as e:
-        print(f"[{datetime.now()}] _generate_response_internal: ERROR during LLM invocation: {e}") # Debug log
-        traceback.print_exc() # Print full traceback to logs
-        raise # Re-raise the exception so it can be caught by handle_single_reply_request
+Email:
+{email}
+
+Write only the reply body. Do not repeat the email or mention any instruction.
+"""
+    prompt = PromptTemplate(
+        input_variables=["email", "language", "length", "style", "tone", "emoji"],
+        template=prompt_template_str
+    )
+    chain = prompt | llm
+    output = chain.invoke({"email": email_text, "language": language, "length": length, "style": style, "tone": tone, "emoji": emoji})
+    return output.content.strip()
 
-# --- Batching Configuration (Caching/Storage logic removed) ---
+# --- Batching and Caching Configuration ---
 MAX_BATCH_SIZE = 20
-BATCH_TIMEOUT = 0.5 # seconds (Adjust based on expected LLM response time and desired latency)
+BATCH_TIMEOUT = 0.5 
 
 reply_request_queue: List[Tuple[GenerateReplyRequest, asyncio.Future, float]] = []
 reply_queue_lock = asyncio.Lock()
@@ -402,101 +305,103 @@ reply_queue_condition = asyncio.Condition(lock=reply_queue_lock)
 batch_processor_task: Optional[asyncio.Task] = None
 
 
-# --- Batch Processor and Handler (Simplified) ---
+# --- Batch Processor and Handler ---
 async def handle_single_reply_request(request_data: GenerateReplyRequest, future: asyncio.Future):
-    """Handles a single request: calls LLM, and sets future with the reply."""
-    print(f"[{datetime.now()}] Handle single reply: Starting for email_text_start='{request_data.email_text[:50]}'...")
+    """Handles a single request: checks cache, calls LLM, stores result, and sets future."""
     if future.cancelled():
-        print(f"[{datetime.now()}] Handle single reply: Future cancelled. Aborting.")
         return
     try:
-        # Directly call LLM (no cache check or storage)
-        print(f"[{datetime.now()}] Handle single reply: Calling LLM for reply generation...")
+        if generated_replies_collection is None:
+            raise HTTPException(status_code=503, detail="Database service not available for caching/storage.")
+
+        cache_query = {
+            "original_email_text": request_data.email_text,
+            "language": request_data.language,
+            "length": request_data.length,
+            "style": request_data.style,
+            "tone": request_data.tone,
+            "emoji": request_data.emoji,
+        }
+        # Use 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:
+            response = {
+                "reply": cached_reply_doc["generated_reply_text"],
+                "stored_id": str(cached_reply_doc["_id"]),
+                "cached": True
+            }
+            if not future.done(): future.set_result(response)
+            return
+
         reply_content = await asyncio.to_thread(
             _generate_response_internal,
             request_data.email_text,
-            request_data.groq_api_key,
+            request_data.groq_api_key, # Groq API key is passed here
             request_data.language,
             request_data.length,
             request_data.style,
             request_data.tone,
             request_data.emoji
         )
-        print(f"[{datetime.now()}] Handle single reply: LLM call completed. Reply length: {len(reply_content)}.")
 
-        # Simplified response as no storage/cache ID
+        reply_data_to_store = GeneratedReplyData(
+            original_email_text=request_data.email_text,
+            generated_reply_text=reply_content,
+            language=request_data.language,
+            length=request_data.length,
+            style=request_data.style,
+            tone=request_data.tone,
+            emoji=request_data.emoji
+        )
+        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)
+        stored_id = str(insert_result.inserted_id)
+
         final_response = {
             "reply": reply_content,
-            "stored_id": "N/A - Caching disabled", # Indicate that ID is not available
+            "stored_id": stored_id,
             "cached": False
         }
-        if not future.done():
-            future.set_result(final_response)
-        print(f"[{datetime.now()}] Handle single reply: Final result set on future.")
+        if not future.done(): future.set_result(final_response)
 
     except Exception as e:
-        print(f"[{datetime.now()}] Handle single reply: EXCEPTION: {e}")
-        traceback.print_exc() # Print full traceback to logs
+        traceback.print_exc()
         if not future.done():
-            # Set the exception on the future so the client can catch it
-            future.set_exception(HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to generate reply: {e}"))
-        print(f"[{datetime.now()}] Handle single reply: Exception set on future.")
-
+            future.set_exception(e)
 
 async def process_reply_batches():
     """Continuously processes requests from the reply_request_queue in batches."""
     global reply_request_queue
-    print(f"[{datetime.now()}] Batch processor task started.")
     while True:
         batch_to_fire: List[Tuple[GenerateReplyRequest, asyncio.Future]] = []
         async with reply_queue_condition:
             if not reply_request_queue:
-                print(f"[{datetime.now()}] Batch processor: Queue empty, waiting for requests...")
-                # Wait for new requests or timeout
                 await reply_queue_condition.wait()
-                # After waking up, re-check if queue is still empty
                 if not reply_request_queue:
-                    print(f"[{datetime.now()}] Batch processor: Woke up, queue still empty. Continuing loop.")
                     continue
 
             now = asyncio.get_event_loop().time()
-            # 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
-                print(f"[{datetime.now()}] Batch processor: Queue became empty before processing. Restarting loop.")
-                continue
-
-            print(f"[{datetime.now()}] Batch processor: Woke up. Queue size: {len(reply_request_queue)}. Oldest item age: {now - oldest_item_timestamp:.2f}s")
+            oldest_item_timestamp = reply_request_queue[0][2]
 
-            # 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):
-                    # 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))
-                print(f"[{datetime.now()}] Batch processor: Firing batch of {len(batch_to_fire)} requests.")
+                    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)
-                print(f"[{datetime.now()}] Batch processor: Not enough requests or timeout not reached. Waiting for {time_to_wait:.2f}s.")
                 try:
                     await asyncio.wait_for(reply_queue_condition.wait(), timeout=time_to_wait)
                 except asyncio.TimeoutError:
-                    print(f"[{datetime.now()}] Batch processor: wait timed out.")
-                    pass # Loop will re-evaluate and likely fire the batch
+                    pass
 
         if batch_to_fire:
             tasks = [handle_single_reply_request(req_data, fut) for req_data, fut in batch_to_fire]
-            print(f"[{datetime.now()}] Batch processor: Awaiting completion of {len(tasks)} single reply tasks.")
             await asyncio.gather(*tasks)
-            print(f"[{datetime.now()}] Batch processor: Batch processing complete.")
         else:
-            # Short sleep to prevent busy-waiting if queue is empty but not waiting
             await asyncio.sleep(0.001)
 
 
@@ -505,62 +410,53 @@ 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",
-    docs_url="/", # Sets Swagger UI to be the root path
+    docs_url="/",
     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):
-    """Handles FastAPI's internal HTTP exceptions."""
-    print(f"[{datetime.now()}] Caught StarletteHTTPException: {exc.status_code} - {exc.detail}")
     return await http_exception_handler(request, exc)
 
-# Catch all other unhandled exceptions
 @app.exception_handler(Exception)
 async def global_exception_handler_wrapper(request, exc):
-    """Handles all unhandled exceptions and returns a consistent JSON error response."""
-    print(f"[{datetime.now()}] Unhandled exception caught by global handler for request: {request.url}")
-    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)}", "type": "unhandled_exception"}),
-        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-        media_type="application/json"
-    )
+    print(f"Unhandled exception caught by global handler for request: {request.url}")
+    traceback.print_exc()
+    return Response(content=json.dumps({"detail": f"Internal Server Error: {str(exc)}"}), status_code=500, media_type="application/json")
 
 
 # --- FastAPI Event Handlers for MongoDB & Batch Processor ---
 @app.on_event("startup")
 async def startup_event():
     global client, db, extracted_emails_collection, generated_replies_collection, batch_processor_task
-    print(f"[{datetime.now()}] FastAPI app startup sequence initiated.")
     try:
-        # Connect to MongoDB
+        # Check if MONGO_URI is set before attempting connection
+        if not MONGO_URI or MONGO_URI == "mongodb+srv://user:pass@cluster.mongodb.net/dbname":
+            print("WARNING: MONGO_URI environment variable not set or using default. MongoDB connection will fail.")
+            # Optionally raise an exception here or set a flag to disable DB functionality
+            # For now, we'll let the connection attempt proceed and catch its failure.
+
         client = MongoClient(MONGO_URI, serverSelectionTimeoutMS=5000)
-        client.admin.command('ping') # Test connection
+        client.admin.command('ping')
         db = client[DB_NAME]
         extracted_emails_collection = db[EXTRACTED_EMAILS_COLLECTION]
-        # Keep generated_replies_collection definition if /query-generated-replies is still desired
         generated_replies_collection = db[GENERATED_REPLIES_COLLECTION]
-        print(f"[{datetime.now()}] Successfully connected to MongoDB: {DB_NAME}")
+        print(f"Successfully connected to MongoDB: {DB_NAME}")
 
-        # Start the batch processor task if not already running
-        if batch_processor_task is None or batch_processor_task.done():
-            batch_processor_task = asyncio.create_task(process_reply_batches())
-            print(f"[{datetime.now()}] Batch processor task for replies started.")
-        else:
-            print(f"[{datetime.now()}] Batch processor task for replies is already running or being initialized.")
+        if batch_processor_task is None:
+            loop = asyncio.get_event_loop()
+            batch_processor_task = loop.create_task(process_reply_batches())
+            print("Batch processor task for replies started.")
 
     except (ConnectionFailure, OperationFailure) as e:
-        print(f"[{datetime.now()}] ERROR: MongoDB Connection/Operation Failure: {e}")
+        print(f"ERROR: MongoDB Connection/Operation Failure: {e}")
         client = None
         db = None
         extracted_emails_collection = None
         generated_replies_collection = None
     except Exception as e:
-        print(f"[{datetime.now()}] ERROR: An unexpected error occurred during MongoDB connection or batch startup: {e}")
+        print(f"ERROR: An unexpected error occurred during MongoDB connection or batch startup: {e}")
         traceback.print_exc()
         client = None
         db = None
@@ -570,208 +466,147 @@ async def startup_event():
         if client is not None and db is not None:
             try:
                 client.admin.command('ping')
-            except Exception as e:
-                print(f"[{datetime.now()}] MongoDB ping failed after initial connection attempt during finally block: {e}")
+            except Exception:
+                print("MongoDB ping failed after initial connection attempt during finally block.")
                 client = None; db = None; extracted_emails_collection = None; generated_replies_collection = None
         else:
-            print(f"[{datetime.now()}] MongoDB client or db object is None after connection attempt in startup. Database likely not connected.")
+            print("MongoDB client or db object is None after connection attempt in startup.")
             if client is None or db is None:
                 client = None; db = None; extracted_emails_collection = None; generated_replies_collection = None
-        print(f"[{datetime.now()}] FastAPI app startup sequence completed for MongoDB client & Batch Processor initialization.")
+        print("FastAPI app starting up. MongoDB client & Batch Processor initialization attempted.")
 
 
 @app.on_event("shutdown")
 async def shutdown_event():
     global client, batch_processor_task
-    print(f"[{datetime.now()}] FastAPI app shutting down.")
     if batch_processor_task:
         batch_processor_task.cancel()
         try:
             await batch_processor_task
-            print(f"[{datetime.now()}] Batch processor task awaited.")
         except asyncio.CancelledError:
-            print(f"[{datetime.now()}] Batch processor task for replies cancelled during shutdown.")
+            print("Batch processor task for replies cancelled.")
         except Exception as e:
-            print(f"[{datetime.now()}] Error during batch processor task shutdown: {e}")
+            print(f"Error during batch processor task shutdown: {e}")
             traceback.print_exc()
         batch_processor_task = None
 
     if client:
         client.close()
-        print(f"[{datetime.now()}] MongoDB client closed.")
+        print("FastAPI app shutting down. MongoDB client closed.")
 
 
-# --- API Endpoints ---
 @app.get("/health", summary="Health Check")
 async def health_check():
-    """
-    Checks the health of the API, including MongoDB connection and batch processor status.
-    """
-    db_status = "MongoDB not connected."
+    db_status = "MongoDB not connected. Check server startup logs."
     db_ok = False
     if client is not None and db is not None:
         try:
-            # Use asyncio.to_thread for blocking MongoDB call
-            await asyncio.to_thread(db.list_collection_names)
+            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
 
-    batch_processor_status = "Batch processor not running."
-    if batch_processor_task is not None:
+    batch_processor_status = "Batch processor not running or state unknown."
+    if batch_processor_task is not None :
         if not batch_processor_task.done():
             batch_processor_status = "Batch processor is running."
         else:
-            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."
-
+            batch_processor_status = "Batch processor task is done (may have completed or errored)."
+    
     if db_ok:
         return {"status": "ok", "message": "Email Assistant API is up.", "database": db_status, "batch_processor": batch_processor_status}
     else:
         raise HTTPException(
-            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
-            detail={"message": "Service unavailable due to issues.", "database": db_status, "batch_processor": batch_processor_status}
+            status_code=503,
+            detail={"message": "Service unavailable.", "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.
-    """
-    print(f"[{datetime.now()}] /extract-data: Received request.")
     if extracted_emails_collection is None:
-        print(f"[{datetime.now()}] /extract-data: MongoDB collection is None.")
-        raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="MongoDB not available for extracted email storage. Check server startup logs.")
+        raise HTTPException(status_code=503, detail="MongoDB not available for extracted_emails. Check server logs for connection errors.")
     try:
         current_date_val = date.today()
-        print(f"[{datetime.now()}] /extract-data: Calling internal processing function.")
         extracted_data = await asyncio.to_thread(
             _process_email_internal, request.email_text, request.groq_api_key, current_date_val
         )
-        print(f"[{datetime.now()}] /extract-data: Internal processing complete. Preparing for DB insert.")
-
         extracted_data_dict = extracted_data.model_dump(by_alias=True, exclude_none=True)
-        # Convert date objects to datetime for MongoDB storage if they are just date objects
-        # Pydantic's default `date` handling might serialize to ISO string, but for
-        # internal MongoDB storage, sometimes `datetime` is preferred for consistency.
-        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())
+        # Convert date objects to datetime for MongoDB storage
+        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())
 
-        print(f"[{datetime.now()}] /extract-data: Inserting into MongoDB...")
         result = await asyncio.to_thread(extracted_emails_collection.insert_one, extracted_data_dict)
-        print(f"[{datetime.now()}] /extract-data: Data inserted into MongoDB. ID: {result.inserted_id}")
-
-        extracted_data.id = result.inserted_id
+        extracted_data.id = str(result.inserted_id) if isinstance(result.inserted_id, ObjectId) else result.inserted_id
         return extracted_data
     except ValueError as e:
-        print(f"[{datetime.now()}] /extract-data: ValueError: {e}")
-        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e))
+        raise HTTPException(status_code=400, detail=str(e))
     except Exception as e:
-        print(f"[{datetime.now()}] /extract-data: Unhandled Exception: {e}")
         traceback.print_exc()
-        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Internal server error during data extraction: {e}")
+        raise HTTPException(status_code=500, detail=f"Internal server error during data extraction: {e}")
 
 
 @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=status.HTTP_501_NOT_IMPLEMENTED, detail="Excel functionality is currently disabled.")
+    raise HTTPException(status_code=501, detail="Excel functionality is currently disabled.")
 
 
-@app.post("/generate-reply", response_model=GenerateReplyResponse, summary="Generate a smart reply to an email (batched)")
+@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. Caching and database storage for replies are disabled.
-    """
-    print(f"[{datetime.now()}] /generate-reply: Received request.")
-    # generated_replies_collection check is no longer relevant for this endpoint's logic
-    if batch_processor_task is None or reply_queue_condition is None:
-        print(f"[{datetime.now()}] /generate-reply: Service not fully initialized. batch_task={batch_processor_task is not None}, queue_cond={reply_queue_condition is not None}")
-        raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Reply generation service not fully initialized. Check server logs for batch processor issues.")
+    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.")
 
     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() # Notify the batch processor that a new request is available
-        print(f"[{datetime.now()}] /generate-reply: Request added to queue, notifying batch processor. Queue size: {len(reply_request_queue)}")
+        reply_queue_condition.notify()
 
     try:
-        # Debugging: Use a very long timeout for now to ensure server-side logs complete
-        client_timeout = BATCH_TIMEOUT + 300.0 # 5 minutes (0.5s batch + 300s buffer)
-        print(f"[{datetime.now()}] /generate-reply: Waiting for future result with timeout {client_timeout}s.")
+        client_timeout = BATCH_TIMEOUT + 10.0
         result = await asyncio.wait_for(future, timeout=client_timeout)
-        print(f"[{datetime.now()}] /generate-reply: Future result received. Returning data.")
         return result
     except asyncio.TimeoutError:
-        print(f"[{datetime.now()}] /generate-reply: Client timeout waiting for future after {client_timeout}s. Future done: {future.done()}")
         if not future.done():
-            future.cancel() # Cancel if it's still pending
-        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. Check server logs for more details.")
+            future.cancel()
+        raise HTTPException(status_code=504, detail=f"Request timed out after {client_timeout}s waiting for batch processing.")
     except Exception as e:
         if isinstance(e, HTTPException):
-            print(f"[{datetime.now()}] /generate-reply: Caught HTTPException: {e.status_code} - {e.detail}")
-            raise e # Re-raise FastAPI HTTPExceptions
-        print(f"[{datetime.now()}] /generate-reply: Unhandled Exception: {e}")
+            raise e
         traceback.print_exc()
-        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error processing your reply request: {str(e)}. Check server logs for more details.")
+        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()):
-    print(f"[{datetime.now()}] /query-extracted-emails: Received request with params: {query_params.model_dump_json()}")
     if extracted_emails_collection is None:
-        print(f"[{datetime.now()}] /query-extracted-emails: MongoDB collection is None.")
-        raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="MongoDB not available for querying extracted emails.")
+        raise HTTPException(status_code=503, detail="MongoDB not available for extracted_emails. Check server logs for connection errors.")
     mongo_query: Dict[str, Any] = {}
-    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.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.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:
-            # 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["processed_at"] = date_query
-    print(f"[{datetime.now()}] /query-extracted-emails: MongoDB query built: {mongo_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["processed_at"] = date_query
 
     try:
-        # Use await asyncio.to_thread for blocking MongoDB operations
         cursor = extracted_emails_collection.find(mongo_query).sort("processed_at", -1).limit(query_params.limit)
         extracted_docs_raw = await asyncio.to_thread(list, cursor)
-        print(f"[{datetime.now()}] /query-extracted-emails: Found {len(extracted_docs_raw)} documents.")
-
+        
         results = []
         for doc_raw in extracted_docs_raw:
-            # Convert datetime objects back to date for Pydantic model validation if necessary
+            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`
             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()
@@ -780,21 +615,16 @@ async def query_extracted_emails_endpoint(query_params: ExtractedEmailQuery = De
                 for task_item in doc_raw['tasks']:
                     if isinstance(task_item.get('due_date'), datetime): task_item['due_date'] = task_item['due_date'].date()
             results.append(ExtractedData(**doc_raw))
-        print(f"[{datetime.now()}] /query-extracted-emails: Returning {len(results)} results.")
         return results
     except Exception as e:
-        print(f"[{datetime.now()}] /query-extracted-emails: Unhandled Exception during query: {e}")
         traceback.print_exc()
-        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error querying extracted emails: {e}")
+        raise HTTPException(status_code=500, detail=f"Error querying extracted emails: {e}")
 
 
 @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()):
-    print(f"[{datetime.now()}] /query-generated-replies: Received request with params: {query_params.model_dump_json()}")
-    # This endpoint still relies on `generated_replies_collection`
     if generated_replies_collection is None:
-        print(f"[{datetime.now()}] /query-generated-replies: MongoDB collection is None.")
-        raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="MongoDB not available for querying generated replies.")
+        raise HTTPException(status_code=503, detail="MongoDB not available for generated_replies. Check server logs for connection errors.")
     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
@@ -802,25 +632,26 @@ async def query_generated_replies_endpoint(query_params: GeneratedReplyQuery = D
 
     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
-    print(f"[{datetime.now()}] /query-generated-replies: MongoDB query built: {mongo_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:
-        # Use await asyncio.to_thread for blocking MongoDB operations
         cursor = generated_replies_collection.find(mongo_query).sort("generated_at", -1).limit(query_params.limit)
         generated_docs_raw = await asyncio.to_thread(list, cursor)
-        print(f"[{datetime.now()}] /query-generated-replies: Found {len(generated_docs_raw)} documents.")
         results = []
         for doc_raw in generated_docs_raw:
+            if isinstance(doc_raw.get("_id"), ObjectId):
+                doc_raw["_id"] = str(doc_raw["_id"])
             results.append(GeneratedReplyData(**doc_raw))
-        print(f"[{datetime.now()}] /query-generated-replies: Returning {len(results)} results.")
         return results
     except Exception as e:
-        print(f"[{datetime.now()}] /query-generated-replies: Unhandled Exception during query: {e}")
         traceback.print_exc()
-        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error querying generated replies: {e}")
+        raise HTTPException(status_code=500, detail=f"Error querying generated replies: {e}")
+
+# --- Main execution for Hugging Face Spaces ---
+# Hugging Face Spaces expects the `app` object to be directly available.
+# The `if __name__ == "__main__":` block is usually not needed for deployment.
+# However, for local testing, you can keep it or use `uvicorn app:app --host 0.0.0.0 --port 8000`
+# if __name__ == "__main__":
+#    uvicorn.run(app, host="0.0.0.0", port=8000)