""" AUTHOR: Khushal P Soonderji DATE: tuesday, 3rd Dec., 2024 OBJECTIVE: From here we sync all mails between the mail client's server and TheCAOffice's database. REFERENCES: N/A DOWNLOADS: N/A """ # ***************************************************************************************************************** # ***** **** # *** IMPORT *** # ***** **** # ***************************************************************************************************************** # To make sibling directories accessible for imports: import sys from logging import exception sys.path.append(".") sys.path.append("..") # For Quart: from quart import current_app # 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, AsyncMongoStorage # 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 # Data models: from models.data.mail.sync import MailSyncOneResult, MailSyncManyResults # To work with MongoDB: from bson import ObjectId from pymongo import InsertOne, UpdateOne, ReplaceOne # To work with LLMs: from langchain_openai import ChatOpenAI from langchain_core.prompts import ChatPromptTemplate # To work with datatypes: from typing import Literal, List, Dict, Any # To make deep-copies: import copy # To work with base-64 encoding: import base64 # 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}\"\"\"" ) ]) @staticmethod async def __save_one_attachment( session_token: str, attachment: Dict[str, Any], attachment_tags: List[str], attachment_metadata: dict, retry_count: int = 1, retry_delay: int = 1, backoff_multiplier: float = 1.1 ) -> Dict[str, Any]: """ Saves one attachment and generates a URL that can be later used to retrieve it. :param session_token: The session token of the uer who is trying to upload this file. :param attachment: The JSON that describes the attachment. :param attachment_tags: Any tags to put on the file for easy search later. :param attachment_metadata: Any metadata to put on the file for easy search later. :param retry_count: How many max. retries to do in case of failure. :param retry_delay: The interval between the delays. :param backoff_multiplier: By what rate the delay between 2 attempts must change. :return: The JSON that describes the same attachment, except that the payload's data is replaced by the id and url of where to find the attachment. """ # Make a deep-copy of the attachment JSON, # and process the payload in advance: attachment_copy = copy.deepcopy(attachment) attachment_payload = attachment_copy.pop("payload").encode() if attachment_copy.pop("contentTransferEncoding", "?").strip().lower() == "base64": attachment_payload = base64.b64decode(attachment_payload) # Start by assuming failure, # and retry as many times as asked: attachment_copy["id"] = None attachment_copy["url"] = None for _ in range(retry_count): # Make the upload: api_response = await current_app.http_client.post( url = current_app.script_data["fileUpload"]["url"], headers = { "X-Session-Token": session_token, "X-File-Name": attachment["filename"], "X-File-Private": "false", "X-File-Tags": json.to_string(attachment_tags, no_space = True), "X-File-Metadata": json.to_string(attachment_metadata, no_space = True) }, data = attachment_payload ) # If the upload was successful: if api_response.status_code in [200]: api_data = api_response.json()["data"] attachment_copy["id"] = api_data["id"] attachment_copy["url"] = api_data["url"] break # Done here: return attachment_copy async def __save_many_attachments( self, session_token: str, attachments: List[Dict[str, Any]], attachment_tags: List[str], attachment_metadata: dict, retry_count: int = 1, retry_delay:int = 1, backoff_multiplier: float = 1.1 ) -> List[Dict[str, Any]]: """ Saves all the attachments received in the mail (whether inline or otherwise) and makes them available through simple download URLs. :param session_token: The session token of the uer who is trying to upload this file. :param attachments: The JSON that describes the attachments. :param attachment_tags: Any tags to put on the file for easy search later. :param attachment_metadata: Any metadata to put on the file for easy search later. :param retry_count: How many max. retries to do in case of failure. :param retry_delay: The interval between the delays. :param backoff_multiplier: By what rate the delay between 2 attempts must change. :return: The JSON that describes the same attachments, except that the payload's data is replaced by the id and url of where to find each attachment. """ # Create and fire all the tasks # needed to save the files: tasks = [ self.__save_one_attachment( session_token = session_token, attachment = attachment, attachment_tags = attachment_tags, attachment_metadata = attachment_metadata, retry_count = retry_count, retry_delay = retry_delay, backoff_multiplier = backoff_multiplier ) for attachment in attachments ] uploaded_attachments = await asyncio.gather(*tasks) # Done here: return uploaded_attachments async def __sync_one_gmail( self, session_token: str, mongo_conn: AsyncMongo, mail_client: AsyncGMailClient, tokens: GoogleAuthTokens, message_id: str, llm: ChatOpenAI = None, force_sync: bool = False ) -> MailSyncOneResult: """ Sync on mail from GMail. :param session_token: The session token of the uer who is trying to upload this file. :param mongo_conn: The instance of the connection to the database to use. :param mail_client: The instance of the mail client to use to perform the action. :param tokens: The tokens to use to fetch the mails. :param message_id: The id that Google uses to identify this mail. This will be received in the 'list_messages' method. :param llm: The instance of the LLM to use to summarize the mail's content. :param force_sync: Whether you would like to forcefully re-sync the mail even if it is already present in the database. :return: """ # Start by assuming failure: sync_result = MailSyncOneResult() # If we've not been forced to re-sync the mail message, # we first check if the mail already exists in our database: if not force_sync: mail_record = await mongo_conn.find_one( collection = self.MAIL_COLLECTION, filter = mongo_conn.dict_to_dot_notation({ "payload": { "messageId": message_id } }), projection = { "_id": False, "readTs": "payload.readTs" } ) if mail_record: sync_result.success = True sync_result.message = f"gmail message '{message_id}' already sync'd on '{mail_record['readTs']} (UTC)'" return sync_result # Now that we know that we have to fetch the mail from GMail: client_response = await mail_client.get_message( tokens = tokens, message_id = message_id, return_raw = False ) if not client_response.success: sync_result.message = f"gmail (messageId: '{message_id}'): {client_response.message}" # We upload the attachments: client_response.data["attachments"] = await self.__save_many_attachments( session_token = session_token, attachments = client_response.data["attachments"], attachment_tags = [ "email", "gmail", client_response.data["from"][0]["name"], client_response.data["from"][0]["email"], tokens.email, ], attachment_metadata = { "project": "tcaoff", "serviceType": "email", "client": "gmail", "from": client_response.data["from"][0]["email"], "to": tokens.email }, retry_count = 3 ) # Give a quick indicator of whether this mail is an inbox mail or sent mail: all_recipients = [] for field in ["to", "cc", "bcc"]: all_recipients += [item["email"] for item in client_response.data[field]] if tokens.email in all_recipients: client_response.data["isInbox"] = True else: client_response.data["isInbox"] = False # If an LLM is given, we add an AI summary: llm_json = None if llm: llm_response = response = await llm.ainvoke( self.prompt_template.invoke({ "mail": client_response.data["unformattedText"] }) ) llm_json = { "snippet": llm_response.content, "usage": llm_response.usage_metadata } client_response.data["aiSnippet"] = llm_json # Done here: sync_result.success = True sync_result.mailMessage = client_response.data return sync_result async def __sync_many_gmail( self, session_token: str, mongo_conn: AsyncMongo, account_identifier: ObjectId, 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 ) -> MailSyncManyResults: """ Sync many mails from GMail in one shot. :param session_token: The session token of the uer who is trying to upload this file. :param mongo_conn: The instance of the connection to the database to use. :param account_identifier: The id of the document in the database that holds the tokens to access the account. Needed only for refreshing the tokens and saving them. :param mail_client: The instance of the mail client to use to perform the action. :param tokens: The tokens to use to fetch the mails. :param llm: The instance of the LLM to use to summarize the mail's content. :param force_sync: Whether you would like to forcefully re-sync the mail even if it is already present in the database. :param start_date: The starting date (inclusive) from when to sync the mails. :param end_date: The ending date (inclusive) from when to sync the mails. :param max_count: The max. no. of mails to sync. :return: The result of the sync'ing. """ # Start by assuming failure: sync_results = MailSyncManyResults() # Refresh the tokens (if needed): tokens_refreshed = await tokens.arefresh( http_client = current_app.http_client, client_id = mail_client.client_id, client_secret = mail_client.client_secret ) if tokens_refreshed: await current_app.mail_oauth_model.set_token( db_conn = current_app.sql_writer, mongo_conn = mongo_conn, account_identifier = account_identifier, email_id = tokens.email, token = tokens, session_token = session_token ) # Let's build the query: sub_queries = [] if start_date: sub_queries.append(start_date.strftime("after:%Y/%m/%d")) if end_date: sub_queries.append((end_date + datetime.timedelta(days = 1)).strftime("before:%Y/%m/%d")) query_string = " ".join(sub_queries) # Let's enlist all the mails that fall in the date range: client_response = await mail_client.list_messages( tokens = tokens, max_count = max_count, query = query_string ) if not client_response.success: sync_results["message"] = f"gmail: {client_response.message}" return sync_results messages_list = client_response.data["messages"] # Now, for every mail in the list, we fetch the mail and note the results: tasks = [ self.__sync_one_gmail( session_token = session_token, mongo_conn = mongo_conn, mail_client = mail_client, tokens = tokens, message_id = v["id"], llm = llm, force_sync = force_sync ) for v in messages_list.values() ] individual_sync_results = await asyncio.gather(*tasks) # Now we create operations for each mail, # and maintain success/failure counters: sync_results.totalCount = len(individual_sync_results) mongo_operations = [] for result in individual_sync_results: if result.success: sync_results.successCount += 1 else: sync_results.failureCount += 1 if result.mailMessage: mongo_operations.append(ReplaceOne( filter = { "serviceType": "email", "$or": [ { "client": "gmail", "payload.messageId": result.mailMessage["messageId"] } ] }, replacement = { "version": "1.0.0", "accountId": ObjectId(account_identifier), "serviceType": "email", "client": "gmail", "payload": result.mailMessage }, upsert = True )) # Make the bulk write: if mongo_operations: mongo_count = await mongo_conn.bulk_write( collection = self.MAIL_COLLECTION, requests = mongo_operations ) # Apply the labels to the read messages: try: client_response = await mail_client.modify_messages( tokens = tokens, message_ids = [v["id"] for v in messages_list.values()], add_label_ids = [tokens.labels.get("TCAOFF", {}).get("id")] ) print(client_response) print([tokens.labels.get("TCAOFF")]) except Exception as exception: self._printer(exception) # Done here: sync_results.message = f"{sync_results.successCount}/{sync_results.totalCount} mail(s) sync'd from gmail" return sync_results async def sync( self, session_token: str, mongo_conn: AsyncMongo, account_identifier: ObjectId, llm: ChatOpenAI = None, force_sync: bool = False, start_date: datetime.datetime = None, end_date: datetime.datetime = None, max_count: int = 100 ) -> MailSyncManyResults: """ Sync many mails at once from many types of clients. Use this as a common entry point after which you internally route the request to the appropriate clients. :param session_token: The session token of the uer who is trying to upload this file. :param mongo_conn: The instance of the connection to the database to use. :param account_identifier: The id of the document in the database that holds the tokens to access the account. Needed only for refreshing the tokens and saving them. :param llm: The instance of the LLM to use to summarize the mail's content. :param force_sync: Whether you would like to forcefully re-sync the mail even if it is already present in the database. :param start_date: The starting date (inclusive) from when to sync the mails. :param end_date: The ending date (inclusive) from when to sync the mails. :param max_count: The max. no. of mails to sync. :return: The result of the sync'ing. """ # Start by assuming failure: sync_results = MailSyncManyResults() # ┏┓ ┓ ┏┳┓ ┓ # ┣ ┏┓╋┏┣┓ ┃ ┏┓┃┏┏┓┏┓┏ # ┻ ┗ ┗┗┛┗ ┻ ┗┛┛┗┗ ┛┗┛ # We first load the authorization tokens: auth_json = await current_app.mail_oauth_model.get_token( mongo_conn = mongo_conn, account_identifier = account_identifier, ) # If we failed to load the authorization tokens: if not auth_json: sync_results.message = f"no such account identifier '{account_identifier}'" return sync_results # ┏┓ ┏┓┳┳┓ •┓ # ┣ ┏┓┏┓ ┃┓┃┃┃┏┓┓┃ # ┻ ┗┛┛ ┗┛┛ ┗┗┻┗┗ if auth_json["client"] == "gmail": return await self.__sync_many_gmail( session_token = session_token, mongo_conn = mongo_conn, account_identifier = account_identifier, mail_client = current_app.gmail_client, tokens = GoogleAuthTokens(**auth_json["token"]), llm = llm, force_sync = force_sync, start_date = start_date, end_date = end_date, max_count = max_count ) # ┳ ┓• ┓ ┏┓┓• # ┃┏┓┓┏┏┓┃┓┏┫ ┃ ┃┓┏┓┏┓╋ # ┻┛┗┗┛┗┻┗┗┗┻ ┗┛┗┗┗ ┛┗┗ # If we haven't been able to sync mail due to not entering any 'if' condition: sync_results.message = f"no such mail client '{auth_json['client']}'" return sync_results # ***************************************************************************************************************** # ***** **** # *** MAIN PROGRAM *** # ***** **** # ***************************************************************************************************************** if __name__ == "__main__": pass