328 lines
12 KiB
Python
328 lines
12 KiB
Python
"""
|
|
|
|
AUTHOR:
|
|
|
|
Khushal P Soonderji
|
|
|
|
DATE:
|
|
|
|
Wednesday, 27th Nov., 2024
|
|
|
|
OBJECTIVE:
|
|
|
|
To define the interaction between the UI layer and the database connectivity in one place. Here we shall handle
|
|
all the activities for OAuth2.0 authorization requests for all the users of our service.
|
|
|
|
REFERENCES:
|
|
|
|
N/A
|
|
|
|
DOWNLOADS:
|
|
|
|
N/A
|
|
|
|
"""
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** IMPORT ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# To make sibling directories accessible for imports:
|
|
import sys
|
|
|
|
from langchain.chains.summarize.stuff_prompt import prompt_template
|
|
from sqlalchemy.orm.collections import collection
|
|
|
|
sys.path.append(".")
|
|
sys.path.append("..")
|
|
|
|
# My async utils:
|
|
from utils_v2.string import json
|
|
from utils_v2.date_time import date_time
|
|
from utils_v2.database.async_mysql_v2 import AsyncMySQL
|
|
from utils_v2.database.async_mongo_v2 import AsyncMongo
|
|
|
|
# Mail Clients:
|
|
from utils_v2.goog.gmail.gmail_client import AsyncGMailClient
|
|
from utils_v2.goog.models.data.auth_tokens import GoogleAuthTokens
|
|
|
|
# Base model:
|
|
from models.behaviour.base import BaseModel
|
|
|
|
# To work with MongoDB:
|
|
from bson import ObjectId
|
|
from pymongo import InsertOne, UpdateOne
|
|
|
|
# To work with LLMs:
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain_core.prompts import ChatPromptTemplate
|
|
|
|
# To work with datatypes:
|
|
from typing import Literal
|
|
|
|
# To make deep-copies:
|
|
import copy
|
|
|
|
# To work with date and time:
|
|
import datetime
|
|
|
|
# For asynchronous activities:
|
|
import asyncio
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** MACROS / ONE-TIME INIT ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** VARIABLES ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** FUNCTIONS ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** CLASSES ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
class MailSyncModel(BaseModel):
|
|
|
|
# For MongoDB:
|
|
AUTH_COLLECTION = "_authTokens"
|
|
MAIL_COLLECTION = "_messages"
|
|
|
|
# For AI Magic through LLMs:
|
|
prompt_template = ChatPromptTemplate.from_messages([
|
|
(
|
|
"system",
|
|
"You're a mail summary expert that summarizes mails in 150 chars or less. HIDE SENSITIVE INFO (LIKE OTPs) FROM THE SUMMARY."
|
|
),
|
|
(
|
|
"user",
|
|
"Please summarize this mail: \"\"\"{mail}\"\"\""
|
|
)
|
|
])
|
|
|
|
async def __sync_one(
|
|
self,
|
|
mongo_conn: AsyncMongo,
|
|
user_info: dict,
|
|
mail_client: AsyncGMailClient,
|
|
tokens: GoogleAuthTokens,
|
|
message_id: str,
|
|
llm: ChatOpenAI = None,
|
|
force_sync: bool = False
|
|
) -> UpdateOne | None:
|
|
|
|
# ┏┓ ┓┏ • ┓ ┓
|
|
# ┃┃┏┓┏┓┏┓┏┓┏┓┏┓ ┃┃┏┓┏┓┓┏┓┣┓┃┏┓┏
|
|
# ┣┛┛ ┗ ┣┛┗┻┛ ┗ ┗┛┗┻┛ ┗┗┻┗┛┗┗ ┛
|
|
# ┛
|
|
|
|
mail_payload = None
|
|
mail_client_name = None
|
|
|
|
# ┏┓┓ ┓ ┏┓ • • ┳┓ ┓
|
|
# ┃ ┣┓┏┓┏┃┏ ┣ ┓┏┓┏╋┓┏┓┏┓ ┣┫┏┓┏┏┓┏┓┏┫┏
|
|
# ┗┛┛┗┗ ┗┛┗ ┗┛┛┗┗┛┗┗┛┗┗┫ ┛┗┗ ┗┗┛┛ ┗┻┛
|
|
# ┛
|
|
|
|
# Check if you already have that mail in your database:
|
|
existing_record = await mongo_conn.find_one(
|
|
collection = self.MAIL_COLLECTION,
|
|
filter = {
|
|
"messageType": "email",
|
|
"$or": [
|
|
{"payload.messageId": message_id}
|
|
]
|
|
},
|
|
projection = {"_id": True}
|
|
)
|
|
|
|
# If there already exists such a record, and we haven't been forced to re-sync it:
|
|
if existing_record and not force_sync: return mail_payload
|
|
|
|
# ┏┓┳┳┓ •┓
|
|
# ┃┓┃┃┃┏┓┓┃
|
|
# ┗┛┛ ┗┗┻┗┗
|
|
|
|
if isinstance(mail_client, AsyncGMailClient):
|
|
|
|
# Note down the name of the mail client:
|
|
mail_client_name = "gmail"
|
|
|
|
# Fetch the formatted mail message:
|
|
client_response = await mail_client.get_message(
|
|
tokens = tokens,
|
|
message_id = message_id,
|
|
return_raw = False
|
|
)
|
|
|
|
# If the fetch was successful:
|
|
if client_response.success:
|
|
|
|
# Summarize the content:
|
|
if llm:
|
|
prompt = self.prompt_template.invoke({"mail": client_response.data["unformattedText"]})
|
|
llm_response = await llm.ainvoke(prompt)
|
|
client_response.data["aiSnippet"] = llm_response.content
|
|
|
|
# Note down the response:
|
|
mail_payload = client_response.data
|
|
|
|
# ┳┓
|
|
# ┣┫┏┓┏┏┓┏┓┏┓┏┏┓
|
|
# ┛┗┗ ┛┣┛┗┛┛┗┛┗
|
|
# ┛
|
|
|
|
if mail_payload:
|
|
return UpdateOne(
|
|
filter = {
|
|
"messageType": "email",
|
|
"$or": [
|
|
{"payload.messageId": message_id}
|
|
]
|
|
},
|
|
update = {
|
|
"$set": {
|
|
"readTs": date_time.get_current_utc_date_time(),
|
|
"user": user_info,
|
|
"messageType": "email",
|
|
"connector": mail_client_name,
|
|
"payload": mail_payload
|
|
}
|
|
},
|
|
upsert = True
|
|
)
|
|
|
|
# Done here:
|
|
return mail_payload
|
|
|
|
async def sync(
|
|
self,
|
|
mongo_conn: AsyncMongo,
|
|
user_info: dict,
|
|
mail_client: AsyncGMailClient,
|
|
tokens: GoogleAuthTokens,
|
|
llm: ChatOpenAI = None,
|
|
force_sync: bool = False,
|
|
start_date: datetime.datetime = None,
|
|
end_date: datetime.datetime = None,
|
|
max_count: int = 100
|
|
) -> int:
|
|
|
|
# Start by assuming failure:
|
|
mails_count = 0
|
|
|
|
# ┏┓┳┳┓ •┓
|
|
# ┃┓┃┃┃┏┓┓┃
|
|
# ┗┛┛ ┗┗┻┗┗
|
|
|
|
if isinstance(mail_client, AsyncGMailClient):
|
|
|
|
# Enlist all the labels, we need to find the label that indicates that we've read the mail:
|
|
client_response = await mail_client.list_labels(tokens = tokens)
|
|
if not client_response.success: return mails_count
|
|
labels = client_response.data
|
|
custom_label = "Sync'd with TheCAOffice"
|
|
custom_label_id = labels.get(custom_label)
|
|
if custom_label_id is None:
|
|
client_response = await mail_client.create_label(
|
|
tokens = tokens,
|
|
label_name = custom_label,
|
|
label_visibility = "labelHide"
|
|
)
|
|
if not client_response.success: return mails_count
|
|
custom_label_id = client_response.data["id"]
|
|
|
|
# Build the query:
|
|
sub_queries = [f"-label:\"{custom_label}\""]
|
|
if start_date: sub_queries.append(start_date.strftime("after:%Y/%m/%d"))
|
|
if end_date: sub_queries.append(end_date.strftime("before:%Y/%m/%d"))
|
|
print("Q:", " ".join(sub_queries))
|
|
|
|
# Get a list of all the mails:
|
|
client_response = await mail_client.list_messages(
|
|
tokens = tokens,
|
|
max_count = max_count,
|
|
# query = " ".join(sub_queries)
|
|
)
|
|
if not client_response.success: return mails_count
|
|
messages_list = client_response.data["messages"]
|
|
print(messages_list)
|
|
|
|
# Create MongoDB operations for all the mails:
|
|
tasks = [
|
|
self.__sync_one(
|
|
mongo_conn = mongo_conn,
|
|
user_info = user_info,
|
|
mail_client = mail_client,
|
|
tokens = tokens,
|
|
message_id = v["id"],
|
|
llm = llm,
|
|
force_sync = force_sync
|
|
)
|
|
for k, v in messages_list.items()
|
|
]
|
|
mongo_operations = await asyncio.gather(*tasks)
|
|
mongo_operations = [mo for mo in mongo_operations if mo is not None]
|
|
|
|
# Write the mails to MongoDB:
|
|
mails_count = await mongo_conn.bulk_write(
|
|
collection = self.MAIL_COLLECTION,
|
|
requests = mongo_operations
|
|
)
|
|
print("MAILS COUNT:", mails_count)
|
|
|
|
# If all the mails were sync'd properly:
|
|
client_response = await mail_client.modify_messages(
|
|
tokens = tokens,
|
|
message_ids = [v["id"] for k, v in messages_list.items()],
|
|
add_label_ids = [custom_label_id]
|
|
)
|
|
|
|
# ┳┓
|
|
# ┣┫┏┓┏┏┓┏┓┏┓┏┏┓
|
|
# ┛┗┗ ┛┣┛┗┛┛┗┛┗
|
|
# ┛
|
|
|
|
# Done here:
|
|
return mails_count
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** MAIN PROGRAM ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
pass
|