(20241202) Testing out the modified auth process.
This commit is contained in:
+144
-25
@@ -34,6 +34,7 @@
|
||||
import sys
|
||||
|
||||
from langchain.chains.summarize.stuff_prompt import prompt_template
|
||||
from sqlalchemy.orm.collections import collection
|
||||
|
||||
sys.path.append(".")
|
||||
sys.path.append("..")
|
||||
@@ -53,6 +54,7 @@ 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
|
||||
@@ -64,6 +66,12 @@ from typing import Literal
|
||||
# To make deep-copies:
|
||||
import copy
|
||||
|
||||
# To work with date and time:
|
||||
import datetime
|
||||
|
||||
# For asynchronous activities:
|
||||
import asyncio
|
||||
|
||||
|
||||
# *****************************************************************************************************************
|
||||
# ***** ****
|
||||
@@ -112,7 +120,7 @@ class MailSyncModel(BaseModel):
|
||||
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."
|
||||
"You're a mail summary expert that summarizes mails in 150 chars or less. HIDE SENSITIVE INFO (LIKE OTPs) FROM THE SUMMARY."
|
||||
),
|
||||
(
|
||||
"user",
|
||||
@@ -120,17 +128,24 @@ class MailSyncModel(BaseModel):
|
||||
)
|
||||
])
|
||||
|
||||
async def sync_one(
|
||||
async def __sync_one(
|
||||
self,
|
||||
mongo_conn: AsyncMongo,
|
||||
user_identifier: str | ObjectId,
|
||||
user_info: dict,
|
||||
mail_client: AsyncGMailClient,
|
||||
tokens: GoogleAuthTokens,
|
||||
message_id: str,
|
||||
llm: ChatOpenAI = None,
|
||||
session_token: str = None,
|
||||
force_sync: bool = False
|
||||
) -> ObjectId:
|
||||
) -> UpdateOne | None:
|
||||
|
||||
# ┏┓ ┓┏ • ┓ ┓
|
||||
# ┃┃┏┓┏┓┏┓┏┓┏┓┏┓ ┃┃┏┓┏┓┓┏┓┣┓┃┏┓┏
|
||||
# ┣┛┛ ┗ ┣┛┗┻┛ ┗ ┗┛┗┻┛ ┗┗┻┗┛┗┗ ┛
|
||||
# ┛
|
||||
|
||||
mail_payload = None
|
||||
mail_client_name = None
|
||||
|
||||
# ┏┓┓ ┓ ┏┓ • • ┳┓ ┓
|
||||
# ┃ ┣┓┏┓┏┃┏ ┣ ┓┏┓┏╋┓┏┓┏┓ ┣┫┏┓┏┏┓┏┓┏┫┏
|
||||
@@ -150,15 +165,7 @@ class MailSyncModel(BaseModel):
|
||||
)
|
||||
|
||||
# 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 existing_record["_id"]
|
||||
|
||||
# ┏┓ ┓┏ • ┓ ┓
|
||||
# ┃┃┏┓┏┓┏┓┏┓┏┓┏┓ ┃┃┏┓┏┓┓┏┓┣┓┃┏┓┏
|
||||
# ┣┛┛ ┗ ┣┛┗┻┛ ┗ ┗┛┗┻┛ ┗┗┻┗┛┗┗ ┛
|
||||
# ┛
|
||||
|
||||
mail_payload = None
|
||||
mail_id = existing_record["_id"] if existing_record else None
|
||||
if existing_record and not force_sync: return mail_payload
|
||||
|
||||
# ┏┓┳┳┓ •┓
|
||||
# ┃┓┃┃┃┏┓┓┃
|
||||
@@ -166,10 +173,10 @@ class MailSyncModel(BaseModel):
|
||||
|
||||
if isinstance(mail_client, AsyncGMailClient):
|
||||
|
||||
# Refresh the tokens:
|
||||
tokens_refreshed
|
||||
# Note down the name of the mail client:
|
||||
mail_client_name = "gmail"
|
||||
|
||||
# Fetch the mail formatted message:
|
||||
# Fetch the formatted mail message:
|
||||
client_response = await mail_client.get_message(
|
||||
tokens = tokens,
|
||||
message_id = message_id,
|
||||
@@ -180,20 +187,132 @@ class MailSyncModel(BaseModel):
|
||||
if client_response.success:
|
||||
|
||||
# Summarize the content:
|
||||
prompt = self.prompt_template.invoke({"mail": client_response.data.pop["unformattedText"]})
|
||||
llm_response = await llm.ainvoke(prompt)
|
||||
client_response.data["aiSnippet"] = llm_response.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:
|
||||
pass
|
||||
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
|
||||
|
||||
|
||||
# *****************************************************************************************************************
|
||||
|
||||
Reference in New Issue
Block a user