From 81550898eb367b530b592780c4740b9c563b90b4 Mon Sep 17 00:00:00 2001 From: khushal Date: Thu, 5 Dec 2024 18:09:32 +0530 Subject: [PATCH] (20241205) LLM endpoint active now. --- api/blueprints/ai/llm/invoke.py | 98 +++----- api/blueprints/mail/oauth_callback.py | 2 +- api/blueprints/mail/sync.py | 4 + api/blueprints/sms/auth.py | 6 +- api/main.py | 22 +- models/behaviour/ai/llm/open_ai.py | 321 +++++++++----------------- models/behaviour/mail/retrieve.py | 11 +- models/behaviour/mail/sync_v2.py | 89 ++++--- models/data/ai/llm.py | 175 +++++++++----- models/data/sms/auth.py | 2 +- 10 files changed, 347 insertions(+), 383 deletions(-) diff --git a/api/blueprints/ai/llm/invoke.py b/api/blueprints/ai/llm/invoke.py index 7693d40..b2c0e5b 100644 --- a/api/blueprints/ai/llm/invoke.py +++ b/api/blueprints/ai/llm/invoke.py @@ -10,7 +10,7 @@ OBJECTIVE: - To receive auth details for various SMS client APIs. + To use LLMs to perform activities like chat completion, text summarization, etc. REFERENCES: @@ -59,15 +59,11 @@ from utils_v2.api.async_quart import ( handle_cancelled_request ) -# SMS-related utils: -from utils_v2.sms.nimbus.async_nimbus import AsyncNimbusSMS -from utils_v2.sms.savvy_bulk_sms.async_savvy_bulk_sms import AsyncSavvyBulkSMS - # Common: from shared import constants # Data Models: -from models.data.sms.auth import SMSAuthRequestHeaders, SMSAuthRequestData +from models.data.ai.llm import LLMRequestHeaders, LLMInput # For asynchronous activities: import asyncio @@ -81,7 +77,7 @@ import asyncio # Related to Quart: -sms_auth_bp = Blueprint("sms_auth", __name__) +llm_invoke_bp = Blueprint("llm_invoke", __name__) # ***************************************************************************************************************** @@ -101,7 +97,7 @@ sms_auth_bp = Blueprint("sms_auth", __name__) # ***************************************************************************************************************** -@sms_auth_bp.record_once +@llm_invoke_bp.record_once def init(blueprint_setup_state): # This gets called when the blueprint is registered. @@ -112,7 +108,7 @@ def init(blueprint_setup_state): # --------------------------------------------------------------------------------------------------------------------- -@sms_auth_bp.route("/auth", methods = ["POST"]) +@llm_invoke_bp.route("/llm/invoke", methods = ["POST"]) @set_api_version(api_version = "1.0.0") @read_input(sanitize_headers = False, sanitize_data = False) @get_session_info(key = "X-Session-Token", session_coro = "get_session") @@ -120,7 +116,7 @@ def init(blueprint_setup_state): attr_name = "logs_mongo", project = constants.PROJECT_NAME, log_type = constants.MODULE_NAME, - operation = "smsAuthApi", + operation = "llmInvokeApi", log_input = True, log_output = True, sensitive_keys = ["sessionToken", "X-Session-Token"] @@ -128,19 +124,19 @@ def init(blueprint_setup_state): @log_chain_to_mongo(attr_name = "logs_mongo") @should_not_be_under_maintenance(attr_name = "is_under_maintenance") @validate_input( - header_validator = lambda x: SMSAuthRequestHeaders(**x).model_dump(), - data_validator = lambda x: SMSAuthRequestData(**x) + header_validator = lambda x: LLMRequestHeaders(**x).model_dump(), + data_validator = lambda x: LLMInput(**x) ) @handle_cancelled_request() -async def request_oauth_authorization_url( - inbound_headers: dict | SMSAuthRequestHeaders = None, - inbound_data: dict | SMSAuthRequestData = None, +async def invoke_llm( + inbound_headers: dict | LLMRequestHeaders = None, + inbound_data: dict | LLMInput = None, inbound_files: dict = None, **kwargs ): """ - Use this when a user wants to register a third-party SMS client with your service. + Use this to invoke an LLM for text completion kind of activities. :param inbound_headers: auto-extracted by the decorators. :param inbound_data: auto-extracted by the decorators. :param inbound_files: auto-extracted by the decorators. @@ -160,54 +156,17 @@ async def request_oauth_authorization_url( http_code = HttpCodes.UNAUTHORIZED ) - # Start by assuming failure: - token_id = None + # ┳ ┓ + # ┃┏┓┓┏┏┓┃┏┏┓ + # ┻┛┗┗┛┗┛┛┗┗ - # ┏┓ ┳┓• ┓ ┏┓┳┳┓┏┓ ┳ ┓• - # ┣ ┏┓┏┓ ┃┃┓┏┳┓┣┓┓┏┏ ┗┓┃┃┃┗┓ ┃┏┓┏┫┓┏┓ - # ┻ ┗┛┛ ┛┗┗┛┗┗┗┛┗┻┛ ┗┛┛ ┗┗┛ ┻┛┗┗┻┗┗┻ - - if inbound_data.messageClient == "nimbusSmsIndia": - - token_id = await current_app.sms_auth_model.set( - db_conn = current_app.sql_writer, - mongo_conn = current_app.data_mongo, - user_info = kwargs["session_info"], - client_user_id = { - "userId": inbound_data.auth.userId, - "senderId": inbound_data.auth.senderId, - "entityId": inbound_data.auth.entityId - }, - auth = inbound_data.auth.model_dump(), - token = None, - service_client = inbound_data.messageClient, - auth_type = "auth", - sync_freq = 300, - session_token = inbound_headers["X-Session-Token"] - ) - - # ┏┓ ┏┓ ┳┓ ┓┓ ┏┓┳┳┓┏┓ ┓┏┓ - # ┣ ┏┓┏┓ ┗┓┏┓┓┏┓┏┓┏ ┣┫┓┏┃┃┏ ┗┓┃┃┃┗┓ ┃┫ ┏┓┏┓┓┏┏┓ - # ┻ ┗┛┛ ┗┛┗┻┗┛┗┛┗┫ ┻┛┗┻┗┛┗ ┗┛┛ ┗┗┛ ┛┗┛┗ ┛┗┗┫┗┻ - # ┛ ┛ - - if inbound_data.messageClient == "savvyBulkSmsKenya": - - token_id = await current_app.sms_auth_model.set( - db_conn = current_app.sql_writer, - mongo_conn = current_app.data_mongo, - user_info = kwargs["session_info"], - client_user_id = { - "partnerId": inbound_data.auth.partnerId, - "shortCode": inbound_data.auth.shortCode - }, - auth = inbound_data.auth.model_dump(), - token = None, - service_client = inbound_data.messageClient, - auth_type = "auth", - sync_freq = 300, - session_token = inbound_headers["X-Session-Token"] - ) + # Call the LLM and see if its service worked or not: + llm_response = await current_app.llm.invoke( + mongo_conn = current_app.data_mongo, + user_info = kwargs["session_info"], + llm_input = inbound_data + ) + success = False if llm_response.output is None else True # ┳┓ # ┣┫┏┓┏┏┓┏┓┏┓┏┏┓ @@ -216,12 +175,15 @@ async def request_oauth_authorization_url( # Done here: return ResponseModel( - status_code = StatusCodes.OK if token_id else StatusCodes.FAILED, - http_code = HttpCodes.SUCCESS if token_id else HttpCodes.INTERNAL_SERVER_ERROR, + status_code = StatusCodes.OK if success else StatusCodes.FAILED, + http_code = HttpCodes.SUCCESS if success else HttpCodes.INTERNAL_SERVER_ERROR, data = { - "messageClient": inbound_data.messageClient, - "authorized": True - } + "ts": llm_response.ts.isoformat(), + "client": llm_response.client, + "model": llm_response.model, + "output": llm_response.output, + "tokens": llm_response.tokens.model_dump(), + } if success else None ) diff --git a/api/blueprints/mail/oauth_callback.py b/api/blueprints/mail/oauth_callback.py index b773c96..e479d85 100644 --- a/api/blueprints/mail/oauth_callback.py +++ b/api/blueprints/mail/oauth_callback.py @@ -247,7 +247,7 @@ async def handle_gmail_callback() -> render_template: @log_chain_to_mongo(attr_name = "logs_mongo") @should_not_be_under_maintenance(attr_name = "is_under_maintenance") @handle_cancelled_request() -async def mail_callback( +async def mail_auth_callback( mail_client: str = None, inbound_headers: dict = None, inbound_data: dict = None, diff --git a/api/blueprints/mail/sync.py b/api/blueprints/mail/sync.py index 09add59..9bb83d4 100644 --- a/api/blueprints/mail/sync.py +++ b/api/blueprints/mail/sync.py @@ -126,6 +126,7 @@ def init(blueprint_setup_state): async def sync_mails( + user_info: dict, mongo_conn: AsyncMongo, llm: ChatOpenAI, inbound_headers: dict, @@ -145,6 +146,7 @@ async def sync_mails( # Try to sync the mails: return await current_app.mail_sync_model.sync( session_token = inbound_headers["X-Session-Token"], + user_info = user_info, mongo_conn = mongo_conn, token_id = inbound_data.tokenId, llm = llm, @@ -213,6 +215,7 @@ async def sync_mail( if mode in ["background", "bg"]: current_app.add_background_task( sync_mails, + user_info = kwargs["session_info"], mongo_conn = current_app.data_mongo, llm = current_app.llm, inbound_headers = inbound_headers, @@ -226,6 +229,7 @@ async def sync_mail( # Otherwise we process it right here: sync_results = await sync_mails( + user_info = kwargs["session_info"], mongo_conn = current_app.data_mongo, llm = current_app.llm, inbound_headers = inbound_headers, diff --git a/api/blueprints/sms/auth.py b/api/blueprints/sms/auth.py index 7693d40..b528c3e 100644 --- a/api/blueprints/sms/auth.py +++ b/api/blueprints/sms/auth.py @@ -180,7 +180,7 @@ async def request_oauth_authorization_url( }, auth = inbound_data.auth.model_dump(), token = None, - service_client = inbound_data.messageClient, + service_client = inbound_data.smsClient, auth_type = "auth", sync_freq = 300, session_token = inbound_headers["X-Session-Token"] @@ -203,7 +203,7 @@ async def request_oauth_authorization_url( }, auth = inbound_data.auth.model_dump(), token = None, - service_client = inbound_data.messageClient, + service_client = inbound_data.smsClient, auth_type = "auth", sync_freq = 300, session_token = inbound_headers["X-Session-Token"] @@ -219,7 +219,7 @@ async def request_oauth_authorization_url( status_code = StatusCodes.OK if token_id else StatusCodes.FAILED, http_code = HttpCodes.SUCCESS if token_id else HttpCodes.INTERNAL_SERVER_ERROR, data = { - "messageClient": inbound_data.messageClient, + "smsClient": inbound_data.messageClient, "authorized": True } ) diff --git a/api/main.py b/api/main.py index 8c37843..59487be 100644 --- a/api/main.py +++ b/api/main.py @@ -75,6 +75,7 @@ from models.behaviour.mail.oauth_v2 import MailOAuthModel from models.behaviour.mail.sync_v2 import MailSyncModel from models.behaviour.mail.retrieve import MailRetrieveModel from models.behaviour.sms.auth import SMSAuthModel +from models.behaviour.ai.llm.open_ai import LLMOpenAI # To make REST API calls: import httpx @@ -91,13 +92,11 @@ from api.blueprints.mail.retrieve import mail_retrieve_bp from api.blueprints.sms.auth import sms_auth_bp from api.blueprints.tech.chat_alerts import tech_chat_alert_bp from api.blueprints.test.callback import test_callback_bp +from api.blueprints.ai.llm.invoke import llm_invoke_bp # All the helpers: from api.helpers.user import session -# To work with LLMs: -from langchain_openai import ChatOpenAI - # ***************************************************************************************************************** # ***** **** @@ -129,6 +128,7 @@ app.register_blueprint(mail_retrieve_bp, url_prefix = f"/{MODULE_BASE}/mail") app.register_blueprint(sms_auth_bp, url_prefix = f"/{MODULE_BASE}/sms") app.register_blueprint(tech_chat_alert_bp, url_prefix = f"/{MODULE_BASE}/tech/alert") app.register_blueprint(test_callback_bp, url_prefix = f"/{MODULE_BASE}/test") +app.register_blueprint(llm_invoke_bp, url_prefix = f"/{MODULE_BASE}/ai") # ***************************************************************************************************************** @@ -200,7 +200,7 @@ async def app_startup(**kwargs): pool = 120.0, # .... Time to wait for a free connection from the pool. connect = 2.5, # ... Time to wait for establishing a connection to the server. write = 10.0, # .... Time to wait for sending data. - read = 2.5 # ....... Time to wait for receiving data. + read = 9.9 # ....... Time to wait for receiving data. ) ) @@ -374,9 +374,17 @@ async def app_startup(**kwargs): # ┛ # For LLMs: - current_app.llm = ChatOpenAI( - model = script_cred["openAi"]["model"], - openai_api_key = script_cred["openAi"]["openai_api_key"] + current_app.llm = LLMOpenAI( + llm_creds = { + "model": script_cred["openAi"]["model"], + "openai_api_key": script_cred["openAi"]["openai_api_key"] + }, + cache = current_app.module_cache, + alert_url = current_app.script_data["alerts"]["url"], + http_client = current_app.http_client, + debug = enable_debugging, + debug_prefix = "AI (LLM) | ", + debug_only_errors = True ) current_app.printer("AI ready.") diff --git a/models/behaviour/ai/llm/open_ai.py b/models/behaviour/ai/llm/open_ai.py index 8216df9..b8ad019 100644 --- a/models/behaviour/ai/llm/open_ai.py +++ b/models/behaviour/ai/llm/open_ai.py @@ -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()) diff --git a/models/behaviour/mail/retrieve.py b/models/behaviour/mail/retrieve.py index a0f01c9..55af60c 100644 --- a/models/behaviour/mail/retrieve.py +++ b/models/behaviour/mail/retrieve.py @@ -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 diff --git a/models/behaviour/mail/sync_v2.py b/models/behaviour/mail/sync_v2.py index fd93f4b..451d88b 100644 --- a/models/behaviour/mail/sync_v2.py +++ b/models/behaviour/mail/sync_v2.py @@ -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, diff --git a/models/data/ai/llm.py b/models/data/ai/llm.py index b7d0faf..4df854e 100644 --- a/models/data/ai/llm.py +++ b/models/data/ai/llm.py @@ -10,7 +10,7 @@ OBJECTIVE: - To provide a structure to receive auth details of various SMS providers. + To provide a structure to normalize input to and output from a standardized LLM wrapper. REFERENCES: @@ -36,8 +36,8 @@ sys.path.append(".") sys.path.append("..") # For making data behaviour_models: -from pydantic import BaseModel, Field, field_validator, PastDatetime -from typing import Optional, Literal, Union +from pydantic import BaseModel, Field, field_validator, PastDatetime, AwareDatetime +from typing import Optional, Literal, Union, List # My utils: from utils_v2.string import regex @@ -75,30 +75,15 @@ REGEX_SESSION_TOKEN = r"^[a-f0-9]{8}-[a-f0-9]{4}-[1-5][a-f0-9]{3}-[89ab][a-f0-9] # ***************************************************************************************************************** -class NimbusSMSIndiaAuth(BaseModel): +class LLMInputMessage(BaseModel): - entityId: str = Field( - description = "the entity id as registered with DLT", - min_length = 1, + role: Literal["system", "ai", "human"] = Field( + description = "the role of this message", frozen = True ) - senderId: str = Field( - description = "the 6-char code that you see in your SMS inbox", - min_length = 1, - frozen = True, - examples = ["HDFCBK", "NSESMS", "ZRODHA"] - ) - - userId: str = Field( - description = "the 6-digit id that Nimbus has assigned to you", - min_length = 1, - frozen = True - ) - - apiKey: str = Field( - description = "the key generated through Nimbus's portal", - min_length = 1, + content: str = Field( + description = "the message sent by the 'role'", frozen = True ) @@ -114,23 +99,63 @@ class NimbusSMSIndiaAuth(BaseModel): # --------------------------------------------------------------------------------------------------------------------- -class SavvyBulkSMSKenyaAuth(BaseModel): +class LLMInput(BaseModel): - apiKey: str = Field( - description = "the key generated through Savvy's portal", - min_length = 1, + messages: List[LLMInputMessage] + + # ┏┓ ┏• + # ┃ ┏┓┏┓╋┓┏┓ + # ┗┛┗┛┛┗┛┗┗┫ + # ┛ + + class Config: + extra = "forbid" + + # ┓┏ ┓• ┓ • + # ┃┃┏┓┃┓┏┫┏┓╋┓┏┓┏┓ + # ┗┛┗┻┗┗┗┻┗┻┗┗┗┛┛┗ + + @field_validator("messages") + def validate_messages(cls, value): + + # Maintain counter(s): + system_message_index = -1 + system_message_count = 0 + + # Loop through the messages and check them: + for index, message in enumerate(value): + + # For 'system' messages: + if message.role == "system": + system_message_index = index + system_message_count += 1 + + # Verify that there is AT MOST ONE 'system' message, + # and verify that the 'system' message is the first message: + if system_message_count > 1: raise ValueError(f"there can be at most 1 'system' message, found {system_message_count}") + if system_message_index > 0: raise ValueError(f"'system' message must always be at index 0, found it at index {system_message_index}") + + # Done here: + return value + + +# --------------------------------------------------------------------------------------------------------------------- + + +class LLMUsageTokens(BaseModel): + + input: int = Field( + description = "how many tokens were given in the input", frozen = True ) - partnerId: str = Field( - description = "the key generated through Savvy's portal", - min_length = 1, + output: int = Field( + description = "how many tokens were generated as the output", frozen = True ) - shortCode: str = Field( - description = "your short code with Savvy", - min_length = 1, + total: int = Field( + description = "the sum of the input and output tokens", frozen = True ) @@ -146,7 +171,54 @@ class SavvyBulkSMSKenyaAuth(BaseModel): # --------------------------------------------------------------------------------------------------------------------- -class SMSAuthRequestHeaders(BaseModel): +class LLMOutput(BaseModel): + + ts: AwareDatetime = Field( + description = "the time at which the llm was invoked", + default_factory = date_time.get_current_utc_date_time, + frozen = True + ) + + messages: List[LLMInputMessage] = Field( + description = "the messages that came in that invoked the llm", + frozen = True + ) + + output: str | None = Field( + description = "what the llm generated", + default = None, + frozen = True + ) + + client: Literal["openai"] = Field( + description = "the co./brand that was used to use an llm", + frozen = True + ) + + model: str = Field( + description = "to know which model used in the process", + frozen = True + ) + + tokens: LLMUsageTokens = Field( + description = "to know how many tokens were used in the process", + default = LLMUsageTokens(input = 0, output = 0, total = 0), + frozen = True + ) + + # ┏┓ ┏• + # ┃ ┏┓┏┓╋┓┏┓ + # ┗┛┗┛┛┗┛┗┗┫ + # ┛ + + class Config: + extra = "forbid" + + +# --------------------------------------------------------------------------------------------------------------------- + + +class LLMRequestHeaders(BaseModel): sessionToken: str = Field( description = "the session token of the user who is requesting the service", @@ -167,23 +239,6 @@ class SMSAuthRequestHeaders(BaseModel): return super().model_dump(*args, by_alias = True, **kwargs) -# --------------------------------------------------------------------------------------------------------------------- - - -class SMSAuthRequestData(BaseModel): - - messageClient: Literal["nimbusSmsIndia", "savvyBulkSmsKenya"] - auth: Union[NimbusSMSIndiaAuth, SavvyBulkSMSKenyaAuth] - - # ┏┓ ┏• - # ┃ ┏┓┏┓╋┓┏┓ - # ┗┛┗┛┛┗┛┗┗┫ - # ┛ - - class Config: - extra = "forbid" - - # ***************************************************************************************************************** # ***** **** # *** MAIN PROGRAM *** @@ -193,4 +248,20 @@ class SMSAuthRequestData(BaseModel): if __name__ == "__main__": - pass + 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..." + } + ] + + llm_input = LLMInput(messages = llm_messages) + print(llm_input) diff --git a/models/data/sms/auth.py b/models/data/sms/auth.py index b7d0faf..15223fe 100644 --- a/models/data/sms/auth.py +++ b/models/data/sms/auth.py @@ -172,7 +172,7 @@ class SMSAuthRequestHeaders(BaseModel): class SMSAuthRequestData(BaseModel): - messageClient: Literal["nimbusSmsIndia", "savvyBulkSmsKenya"] + smsClient: Literal["nimbusSmsIndia", "savvyBulkSmsKenya"] auth: Union[NimbusSMSIndiaAuth, SavvyBulkSMSKenyaAuth] # ┏┓ ┏•