File size: 32,710 Bytes
99b21b2
 
 
efd4eb0
e99343e
 
 
 
 
 
 
 
 
 
 
 
 
 
99b21b2
e99343e
 
 
 
 
 
99b21b2
e99343e
c313a26
e99343e
 
 
 
99b21b2
 
 
e99343e
 
 
 
99b21b2
e99343e
 
 
 
 
 
 
99b21b2
 
 
 
 
e99343e
 
 
99b21b2
e99343e
 
99b21b2
 
 
e99343e
 
99b21b2
e99343e
 
99b21b2
 
e99343e
 
99b21b2
 
e99343e
 
99b21b2
 
 
e99343e
 
 
 
99b21b2
 
e99343e
99b21b2
e99343e
 
99b21b2
e99343e
 
99b21b2
e99343e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
99b21b2
e99343e
 
 
 
 
 
99b21b2
e99343e
99b21b2
 
 
 
e99343e
 
 
99b21b2
e99343e
 
99b21b2
 
 
e99343e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
99b21b2
e99343e
 
 
 
 
 
 
 
 
99b21b2
e99343e
 
 
99b21b2
e99343e
 
 
 
 
 
 
99b21b2
e99343e
 
 
 
 
99b21b2
e99343e
99b21b2
 
 
 
 
 
e99343e
 
99b21b2
 
 
 
 
 
 
 
e99343e
99b21b2
 
 
 
e99343e
 
99b21b2
 
 
e99343e
99b21b2
 
e99343e
 
 
99b21b2
 
 
 
 
 
e99343e
99b21b2
 
 
 
 
 
 
 
 
 
e99343e
 
99b21b2
 
 
 
e99343e
 
 
 
 
99b21b2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e99343e
 
99b21b2
e99343e
99b21b2
 
 
 
 
 
e99343e
99b21b2
e99343e
 
99b21b2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e99343e
 
99b21b2
e99343e
99b21b2
e99343e
99b21b2
 
 
e99343e
 
99b21b2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e99343e
 
99b21b2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e99343e
99b21b2
e99343e
99b21b2
e99343e
99b21b2
 
 
 
e99343e
fbe81c0
99b21b2
 
 
 
 
 
 
 
 
 
 
e99343e
 
99b21b2
 
e99343e
efd4eb0
99b21b2
e99343e
 
99b21b2
 
 
 
 
 
 
 
 
 
 
 
e99343e
 
20cb3f9
e99343e
66aaed6
99b21b2
e99343e
99b21b2
e99343e
 
 
 
66aaed6
e99343e
 
99b21b2
 
 
 
e99343e
99b21b2
e99343e
99b21b2
 
 
 
e99343e
99b21b2
 
 
 
 
 
 
 
 
 
 
c313a26
 
 
99b21b2
 
e99343e
99b21b2
e99343e
 
99b21b2
c313a26
e99343e
 
c313a26
99b21b2
 
 
e99343e
 
b4edd10
e99343e
99b21b2
 
e99343e
 
 
 
99b21b2
 
e99343e
99b21b2
c313a26
99b21b2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c313a26
99b21b2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fbe81c0
99b21b2
 
 
 
 
fbe81c0
99b21b2
20cb3f9
99b21b2
c313a26
99b21b2
 
 
 
 
 
 
 
 
 
 
 
fbe81c0
 
99b21b2
 
 
 
fbe81c0
 
99b21b2
 
 
e99343e
99b21b2
 
 
604c55b
99b21b2
 
 
604c55b
99b21b2
 
 
c313a26
99b21b2
 
 
2b85b0c
20cb3f9
 
99b21b2
 
 
 
 
 
20cb3f9
99b21b2
 
 
 
 
 
 
 
 
 
 
 
20cb3f9
99b21b2
 
 
 
 
 
 
e99343e
99b21b2
20cb3f9
99b21b2
20cb3f9
99b21b2
 
20cb3f9
99b21b2
 
 
 
e99343e
20cb3f9
 
99b21b2
 
 
 
 
 
20cb3f9
99b21b2
 
 
 
 
 
 
 
