224 lines
8.4 KiB
Python
224 lines
8.4 KiB
Python
"""
|
|
|
|
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
|
|
from models.data.core.user_info import CoreUserInfoModel
|
|
|
|
# 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: CoreUserInfoModel,
|
|
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.model_dump()}
|
|
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
|
|
)
|
|
if inserted_id: llm_response.invocationId = str(inserted_id)
|
|
|
|
# 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())
|