(20241205) LLM endpoint active now.

This commit is contained in:
2024-12-05 18:09:32 +05:30
parent 0ec57089b0
commit 81550898eb
10 changed files with 347 additions and 383 deletions
+107 -214
View File
@@ -10,7 +10,7 @@
OBJECTIVE:
To c
To create an interface between OpenAI and our internal system to perform LLM-based activities.
REFERENCES:
@@ -43,6 +43,12 @@ from utils_v2.database.async_mongo_v2 import AsyncMongo
# Base model:
from models.behaviour.base import BaseModel
# Data Models:
from models.data.ai.llm import LLMInput, LLMOutput, LLMUsageTokens
# To work with LLMs:
from langchain_openai import ChatOpenAI
# To work with MongoDB:
from bson import ObjectId
@@ -90,230 +96,86 @@ import copy
# *****************************************************************************************************************
class MailOAuthModel(BaseModel):
class LLMOpenAI(BaseModel):
AUTH_COLLECTION = "_authTokens"
AI_USAGE_COLLECTION = "_aiUsage"
async def get_token_id(
def __init__(
self,
llm_creds: dict,
cache = None,
alert_url = None,
http_client = None,
debug = True,
debug_prefix = "Model | ",
debug_only_errors = True
):
"""
This is the model that works with OpenAi's LLM to perform tasks like text completion.
:param llm_creds: The JSON that holds the credentials to access your OpenAI account. Should have the keys
'model', and 'openai_api_key'.
:param cache: The object to use for caching results from database calls.
:param alert_url: Which URL to call when something goes wrong.
:param http_client: The instance of an HTTP client to use when trying to send alerts and make other APIs.
:param debug: Whether, or not, you would like to print debugging messages:
:param debug_prefix: The prefix to print with the debugging messages.
:param debug_only_errors: Whether you would like to print only error messages or all messages.
:return: None.
"""
# Initialize the parent:
super().__init__(
cache = cache,
alert_url = alert_url,
http_client = http_client,
debug = debug,
debug_prefix = debug_prefix,
debug_only_errors = debug_only_errors
)
# Create the interface to the LLM:
self.__llm = ChatOpenAI(**llm_creds)
async def invoke(
self,
db_conn: AsyncMySQL,
mongo_conn: AsyncMongo,
user_info: dict,
client_user_id: dict,
auth: dict,
service_client: Literal["gmail"],
auth_type: Literal["oauth"],
sync_freq: Literal[60, 300, 900] = 300,
session_token: str = None
) -> ObjectId:
llm_input: LLMInput
) -> LLMOutput:
"""
Stores params from the session info and gives an identifier to use in the authorization URL. Use this when the
user requests an authorization URL to link your service to another service (like GMail).
:param db_conn: The database connection (MariaDB) to use to perform the action.
:param mongo_conn: The database connection (MongoDB) to use to perform the action.
:param user_info: The dictionary that has the user's session information.
:param client_user_id: The way the third-party client recognizes your user.
:param auth: The authentication details of the account.
:param service_client: The name of the company or brand that is providing this service that is being integrated.
:param auth_type: To identify the type of authentication being done here. This could indicate simple password
authentication, more advance OAuth2.0 authentication, etc.
:param sync_freq: The time interval in which mails need to be sync'd. Specify this in seconds.
:param session_token: The session token of the user who requested this service.
:return: An ObjectId to later store the granted tokens.
"""
# Format the message as per the format of OpenAI:
prompt = [
{
"role": {"system": "system", "ai": "assistant", "human": "user"}[message.role],
"content": message.content
} for message in llm_input.messages
]
# Note down the timestamp at which this event occurred:
request_ts = date_time.get_current_utc_date_time(as_string = False)
# Get the identifier from the database:
mongo_json = await mongo_conn.find_one_and_update(
collection = MailOAuthModel.AUTH_COLLECTION,
filter = mongo_conn.dict_to_dot_notation({
"serviceType": "email",
"user": {
"entityId": user_info["entityId"],
"billingAccountId": user_info["billingAccountId"]
},
"clientUserId": client_user_id
}),
update = {
"$set": {
"lastRequestTs": request_ts,
"status": "active",
"syncFreq": max(sync_freq, 60)
},
"$setOnInsert": {
"version": "1.1.1",
"serviceType": "email",
"client": service_client,
"authType": auth_type,
"user": user_info,
"clientUserId": client_user_id,
"auth": auth,
"token": None,
"firstRefreshTs": None,
"lastRefreshTs": None,
"firstRequestTs": request_ts,
}
},
projection = {
"_id": True
},
upsert = True,
return_updated = True
# Invoke the AI, and format the response:
llm_response = await self.__llm.ainvoke(prompt)
llm_response = LLMOutput(
messages = llm_input.messages,
output = llm_response.content,
client = "openai",
model = llm_response.response_metadata["model_name"],
tokens = LLMUsageTokens(
input = llm_response.usage_metadata["input_tokens"],
output = llm_response.usage_metadata["output_tokens"],
total = llm_response.usage_metadata["total_tokens"],
)
)
# Tell MariaDB that an authorization request was initiated:
db_json = {}
if mongo_json is not None:
db_json = await self.call_procedure(
db_conn = db_conn,
proc_name = "entity_integration_save",
proc_args = (
user_info["entityId"], # ............................................ 'p_entity_id'
service_client, # ................................................... 'p_provider'
"Pending", # ........................................................ 'p_current_status'
"Auth Requested", # ................................................. 'p_last_action'
None, # ............................................................. 'p_display_name'
None, # ............................................................. 'p_display_picture'
str(mongo_json["_id"]), # ........................................... 'p_token_id'
json.to_string(python_data = {"email": None}, no_space = True), # ... 'p_notes'
user_info["userId"] # ............................................... 'p_created_by'
),
session_token = session_token
)
# Store this into MongoDB:
mongo_document = {"user": user_info}
for k, v in llm_response.model_dump().items(): mongo_document[k] = v
inserted_id = await mongo_conn.insert_one(
collection = self.AI_USAGE_COLLECTION,
document = mongo_document
)
# Done here:
return mongo_json["_id"] if mongo_json and db_json.get("status") == 1 else None
async def set_token(
self,
db_conn: AsyncMySQL,
mongo_conn: AsyncMongo,
token_id: ObjectId | str,
client_user_id: dict,
token: dict,
session_token: str = None
) -> bool:
"""
This method is to be called when the end user authorizes your service to connect to his third-party account. For
example, when the end user allows you to access his GMail account. USE THIS FOR UPDATING (REFRESHING) TOKENS
ALSO.
:param db_conn: The database connection (MariaDB) to use to perform the action.
:param mongo_conn: The database connection (MongoDB) to use to perform the action.
:param token_id: The identifier granted by the 'get_token_id' method.
:param client_user_id: The way the third-party client recognizes your user. These details should match the
details furnished while requesting the authorization through 'get_token_id' method.
:param token: The token granted by the third-party service.
:param session_token: The session token of the user who requested this service.
:return: True if saved, False if failed.
"""
# Start by assuming failure:
token_saved = False
# Note down the timestamp at which this event occurred:
request_ts = date_time.get_current_utc_date_time(as_string = False)
# Save the token to MongoDB:
mongo_json = await mongo_conn.find_one_and_update(
collection = MailOAuthModel.AUTH_COLLECTION,
filter = mongo_conn.dict_to_dot_notation({
"_id": ObjectId(token_id),
"clientUserId": client_user_id
}),
update = [{
"$set": {
"token": token,
"status": "active",
"lastRefreshTs": request_ts,
"firstRefreshTs": {
"$cond": {
"if": {
"$or": [
{"$eq": ["$firstRefreshTs", None]},
{"$eq": [{"$type": "$firstRefreshTs"}, "missing"]}
]
},
"then": request_ts,
"else": "$firstRefreshTs"
}
}
}
}],
projection = {"token": False},
return_updated = True,
upsert = False
)
# Tell MariaDB that the token was saved:
if mongo_json is not None:
token_notes = {
"email": token["email"],
"displayName": token.get("displayName"),
"displayPictureUrl": token.get("displayPictureUrl"),
}
db_json = await self.call_procedure(
db_conn = db_conn,
proc_name = "entity_integration_save",
proc_args = (
mongo_json["user"]["entityId"], # ............................... 'p_entity_id'
mongo_json["client"], # ......................................... 'p_provider'
"Active", # ..................................................... 'p_current_status'
"Auth Granted", # ............................................... 'p_last_action'
token["displayName"], # ......................................... 'p_display_name'
token["displayPictureUrl"], # ................................... 'p_display_picture'
token_id, # ..................................................... 'p_token_id'
json.to_string(python_data = token_notes, no_space = True), # ... 'p_notes'
mongo_json["user"]["userId"] # .................................. 'p_created_by'
),
session_token = session_token
)
if db_json["status"] == 1: token_saved = True
# Done here:
return token_saved
async def get_token(
self,
mongo_conn: AsyncMongo,
token_id: ObjectId | str = None,
**kwargs
) -> dict | None:
"""
To retrieve stored tokens from the database.
:param mongo_conn: The database connection (MongoDB) to use to perform the action.
:param token_id: The identifier granted by the 'get_token_id' method.
:param kwargs: Any set of key-value pairs to build custom search criteria. This could be things like the user
info, the client, the type of authentication used, or even the kind of service.
:return: The retrieved record that has the token, and information about the service and client if found, else
None when there is no matching record.
"""
# Build the filter:
filter_json = {k: v for k, v in kwargs.items()}
if token_id: filter_json["_id"] = ObjectId(token_id)
# If there is no search criteria, we exit with failure:
if not filter_json: return None
# If there is some filtering possible,
# we fetch and return the token:
return await mongo_conn.find_one(
collection = self.AUTH_COLLECTION,
filter = filter_json,
projection = {
"_id": True,
"serviceType": True,
"authType": True,
"client": True,
"clientUserId": True,
"token": True
}
)
return llm_response
# *****************************************************************************************************************
@@ -326,3 +188,34 @@ class MailOAuthModel(BaseModel):
if __name__ == "__main__":
pass
# import asyncio
#
# llm_messages = [
# {
# "role": "system",
# "content": "You are an office assistant."
# },
# {
# "role": "ai",
# "content": "Hello, sir. How may I help you today?"
# },
# {
# "role": "human",
# "content": "Please summarize this mail for me..."
# }
# ]
#
# my_llm = LLMOpenAI(
# llm_creds = {
# "model": "gpt-4o-mini",
# "openai_api_key": "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"
# }
# )
#
# async def main():
#
# llm_response = await my_llm.invoke(llm_input = LLMInput(messages = llm_messages))
# print("LLM RESPONSE:", llm_response.model_dump_json(indent = 4))
#
# asyncio.run(main())
+3 -8
View File
@@ -124,7 +124,9 @@ class MailRetrieveModel(BaseModel):
"payload.bcc": True,
"payload.parts": True,
"payload.attachments": True,
"payload.labels": True
"payload.labels": True,
"payload.snippet": True,
"payload.aiSnippet": True,
}
)
@@ -133,8 +135,6 @@ class MailRetrieveModel(BaseModel):
mail_data["mailId"] = str(mail_data.pop("_id"))
mail_data["payload"]["ts"] = mail_data["payload"]["ts"].isoformat()
mail_data["payload"]["readTs"] = mail_data["payload"]["readTs"].isoformat()
if ai_snippet := mail_data["payload"].pop("aiSnippet"):
mail_data["payload"]["aiSnippet"] = ai_snippet["snippet"]
# Done here:
return mail_data
@@ -186,11 +186,6 @@ class MailRetrieveModel(BaseModel):
mail_data["mailId"] = str(mail_data.pop("_id"))
mail_data["payload"]["ts"] = mail_data["payload"]["ts"].isoformat()
mail_data["payload"]["readTs"] = mail_data["payload"]["readTs"].isoformat()
if ai_snippet := mail_data["payload"].pop("aiSnippet"):
mail_data["payload"]["aiSnippet"] = {
"snippet": ai_snippet["snippet"],
"usage": ai_snippet["usage"]
}
# Done here:
return mails_list
+60 -29
View File
@@ -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,