20cb3f9
99b21b2
 
 
 
 
 
 
e99343e
99b21b2
20cb3f9
99b21b2
20cb3f9
99b21b2
 
20cb3f9
99b21b2
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
# 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 traceback
import asyncio

from fastapi import FastAPI, HTTPException, Response, Query, Depends, status
from fastapi.responses import FileResponse
from fastapi.exception_handlers import http_exception_handler
from starlette.exceptions import HTTPException as StarletteHTTPException
from langchain.prompts import PromptTemplate
from langchain_groq import ChatGroq
from pydantic import BaseModel, Field, BeforeValidator, model_serializer
from typing_extensions import Annotated
from pydantic_core import core_schema # Import core_schema for direct use in __get_pydantic_json_schema__

from 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"
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
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):
        # 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
        if not ObjectId.is_valid(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: 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
        return json_schema

# Annotated type for convenience in models
PyObjectId = Annotated[CustomObjectId, BeforeValidator(str)]


# ---------------------- Models ----------------------
class Contact(BaseModel):
    name: str
    last_name: str
    email: Optional[str] = None
    phone_number: Optional[str] = None

class Appointment(BaseModel):
    title: str
    description: str
    start_date: date
    start_time: Optional[str] = None
    end_date: Optional[date] = None
    end_time: Optional[str] = None

class Task(BaseModel):
    task_title: str
    task_description: str
    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]
    tasks: List[Task]
    original_email_text: str
    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

    # 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):
                    appt['start_date'] = appt['start_date'].isoformat()
                if isinstance(appt.get('end_date'), date) and appt.get('end_date') is not None:
                    appt['end_date'] = appt['end_date'].isoformat()
        if 'tasks' in data:
            for task_item in data['tasks']:
                if isinstance(task_item.get('due_date'), date):
                    task_item['due_date'] = task_item['due_date'].isoformat()
        return data

class ProcessEmailRequest(BaseModel):
    email_text: str = Field(..., example="Oggetto: Follow-up progetto “Delta”...")
    groq_api_key: str = Field(..., example="YOUR_GROQ_API_KEY")

class GenerateReplyRequest(BaseModel):
    email_text: str = Field(..., example="Oggetto: Follow-up progetto “Delta”...")
    groq_api_key: str = Field(..., example="YOUR_GROQ_API_KEY")
    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"])
    tone: str = Field("Friendly", examples=["Friendly", "Neutral", "Urgent", "Empathetic"])
    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
    language: str
    length: str
    style: str
    tone: str
    emoji: str
    generated_at: datetime = Field(default_factory=datetime.utcnow)

    class Config:
        populate_by_name = True
        arbitrary_types_allowed = True

    @model_serializer(when_used='json')
    def serialize_model(self):
        data = self.model_dump(by_alias=True, exclude_none=True)
        if "_id" in data and isinstance(data["_id"], ObjectId):
            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).")
    appointment_title: Optional[str] = Query(None, description="Filter by appointment title (case-insensitive partial match).")
    task_title: Optional[str] = Query(None, description="Filter by task title (case-insensitive partial match).")
    from_date: Optional[date] = Query(None, description="Filter by data processed on or after this date (YYYY-MM-DD).")
    to_date: Optional[date] = Query(None, description="Filter by data processed on or before this date (YYYY-MM-DD).")
    limit: int = Query(10, ge=1, le=100, description="Maximum number of results to return.")

class GeneratedReplyQuery(BaseModel):
    language: Optional[Literal["Italian", "English"]] = Query(None, description="Filter by reply language.")
    style: Optional[str] = Query(None, description="Filter by reply style (e.g., Professional, Casual).")
    tone: Optional[str] = Query(None, description="Filter by reply tone (e.g., Friendly, Neutral).")
    from_date: Optional[date] = Query(None, description="Filter by data generated on or after this date (YYYY-MM-DD).")
    to_date: Optional[date] = Query(None, description="Filter by data generated on or before this date (YYYY-MM-DD).")
    limit: int = Query(10, ge=1, le=100, description="Maximum number of results to return.")

