import os

from datetime import datetime, timezone
from dotenv import load_dotenv
from langchain.agents import AgentType, initialize_agent
from langchain.chat_models import ChatOpenAI
from langchain.memory import ConversationBufferMemory
from langchain.prompts import MessagesPlaceholder

from src.xm_group_tools import GroupDetailsTool, GroupPerformanceTool, OneGroupPerformanceTool
from src.xmapi import XMApi

# import langchain
# langchain.debug = True

load_dotenv()
api_key = os.getenv("OPENAI_API_KEY")

# larger context model works better but costs a lot more
# MODEL = "gpt-3.5-turbo-16k-0613"
MODEL = "gpt-3.5-turbo-0613"

MEMORY_KEY = "memory"

tools = [GroupPerformanceTool(), OneGroupPerformanceTool(), GroupDetailsTool()]
llm = ChatOpenAI(model_name = MODEL)
memory = ConversationBufferMemory(memory_key=MEMORY_KEY, return_messages=True)

# give the conversation memory the current date
memory.save_context({"input": "What is today's date?"},{"output": f"Today is {datetime.now(timezone.utc)}"})
memory.save_context({"input": "Main attribute of performance is the ratio of response and the lowest time to first response?"},{"output": f""})
memory.save_context({"input": "timeToFirstResponse is in milliseconds."},{"output": f""})

agent_kwargs = {
    "extra_prompt_messages": [MessagesPlaceholder(variable_name = MEMORY_KEY)],
    "temperature": 0
}

agent_chain = initialize_agent(
    tools,
    llm,
    agent = AgentType.OPENAI_FUNCTIONS,
    verbose = True,
    agent_kwargs = agent_kwargs,
    memory = memory)

while True:
    user_input = input("User: ")
    if not user_input or user_input.lower() == "quit":
        break
    agent_chain.run(input = user_input)

langchain.agents.initialize_agent langchain.memory.ConversationBufferMemory langchain.prompts.MessagesPlaceholder langchain.chat_models.ChatOpenAI
[((458, 471), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (469, 471), False, 'from dotenv import load_dotenv\n'), ((482, 509), 'os.getenv', 'os.getenv', (['"""OPENAI_API_KEY"""'], {}), "('OPENAI_API_KEY')\n", (491, 509), False, 'import os\n'), ((742, 770), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': 'MODEL'}), '(model_name=MODEL)\n', (752, 770), False, 'from langchain.chat_models import ChatOpenAI\n'), ((782, 851), 'langchain.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {'memory_key': 'MEMORY_KEY', 'return_messages': '(True)'}), '(memory_key=MEMORY_KEY, return_messages=True)\n', (806, 851), False, 'from langchain.memory import ConversationBufferMemory\n'), ((1380, 1502), 'langchain.agents.initialize_agent', 'initialize_agent', (['tools', 'llm'], {'agent': 'AgentType.OPENAI_FUNCTIONS', 'verbose': '(True)', 'agent_kwargs': 'agent_kwargs', 'memory': 'memory'}), '(tools, llm, agent=AgentType.OPENAI_FUNCTIONS, verbose=True,\n    agent_kwargs=agent_kwargs, memory=memory)\n', (1396, 1502), False, 'from langchain.agents import AgentType, initialize_agent\n'), ((665, 687), 'src.xm_group_tools.GroupPerformanceTool', 'GroupPerformanceTool', ([], {}), '()\n', (685, 687), False, 'from src.xm_group_tools import GroupDetailsTool, GroupPerformanceTool, OneGroupPerformanceTool\n'), ((689, 714), 'src.xm_group_tools.OneGroupPerformanceTool', 'OneGroupPerformanceTool', ([], {}), '()\n', (712, 714), False, 'from src.xm_group_tools import GroupDetailsTool, GroupPerformanceTool, OneGroupPerformanceTool\n'), ((716, 734), 'src.xm_group_tools.GroupDetailsTool', 'GroupDetailsTool', ([], {}), '()\n', (732, 734), False, 'from src.xm_group_tools import GroupDetailsTool, GroupPerformanceTool, OneGroupPerformanceTool\n'), ((1292, 1337), 'langchain.prompts.MessagesPlaceholder', 'MessagesPlaceholder', ([], {'variable_name': 'MEMORY_KEY'}), '(variable_name=MEMORY_KEY)\n', (1311, 1337), False, 'from langchain.prompts import MessagesPlaceholder\n'), ((979, 1005), 'datetime.datetime.now', 'datetime.now', (['timezone.utc'], {}), '(timezone.utc)\n', (991, 1005), False, 'from datetime import datetime, timezone\n')]