(20241205) LLM endpoint active now.
This commit is contained in:
@@ -31,8 +31,6 @@
|
||||
|
||||
# To make sibling directories accessible for imports:
|
||||
import sys
|
||||
from logging import exception
|
||||
|
||||
sys.path.append(".")
|
||||
sys.path.append("..")
|
||||
|
||||
@@ -60,8 +58,8 @@ 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
|
||||
from models.behaviour.ai.llm.open_ai import LLMOpenAI
|
||||
from models.data.ai.llm import LLMInput
|
||||
|
||||
# To work with datatypes:
|
||||
from typing import Literal, List, Dict, Any
|
||||
@@ -123,16 +121,20 @@ class MailSyncModel(BaseModel):
|
||||
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}\"\"\""
|
||||
)
|
||||
])
|
||||
PROMPT_TEMPLATE = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"You're a mail summary expert that summarizes mails in 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."
|
||||
)
|
||||
}
|
||||
]
|
||||
|
||||
# ┏┓ ┓
|
||||
# ┣┫╋╋┏┓┏┣┓┏┳┓┏┓┏┓╋┏
|
||||
# ┛┗┗┗┗┻┗┛┗┛┗┗┗ ┛┗┗┛
|
||||
|
||||
@staticmethod
|
||||
async def __save_one_attachment(
|
||||
@@ -237,20 +239,26 @@ class MailSyncModel(BaseModel):
|
||||
# Done here:
|
||||
return uploaded_attachments
|
||||
|
||||
# ┏┓ ┏┓┳┳┓ •┓
|
||||
# ┣ ┏┓┏┓ ┃┓┃┃┃┏┓┓┃
|
||||
# ┻ ┗┛┛ ┗┛┛ ┗┗┻┗┗
|
||||
|
||||
async def __sync_one_gmail(
|
||||
self,
|
||||
session_token: str,
|
||||
user_info: dict,
|
||||
mongo_conn: AsyncMongo,
|
||||
mail_client: AsyncGMailClient,
|
||||
tokens: GoogleAuthTokens,
|
||||
message_id: str,
|
||||
llm: ChatOpenAI = None,
|
||||
llm: LLMOpenAI = 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 user_info: The information of the user (derived from his session token).
|
||||
: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.
|
||||
@@ -291,8 +299,11 @@ class MailSyncModel(BaseModel):
|
||||
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
|
||||
|
||||
# We upload the attachments:
|
||||
client_response.data["attachments"] = await self.__save_many_attachments(
|
||||
@@ -321,23 +332,33 @@ class MailSyncModel(BaseModel):
|
||||
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:
|
||||
# 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"]
|
||||
})
|
||||
|
||||
# Invoke the LLM:
|
||||
llm_response = response = await llm.invoke(
|
||||
mongo_conn = mongo_conn,
|
||||
user_info = user_info,
|
||||
llm_input = LLMInput(
|
||||
messages = self.PROMPT_TEMPLATE + [
|
||||
{
|
||||
"role": "human",
|
||||
"content": f"Please summarize this mail: \"\"\"{client_response.data["unformattedText"]}\"\"\""
|
||||
}
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
# Format the response:
|
||||
llm_json = {
|
||||
"snippet": llm_response.content,
|
||||
"usage": {
|
||||
"input": llm_response.usage_metadata["input_tokens"],
|
||||
"output": llm_response.usage_metadata["output_tokens"],
|
||||
"total": llm_response.usage_metadata["total_tokens"],
|
||||
},
|
||||
"rawUsage": llm_response.usage_metadata
|
||||
"ts": llm_response.ts,
|
||||
"snippet": llm_response.output,
|
||||
"tokens": llm_response.tokens.model_dump()
|
||||
}
|
||||
|
||||
# Add the LLM's response to the main data:
|
||||
client_response.data["aiSnippet"] = llm_json
|
||||
|
||||
# Done here:
|
||||
@@ -348,11 +369,12 @@ class MailSyncModel(BaseModel):
|
||||
async def __sync_many_gmail(
|
||||
self,
|
||||
session_token: str,
|
||||
user_info: dict,
|
||||
mongo_conn: AsyncMongo,
|
||||
token_id: ObjectId,
|
||||
mail_client: AsyncGMailClient,
|
||||
tokens: GoogleAuthTokens,
|
||||
llm: ChatOpenAI = None,
|
||||
llm: LLMOpenAI = None,
|
||||
force_sync: bool = False,
|
||||
start_date: datetime.datetime = None,
|
||||
end_date: datetime.datetime = None,
|
||||
@@ -362,6 +384,7 @@ class MailSyncModel(BaseModel):
|
||||
"""
|
||||
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 user_info: The information of the user (derived from his session token).
|
||||
:param mongo_conn: The instance of the connection to the database to use.
|
||||
:param token_id: 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.
|
||||
@@ -415,6 +438,7 @@ class MailSyncModel(BaseModel):
|
||||
tasks = [
|
||||
self.__sync_one_gmail(
|
||||
session_token = session_token,
|
||||
user_info = user_info,
|
||||
mongo_conn = mongo_conn,
|
||||
mail_client = mail_client,
|
||||
tokens = tokens,
|
||||
@@ -473,12 +497,17 @@ class MailSyncModel(BaseModel):
|
||||
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,
|
||||
user_info: dict,
|
||||
mongo_conn: AsyncMongo,
|
||||
token_id: ObjectId,
|
||||
llm: ChatOpenAI = None,
|
||||
llm: LLMOpenAI = None,
|
||||
force_sync: bool = False,
|
||||
start_date: datetime.datetime = None,
|
||||
end_date: datetime.datetime = None,
|
||||
@@ -489,6 +518,7 @@ class MailSyncModel(BaseModel):
|
||||
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 user_info: The information of the user (derived from his session token).
|
||||
:param mongo_conn: The instance of the connection to the database to use.
|
||||
:param token_id: 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.
|
||||
@@ -526,6 +556,7 @@ class MailSyncModel(BaseModel):
|
||||
if auth_json["client"] == "gmail":
|
||||
return await self.__sync_many_gmail(
|
||||
session_token = session_token,
|
||||
user_info = user_info,
|
||||
mongo_conn = mongo_conn,
|
||||
token_id = token_id,
|
||||
mail_client = current_app.gmail_client,
|
||||
|
||||
Reference in New Issue
Block a user