# ---------------------- 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 YYYY-MM-DD format.
    Returns None if input is None or cannot be parsed into a valid 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)
    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

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 ""
        last_name = " ".join(parts[1:]) if len(parts) > 1 else ""
        return name, last_name

    contacts_data = []
    for c in data.get("contacts", []):
        name_val, last_name_val = split_name(c.get("name", ""))
        contacts_data.append(Contact(name=name_val, last_name=last_name_val, email=c.get("email"), phone_number=c.get("phone_number")))

    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
        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"),
            start_date=start_date_val, start_time=a.get("start_time"),
            end_date=end_date_val, end_time=a.get("end_time")
        ))

    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
        tasks_data.append(Task(
            task_title=t.get("task_title", "Untitled"), task_description=t.get("task_description", "No description"),
            due_date=due_date_val
        ))
    return ExtractedData(contacts=contacts_data, appointments=appointments_data, tasks=tasks_data, original_email_text=original_email_text)

# ---------------------- 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.")

    llm = ChatGroq(model="meta-llama/llama-4-scout-17b-16e-instruct", temperature=0, max_tokens=2000, groq_api_key=api_key)

    prompt_today_str = current_date.isoformat()
    prompt_tomorrow_str = (current_date + timedelta(days=1)).isoformat()

    prompt_template_str = f"""
You are an expert email assistant tasked with extracting structured information from an Italian email.

**Your response MUST be a single, complete JSON object, wrapped in a ```json``` block.**
**DO NOT include any conversational text, explanations, or preambles outside the JSON block.**
**The JSON should contain three top-level keys: "contacts", "appointments", and "tasks".**
If a category has no items, its list should be empty (e.g., "contacts": []).

Here is the required JSON schema for each category:

- **contacts**: List of Contact objects.
    Each Contact object must have:
    - `name` (string, full name)
    - `last_name` (string, last name) - You should infer this from the full name.
    - `email` (string, optional, null if not present)
    - `phone_number` (string, optional, null if not present)

- **appointments**: List of Appointment objects.
    Each Appointment object must have:
    - `title` (string, short, meaningful title in Italian based on the meeting's purpose)
    - `description` (string, summary of the meeting's goal)
    - `start_date` (string, YYYY-MM-DD. If not explicitly mentioned, use "{prompt_today_str}" for "today", or "{prompt_tomorrow_str}" for "tomorrow")
    - `start_time` (string, optional, e.g., "10:30 AM", null if not present)
    - `end_date` (string, YYYY-MM-DD, optional, null if unknown or not applicable)
    - `end_time` (string, optional, e.g., "11:00 AM", null if not present)

- **tasks**: List of Task objects.
    Each Task object must have:
    - `task_title` (string, short summary of action item)
    - `task_description` (string, more detailed explanation)
    - `due_date` (string, YYYY-MM-DD. Infer from context, e.g., "entro domani" becomes "{prompt_tomorrow_str}", "today" becomes "{prompt_today_str}")

---

Email:
{{email}}
"""
    prompt_template = PromptTemplate(input_variables=["email", "prompt_today_str", "prompt_tomorrow_str"], template=prompt_template_str)
    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_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}")
        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:
        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()
        raise Exception(f"An error occurred during email processing: {e}")

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."

    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


