(20241209) Better session info capture in the logs.

models/behaviour/core/auth_token.py

@@ -0,0 +1,221 @@
+"""
+
+    AUTHOR:
+
+        Khushal P Soonderji
+
+    DATE:
+
+        Thursday, 5th Dec., 2024
+
+    OBJECTIVE:
+
+        To create an interface between OpenAI and our internal system to perform LLM-based activities.
+
+    REFERENCES:
+
+        N/A
+
+    DOWNLOADS:
+
+        N/A
+
+"""
+
+#   *****************************************************************************************************************
+#  *****                                                                                                          ****
+# ***                                                   IMPORT                                                      ***
+#  *****                                                                                                          ****
+#   *****************************************************************************************************************
+
+
+# To make sibling directories accessible for imports:
+import sys
+sys.path.append(".")
+sys.path.append("..")
+
+# 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
+
+# Base model:
+from models.behaviour.base import BaseModel
+
+# Data Models:
+from models.data.api.ai.llm import LLMInput, LLMOutput, LLMUsageTokens
+
+# To work with LLMs:
+from langchain_openai import ChatOpenAI
+
+# To work with MongoDB:
+from bson import ObjectId
+
+# To work with datatypes:
+from typing import Literal
+
+# To make deep-copies:
+import copy
+
+
+#   *****************************************************************************************************************
+#  *****                                                                                                          ****
+# ***                                           MACROS / ONE-TIME INIT                                              ***
+#  *****                                                                                                          ****
+#   *****************************************************************************************************************
+
+
+# --- Nothing Yet
+
+
+#   *****************************************************************************************************************
+#  *****                                                                                                          ****
+# ***                                                VARIABLES                                                      ***
+#  *****                                                                                                          ****
+#   *****************************************************************************************************************
+
+
+# --- Nothing Yet
+
+
+#   *****************************************************************************************************************
+#  *****                                                                                                          ****
+# ***                                                  FUNCTIONS                                                    ***
+#  *****                                                                                                          ****
+#   *****************************************************************************************************************
+
+
+# --- Nothing Yet
+
+
+#   *****************************************************************************************************************
+#  *****                                                                                                          ****
+# ***                                                   CLASSES                                                     ***
+#  *****                                                                                                          ****
+#   *****************************************************************************************************************
+
+
+class LLMOpenAI(BaseModel):
+
+    AI_USAGE_COLLECTION = "_aiUsage"
+
+    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,
+            mongo_conn: AsyncMongo,
+            user_info: dict,
+            llm_input: LLMInput
+    ) -> LLMOutput:
+
+        # 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
+        ]
+
+        # 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"],
+            )
+        )
+
+        # 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 llm_response
+
+
+#   *****************************************************************************************************************
+#  *****                                                                                                          ****
+# ***                                               MAIN PROGRAM                                                    ***
+#  *****                                                                                                          ****
+#   *****************************************************************************************************************
+
+
+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())