""" AUTHOR: Khushal P Soonderji DATE: Thursday, 12th Dec., 2024 OBJECTIVE: To handle all auth-tokens from one place. REFERENCES: N/A DOWNLOADS: N/A """ # ***************************************************************************************************************** # ***** **** # *** IMPORT *** # ***** **** # ***************************************************************************************************************** # To make sibling directories accessible for imports: import sys 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 # Base model: from controllers.base import BaseModel # Data models: from models.core.user import CoreUserInfoModel from models.core.auth_token import CoreAuthTokenModel from models.core.message import CoreMessageModel from models.api.mail.sync import MailSyncOneResult, MailSyncManyResults, MailSendOneResult from models.api.mail.send import MailSendRequestData # Mail Clients: from utils_v2.goog.controllers.gmail.gmail_client import AsyncGMailClient from utils_v2.goog.models.auth_tokens import GoogleAuthTokens # To work with MongoDB: from bson import ObjectId from pymongo import InsertOne, UpdateOne, ReplaceOne # To work with LLMs: from controllers.core.ai.llm import CoreLLMController from models.core.ai.llm import LLMInput, LLMOutput # To work with datatypes: from typing import Literal, List, Dict, Any # To parse the HTML content in the mail: from bs4 import BeautifulSoup # 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 MailController: # ┏┓┓ ┓┏ # ┃ ┃┏┓┏┏ ┃┃┏┓┏┓┏ # ┗┛┗┗┻┛┛ ┗┛┗┻┛ ┛ # For AI Magic through LLMs: RECEIVED_MAIL_PROMPT_TEMPLATE = [ { "role": "system", "content": ( "You're an expert mail summary program. " "Provide the response in a structured JSON format with two fields: \"summary\" and \"senderType\". " "The summary should be 150 chars or less. " "If available, show login info like username and OTPs in your summary." "If no login info is provided, please don't worry; just summarize what you see." "\"senderType\" must be one of [\"Vendor\", \"Client\", null].\n" "Example output 1: " "{\"summary\":\"Sagar Supplies has shipped your materials. They are expected to reach by Thursday. " "Use OTP 346780 when the delivery agent asks.\",\"senderType\":\"Vendor\"}\n" "Example output 2: " "{\"summary\":\"Mr. Mehta is enquiring about the submission of his tax filings, which are to be done " "today.\",\"senderType\":\"Client\"}\n" "Example output 3: " "{\"summary\":\"JustDial's marketing message. They're offering a 35% discount to new accounts.\"," "\"senderType\":null}\n" "Remember to respond only with the raw JSON string, nothing else. Good luck :)" ) } ] SENT_MAIL_PROMPT_TEMPLATE = [ { "role": "system", "content": ( "You're an expert mail summary assistant. That summarizes sent mails in 150 chars or less. " "The objective of your user is to be able to recollect what a mail they sent was about from a brief " "summary. Reply in a simple string, no formatting is allowed except emojis. Good luck :)" ) } ] # ┓┏ ┓ # ┣┫┏┓┃┏┓┏┓┏┓┏ # ┛┗┗ ┗┣┛┗ ┛ ┛ # ┛ def extract_plaintext_parts( self, payload: dict ) -> List[str]: # Start with just a holder: text_parts = [] # If a direct text/plain part occurs, # we just add it to the list: if ( payload["contentMainType"] == "text" and payload["contentSubType"] == "plain" ): text_parts.append(payload["payload"]) # If a direct text/html part occurs, # we just add it to the list: if ( payload["contentMainType"] == "text" and payload["contentSubType"] == "html" ): html_parser = BeautifulSoup(payload["payload"], "html.parser") text_parts.append(html_parser.get_text()) # If a multipart/alternative part occurs, # we pick just the ready plaintext part: if ( payload["contentMainType"] == "multipart" and payload["contentSubType"] == "alternative" ): for part in payload["payload"]: if part["contentSubType"] == "plain": text_parts.append(part["payload"]) # If a multipart/mixed or multipart/related part occurs, # we use recursion to look for plaintext parts nested inside: if ( payload["contentMainType"] == "multipart" and ( payload["contentSubType"] == "mixed" or payload["contentSubType"] == "related" ) ): for part in payload["payload"]: text_parts += self.extract_plaintext_parts(payload = part) # Done here: return text_parts async def summarize_mail_with_ai( self, mongo_conn: AsyncMongo, user_info: CoreUserInfoModel, llm: CoreLLMController, message: CoreMessageModel, prompt_template: List[dict] ) -> LLMOutput: # Extract the text from the message here: text_parts = self.extract_plaintext_parts(payload = message.message["payload"]) text = "\n".join(text_parts) # Invoke the LLM and return the response: return await llm.invoke( mongo_conn = mongo_conn, user_info = user_info, llm_input = LLMInput( messages = prompt_template + [ { "role": "human", "content": f"Please summarize this mail: \"\"\"{text}\"\"\"" } ] ) ) def drop_attachments( self, payload: dict ): # If the part is some sort of file: if payload["contentMainType"] not in ["multipart", "text"]: payload["payload"] = None payload["payloadId"] = None payload["payloadUrl"] = None # If the payload is of multipart type, # we use recursion to look inside it: elif payload["contentMainType"] == "multipart": for part in payload["payload"]: self.drop_attachments(part) # Done here: return payload # ┏┓┏┓ ┓ ┏┓ ┏┓ # ┃┃┣┫┓┏╋┣┓┏┛ ┃┫ # ┗┛┛┗┗┻┗┛┗┗━•┗┛ @staticmethod async def get_token_key( db_conn: AsyncMySQL, mongo_conn: AsyncMongo, auth_token: CoreAuthTokenModel, session_token: str = None ) -> ObjectId: # Simply call the core model: return await current_app.core_auth_token_controller.get_token_key( sql_conn = db_conn, mongo_data_conn = mongo_conn, auth_token = auth_token, token_notes = { "email": None }, session_token = session_token, ) @staticmethod async def set_token( db_conn: AsyncMySQL, mongo_conn: AsyncMongo, token_key: ObjectId | str, auth_token: CoreAuthTokenModel, session_token: str = None ) -> bool: # Simply call the core model: return await current_app.core_auth_token_controller.set_token( sql_conn = db_conn, mongo_data_conn = mongo_conn, token_key = token_key, auth_token = auth_token, token_notes = { "email": auth_token.token["email"], "displayName": auth_token.token.get("displayName"), "displayPictureUrl": auth_token.token.get("displayPictureUrl"), }, session_token = session_token, ) @staticmethod async def get_token_from_key( mongo_conn: AsyncMongo, token_key: ObjectId | str = None, ) -> CoreAuthTokenModel | None: # Simply call the core model: return await current_app.core_auth_token_controller.get_token_from_key( mongo_data_conn = mongo_conn, token_key = token_key ) @staticmethod async def get_tokens_from_keys( mongo_conn: AsyncMongo, token_keys: ObjectId | str = None, ) -> List[CoreAuthTokenModel] | None: # Simply call the core model: return await current_app.core_auth_token_controller.get_tokens_from_keys( mongo_data_conn = mongo_conn, token_keys = token_keys ) # ┏┓ ┳┳┓ # ┗┓┓┏┏┓┏ ┃┃┃┏┓┏┏┏┓┏┓┏┓┏ # ┗┛┗┫┛┗┗ ┛ ┗┗ ┛┛┗┻┗┫┗ ┛ # ┛ ┛ # In this section, we pull mails from the third-party clients (like GMail), and store them on our server. This makes # those mails available on the platform. async def __sync_one_gmail( self, mongo_conn: AsyncMongo, user_info: CoreUserInfoModel, auth_token: CoreAuthTokenModel, mail_client: AsyncGMailClient, google_tokens: GoogleAuthTokens, message_id: str, llm: CoreLLMController = None, force_sync: bool = False ) -> MailSyncOneResult: # 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_records = await current_app.core_message_controller.get_previews( mongo_conn = mongo_conn, token_ids = [ObjectId(auth_token.authTokenId)], limit = 1, skip = 0, additional_filter = { "tokenId": ObjectId(auth_token.authTokenId), "serviceType": auth_token.serviceType, "client": auth_token.client, "clientMessageId": message_id } ) if mail_records: sync_result.success = True sync_result.message = ( f"gmail message '{message_id}' already " f"sync'd on '{mail_records[0].syncTs} (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 = google_tokens, message_id = message_id, return_raw = False ) # If we didn't get the mail from GMail; if not client_response.success: sync_result.message = f"gmail (messageId: '{message_id}'): {client_response.message}" return sync_result # HANDLE ATTACHMENTS HERE: client_response.data["payload"] = self.drop_attachments(client_response.data["payload"]) # Now we structure the message into the model: mail_message = CoreMessageModel( ts = client_response.data["ts"], syncTs = date_time.get_current_utc_date_time(as_string = False), tokenId = auth_token.authTokenId, serviceType = auth_token.serviceType, client = auth_token.client, clientMessageId = message_id, clientThreadId = client_response.data["threadId"], isSent = False, isBroadcast = False, sentSuccessfully = False, sender = client_response.data["from"][0]["name"], chat = None, message = client_response.data, snippet = client_response.data["subject"], aiSnippet = None, tags = ["Email", "Gmail"] ) # 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 google_tokens.email in all_recipients: mail_message.isSent = False else: mail_message.isSent = True # Invoke the LLM: try: ai_snippet = await self.summarize_mail_with_ai( mongo_conn = mongo_conn, user_info = user_info, llm = llm, message = mail_message, prompt_template = self.PROMPT_TEMPLATE ) ai_json = ai_snippet.json mail_message.aiSnippet = ai_snippet.summary mail_message.aiSnippet["output"] = ai_json["summary"] if ai_json["senderType"] is not None: mail_message.tags.append(ai_json["senderType"]) except Exception as exception: current_app.printer(exception) # Done here: sync_result.success = True sync_result.mailMessage = mail_message return sync_result async def __sync_many_gmail( self, db_conn: AsyncMySQL, mongo_conn: AsyncMongo, user_info: CoreUserInfoModel, auth_token: CoreAuthTokenModel, mail_client: AsyncGMailClient, llm: CoreLLMController = None, force_sync: bool = False, start_date: datetime.datetime = None, end_date: datetime.datetime = None, max_count: int = 100, session_token: str = None ) -> MailSyncManyResults: # Start by assuming failure: sync_results = MailSyncManyResults() # Extract the client's tokens from the full token payload given by the database: google_tokens = GoogleAuthTokens(**auth_token.token) # Refresh the tokens (if needed): tokens_refreshed = await google_tokens.arefresh( http_client = current_app.http_client, client_id = mail_client.client_id, client_secret = mail_client.client_secret ) if tokens_refreshed: auth_token.token = google_tokens.model_dump() auth_token.lastRefreshTs = date_time.get_current_utc_date_time(as_string = True) await self.set_token( db_conn = db_conn, mongo_conn = mongo_conn, token_key = auth_token.key, auth_token = auth_token, session_token = session_token ) # Let's build the query to send to Google: 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 = google_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( mongo_conn = mongo_conn, user_info = user_info, auth_token = auth_token, mail_client = mail_client, google_tokens = google_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: replacement_json = result.mailMessage.model_dump() replacement_json.pop("_id", None) mongo_operations.append( ReplaceOne( filter = { "tokenId": ObjectId(auth_token.authTokenId), "serviceType": auth_token.serviceType, "client": auth_token.client, "clientMessageId": result.mailMessage.clientMessageId }, replacement = replacement_json, upsert = True ) ) # Make the bulk insert operation: if mongo_operations: sync_count = await current_app.core_message_controller.bulk_write( mongo_conn = mongo_conn, mongo_operations = mongo_operations ) # Apply the labels to the read messages: try: client_response = await mail_client.modify_messages( tokens = google_tokens, message_ids = [v["id"] for v in messages_list.values()], add_label_ids = [google_tokens.labels.get("TCAOFF", {}).get("id")] ) except Exception as exception: pass # 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, db_conn: AsyncMySQL, mongo_conn: AsyncMongo, user_info: CoreUserInfoModel, token_key: ObjectId | str, llm: CoreLLMController = None, force_sync: bool = False, start_date: datetime.datetime = None, end_date: datetime.datetime = None, max_count: int = 100, session_token: str = None ) -> MailSyncManyResults: # Start by assuming failure: sync_results = MailSyncManyResults() # ┏┓ ┓ ┏┳┓ ┓ # ┣ ┏┓╋┏┣┓ ┃ ┏┓┃┏┏┓┏┓┏ # ┻ ┗ ┗┗┛┗ ┻ ┗┛┛┗┗ ┛┗┛ # We first load the authorization tokens: auth_token = await self.get_token_from_key( mongo_conn = mongo_conn, token_key = token_key, ) # If we failed to load the authorization tokens: if not auth_token: sync_results.message = f"no such token key" return sync_results # ┏┓ ┏┓┳┳┓ •┓ # ┣ ┏┓┏┓ ┃┓┃┃┃┏┓┓┃ # ┻ ┗┛┛ ┗┛┛ ┗┗┻┗┗ if auth_token.client == "gmail": return await self.__sync_many_gmail( db_conn = db_conn, mongo_conn = mongo_conn, user_info = user_info, auth_token = auth_token, mail_client = current_app.gmail_client, llm = llm, force_sync = force_sync, start_date = start_date, end_date = end_date, max_count = max_count, session_token = session_token, ) # ┳ ┓• ┓ ┏┓┓• # ┃┏┓┓┏┏┓┃┓┏┫ ┃ ┃┓┏┓┏┓╋ # ┻┛┗┗┛┗┻┗┗┗┻ ┗┛┗┗┗ ┛┗┗ # 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_token.client}'" return sync_results # ┏┓ ┓ ┳┳┓ # ┗┓┏┓┏┓┏┫ ┃┃┃┏┓┏┏┏┓┏┓┏┓┏ # ┗┛┗ ┛┗┗┻ ┛ ┗┗ ┛┛┗┻┗┫┗ ┛ # ┛ async def __send_one_gmail( self, db_conn: AsyncMySQL, mongo_conn: AsyncMongo, user_info: CoreUserInfoModel, token_key: ObjectId | str, inbound_data: MailSendRequestData, inbound_files: dict, llm: CoreLLMController = None, session_token: str = None ) -> MailSendOneResult: # Start by assuming failure: sync_results = MailSendOneResult() # ┏┓ ┏┳┓┓ ┳┳┓ # ┃ ┏┓┏┓┏╋┏┓┓┏┓┏┏╋ ┃ ┣┓┏┓ ┃┃┃┏┓┏┏┏┓┏┓┏┓ # ┗┛┗┛┛┗┛┗┛ ┗┻┛┗┗┗ ┻ ┛┗┗ ┛ ┗┗ ┛┛┗┻┗┫┗ # ┛ gmail_message = None async def send_one_mail( self, db_conn: AsyncMySQL, mongo_conn: AsyncMongo, user_info: CoreUserInfoModel, token_key: ObjectId | str, inbound_data: MailSendRequestData, inbound_files: dict, llm: CoreLLMController = None, session_token: str = None ) -> MailSendOneResult: # Start by assuming failure: sync_results = MailSendOneResult() # ┏┓ ┓ ┏┳┓ ┓ # ┣ ┏┓╋┏┣┓ ┃ ┏┓┃┏┏┓┏┓┏ # ┻ ┗ ┗┗┛┗ ┻ ┗┛┛┗┗ ┛┗┛ # We first load the authorization tokens: auth_token = await self.get_token_from_key( mongo_conn = mongo_conn, token_key = token_key, ) # If we failed to load the authorization tokens: if not auth_token: sync_results.message = f"no such token key" return sync_results # ┏┓ ┏┓┳┳┓ •┓ # ┣ ┏┓┏┓ ┃┓┃┃┃┏┓┓┃ # ┻ ┗┛┛ ┗┛┛ ┗┗┻┗┗ # if auth_token.client == "gmail": # return await self.__send_one_gmail( # db_conn = db_conn, # mongo_conn = mongo_conn, # user_info = user_info, # auth_token = auth_token, # mail_client = current_app.gmail_client, # llm = llm, # force_sync = force_sync, # start_date = start_date, # end_date = end_date, # max_count = max_count, # session_token = session_token, # ) # ┳ ┓• ┓ ┏┓┓• # ┃┏┓┓┏┏┓┃┓┏┫ ┃ ┃┓┏┓┏┓╋ # ┻┛┗┗┛┗┻┗┗┗┻ ┗┛┗┗┗ ┛┗┗ # 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_token.client}'" return sync_results # ┓ • ┏┓ ┏┓ ┳┳┓ # ┃ ┓┏╋ ┣╋ ┃┓┏┓╋ ┃┃┃┏┓┏┏┏┓┏┓┏┓┏ # ┗┛┗┛┗ ┗┻ ┗┛┗ ┗ ┛ ┗┗ ┛┛┗┻┗┫┗ ┛ # ┛ # These are simply for retrieving mails. You need to already have them sync'd to the database. These methods don't # fetch the mails from the third-party clients. @staticmethod async def list_mails( mongo_conn: AsyncMongo, token_ids: List[ObjectId | str], limit: int = 100, skip: int = 0, additional_filter: dict = None ) -> List[CoreMessageModel] | None: # Regardless of what additional filter is provided from outside, # we add a mail-selecting filter here: if additional_filter is None: additional_filter = {} additional_filter["serviceType"] = "email" # Simply call the core model: return await current_app.core_message_controller.get_previews( mongo_conn = mongo_conn, token_ids = token_ids, limit = limit, skip = skip, additional_filter = additional_filter ) @staticmethod async def get_one_mail( mongo_conn: AsyncMongo, message_id: ObjectId | str ) -> CoreMessageModel | None: # Simply call the core model: return await current_app.core_message_controller.get_message( mongo_conn = mongo_conn, message_id = message_id ) # ┳┳ ┓ # ┃┃┏┓┏┫┏┓╋┏┓ # ┗┛┣┛┗┻┗┻┗┗ # ┛ @staticmethod async def update_tags( mongo_conn: AsyncMongo, message_id: ObjectId | str, unset_tags: List[str] = None, set_tags: List[str] = None ) -> bool: # Simply call the core model: return await current_app.core_message_controller.update_tags( mongo_conn = mongo_conn, message_id = message_id, unset_tags = unset_tags, set_tags = set_tags ) # ***************************************************************************************************************** # ***** **** # *** MAIN PROGRAM *** # ***** **** # ***************************************************************************************************************** if __name__ == "__main__": pass # from utils_v2.string import json # # file_options = [ # r"/home/developer/Downloads/recursive parts parse - 20241210.json", # r"/home/developer/Downloads/recursive parts parse (no attachment) - 20241210.json", # ] # # raw_mail_json = json.from_file(file_options[1]) # print("FROM FILE:", json.to_string(raw_mail_json["payload"])) # print("\n\n---------\n\n") # mail_controller = MailController() # print(json.to_string(mail_controller.drop_attachments(raw_mail_json["payload"])))