# --- 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, and caching.",
    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):
    """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"
    )


# --- FastAPI Event Handlers for MongoDB ---
@app.on_event("startup")
async def startup_event():
    global client, db, extracted_emails_collection, generated_replies_collection
    print(f"[{datetime.now()}] FastAPI app startup sequence initiated.")
    try:
        # Connect to MongoDB
        client = MongoClient(MONGO_URI, serverSelectionTimeoutMS=5000)
        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"[{datetime.now()}] Successfully connected to MongoDB: {DB_NAME}")

    except (ConnectionFailure, OperationFailure) as e:
        print(f"[{datetime.now()}] 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 startup: {e}")
        traceback.print_exc()
        client = None
        db = None
        extracted_emails_collection = None
        generated_replies_collection = None
    finally:
        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}")
                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.")
            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 initialization.")


@app.on_event("shutdown")
async def shutdown_event():
    global client
    print(f"[{datetime.now()}] FastAPI app shutting down.")
    if client:
        client.close()
        print(f"[{datetime.now()}] MongoDB client closed.")


# --- API Endpoints ---
@app.get("/health", summary="Health Check")
async def health_check():
    """
    Checks the health of the API, including MongoDB connection.
    """
    db_status = "MongoDB not connected."
    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_status = "MongoDB connection OK."
            db_ok = True
        except Exception as e:
            db_status = f"MongoDB connection error: {e}"
            db_ok = False

    if db_ok:
        return {"status": "ok", "message": "Email Assistant API is up.", "database": db_status}
    else:
        raise HTTPException(
            status_code=503,
            detail={"message": "Service unavailable.", "database": db_status}
        )


@app.post("/generate-reply", response_model=GenerateReplyResponse, summary="Generate a smart reply to an email")
async def generate_email_reply(request: GenerateReplyRequest):
    """
    Generates a smart reply to the provided email text using an LLM.
    The generated reply is also stored in MongoDB for caching and historical purposes.
    """
    if generated_replies_collection is None:
        raise HTTPException(status_code=503, detail="MongoDB not available for generated_replies.")

    try:
        # Check cache first
        cache_query = {
            "original_email_text": request.email_text,
            "language": request.language,
            "length": request.length,
            "style": request.style,
            "tone": request.tone,
            "emoji": request.emoji,
        }
        print(f"[{datetime.now()}] /generate-reply: Checking cache for reply...")
        # 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:
            print(f"[{datetime.now()}] /generate-reply: Reply found in cache. ID: {str(cached_reply_doc['_id'])}")
            return GenerateReplyResponse(
                reply=cached_reply_doc["generated_reply_text"],
                stored_id=str(cached_reply_doc["_id"]),
                cached=True
            )

        # If not in cache, directly call the internal LLM function
        print(f"[{datetime.now()}] /generate-reply: Reply not in cache. Calling LLM for generation...")
        reply_content = await asyncio.to_thread(
            _generate_response_internal,
            request.email_text,
            request.groq_api_key,
            request.language,
            request.length,
            request.style,
            request.tone,
            request.emoji
        )
        print(f"[{datetime.now()}] /generate-reply: LLM call completed. Storing newly generated reply in MongoDB.")

        # Prepare data for storage
        reply_data_to_store = GeneratedReplyData(
            original_email_text=request.email_text,
            generated_reply_text=reply_content,
            language=request.language,
            length=request.length,
            style=request.style,
            tone=request.tone,
            emoji=request.emoji
        )
        # Use model_dump for Pydantic v2. Exclude 'id' as it's generated by MongoDB.
        reply_data_dict = reply_data_to_store.model_dump(by_alias=True, exclude_none=True, exclude={'id'})

        # Insert into MongoDB
        insert_result = await asyncio.to_thread(generated_replies_collection.insert_one, reply_data_dict)
        stored_id = str(insert_result.inserted_id) # Convert ObjectId to string for the response

        print(f"[{datetime.now()}] /generate-reply: Reply stored in MongoDB. ID: {stored_id}")

        # Return the response as per GenerateReplyResponse model
        return GenerateReplyResponse(
            reply=reply_content,
            stored_id=stored_id,
            cached=False # Always False since we just generated it
        )
    except Exception as e:
        traceback.print_exc()
        # Ensure consistent error response
        raise HTTPException(status_code=500, detail=f"Error generating or storing reply: {str(e)}")

@app.post("/extract-data", response_model=ExtractedData, summary="Extract structured data from an email")
async def extract_email_data(request: ProcessEmailRequest):
    """
    Extracts contacts, appointments, and tasks from the provided email text.
    """
    if extracted_emails_collection is None:
        raise HTTPException(status_code=503, detail="MongoDB not available.")

    current_date = date.today() # Get current date for context

    print(f"[{datetime.now()}] /extract-data: Received request.")
    try:
        print(f"[{datetime.now()}] /extract-data: Calling internal processing function.")
        # Run blocking LLM call in a thread pool
        extracted_data = await asyncio.to_thread(_process_email_internal, request.email_text, request.groq_api_key, current_date)

        print(f"[{datetime.now()}] /extract-data: Internal processing complete. Preparing for DB insert.")
        # Convert Pydantic model to dictionary for MongoDB insert, handling _id alias
        # Use model_dump for Pydantic v2
        data_to_insert = extracted_data.model_dump(by_alias=True, exclude_none=True, exclude={'id'})

        # --- NEW CONVERSION FOR MONGODB ---
        # MongoDB's BSON doesn't natively support Python's datetime.date type.
        # It expects datetime.datetime. Convert all date fields to datetime.datetime.
        if 'appointments' in data_to_insert:
            for appt in data_to_insert['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 data_to_insert:
            for task_item in data_to_insert['tasks']:
                if isinstance(task_item.get('due_date'), date):
                    task_item['due_date'] = datetime.combine(task_item['due_date'], datetime.min.time())
        # --- END NEW CONVERSION ---

        print(f"[{datetime.now()}] /extract-data: Inserting into MongoDB... Data: {data_to_insert}") # Add data logging
        # Use asyncio.to_thread for blocking MongoDB insert operation
        insert_result = await asyncio.to_thread(extracted_emails_collection.insert_one, data_to_insert)

        # Update the extracted_data object with the MongoDB-generated ID
        extracted_data.id = str(insert_result.inserted_id)
        print(f"[{datetime.now()}] /extract-data: Data inserted into MongoDB. ID: {extracted_data.id}")

        return extracted_data
    except ValueError as ve:
        raise HTTPException(status_code=400, detail=str(ve))
    except Exception as e:
        traceback.print_exc() # Print full traceback for debugging
        raise HTTPException(status_code=500, detail=f"An internal server error occurred: {e}")


@app.get("/query-extracted-emails", response_model=List[ExtractedData], summary="Query stored extracted email data")
async def query_extracted_emails(query_params: ExtractedEmailQuery = Depends()):
    """
    Queries extracted email data from MongoDB based on various filters.
    """
    if extracted_emails_collection is None:
        raise HTTPException(status_code=503, detail="MongoDB not available.")

    mongo_query = {}
    if query_params.contact_name:
        # Case-insensitive partial match on contact name or last name
        mongo_query["$or"] = [
            {"contacts.name": {"$regex": query_params.contact_name, "$options": "i"}},
            {"contacts.last_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"}

    # Date range filtering for processed_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["$lte"] = datetime.combine(query_params.to_date, datetime.max.time())
    if date_query:
        mongo_query["processed_at"] = date_query

    try:
        # Use asyncio.to_thread for blocking MongoDB find operation
        cursor = await asyncio.to_thread(extracted_emails_collection.find, mongo_query)
        # Use to_list to limit results and convert to list
        results = await asyncio.to_thread(lambda: list(cursor.limit(query_params.limit)))

        # Convert MongoDB documents to ExtractedData Pydantic models
        return [ExtractedData(**doc) for doc in results]
    except Exception as e:
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Error querying extracted emails: {e}")


@app.get("/query-generated-replies", response_model=List[GeneratedReplyData], summary="Query stored generated replies")
async def query_generated_replies(query_params: GeneratedReplyQuery = Depends()):
    """
    Queries generated email replies from MongoDB based on various filters.
    """
    if generated_replies_collection is None:
        raise HTTPException(status_code=503, detail="MongoDB not available.")

    mongo_query = {}
    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

    # Date range filtering for 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["$lte"] = datetime.combine(query_params.to_date, datetime.max.time())
    if date_query:
        mongo_query["generated_at"] = date_query

    try:
        # Use asyncio.to_thread for blocking MongoDB find operation
        cursor = await asyncio.to_thread(generated_replies_collection.find, mongo_query)
        # Use to_list to limit results and convert to list
        results = await asyncio.to_thread(lambda: list(cursor.limit(query_params.limit)))

        # Convert MongoDB documents to GeneratedReplyData Pydantic models
        return [GeneratedReplyData(**doc) for doc in results]
    except Exception as e:
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Error querying generated replies: {e}")