""" A simple cloud consultant bot that can answer questions 
about kubernetes, aws and cloud native."""
import langchain
from langchain.agents import Tool, AgentType, initialize_agent
from langchain.tools import HumanInputRun
from langchain.callbacks import HumanApprovalCallbackHandler
from langchain.vectorstores import Chroma
from langchain.chat_models import ChatOpenAI
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
from termcolor import colored
from cloud_tool import CloudTool
from approval import ApprovalCallBackHandler

langchain.debug = False

MODEL = "gpt-3.5-turbo"

cloud_tool = CloudTool(callbacks=[ApprovalCallBackHandler()])
cloud_tool.description = cloud_tool.description + f"args {cloud_tool.args}".replace(
    "{", "{{"
).replace("}", "}}")

human = HumanInputRun()


llm = ChatOpenAI(temperature=0, model=MODEL)
embeddings = OpenAIEmbeddings()
vectorstore = Chroma(persist_directory="./", embedding_function=embeddings)
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
kubememory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
awsmemory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

docs = ConversationalRetrievalChain.from_llm(
    llm, vectorstore.as_retriever(), memory=memory
)

cloud_tools = [cloud_tool, human]
kubectl_agent_chain = initialize_agent(
    tools=cloud_tools,
    llm=llm,
    agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
    memory=kubememory,
    verbose=False,
    agent_kwargs={
        "prefix": """
You are a Kubernetes Command line tool (kubectl) expert. 
Given an input question, first create a syntactically correct kubectl command to run, then look at the results of the command and return the answer to the input question.
If there is no namespace name given please use the "default" namespace.
Only return the command. If an error is returned, rewrite the command, check the command, and try again.

""",
    },
)

aws_agent_chain = initialize_agent(
    cloud_tools,
    llm,
    agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
    memory=awsmemory,
    verbose=False,
    agent_kwargs={
        "prefix": """
You are a AWS Command line tool (aws cli) expert. 
Given an input question, first create a syntactically correct aws cli command to run, then look at the results of the query and return the answer to the input question.
You must generate the correct aws cli command to answer he question. 
Only return the command. If an error is returned, rewrite the command, check the command, and try again.
""",
    },
)


tools = [
    Tool(
        name="Kubernetes QA System",
        func=docs.run,
        description="useful for when you need to answer questions about kubernetes or cloud native and from the kubernetes or cloud native documentation. input should be a fully formed question.",
    ),
    Tool(
        name="Kubectl",
        func=kubectl_agent_chain.run,
        description="useful for when you need to use kubectl to look up, change or update your kubernetes cluster.",
    ),
    Tool(
        name="Aws CLI",
        func=aws_agent_chain.run,
        description="useful for when you need to use aws cli to look up, change or update your AWS setup.",
    ),
    human,
]


agent_chain = initialize_agent(
    tools,
    llm,
    agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
    verbose=False,
    memory=memory,
)


def ask_ai():
    """Main method to talk to the ai"""
    print(
        colored(
            "Welcome, i am Your AI cloud consultant. How can i help You today?", "green"
        )
    )
    try:
        while True:
            query = input(colored("You: ", "white", attrs=["bold"]))
            result = agent_chain.run(input=query)
            print(
                colored("Answer: ", "green", attrs=["bold"]),
                colored(result, "light_green"),
            )
    except (EOFError, KeyboardInterrupt):
        print("kthxbye")
        exit()


if __name__ == "__main__":
    ask_ai()

langchain.agents.initialize_agent langchain.tools.HumanInputRun langchain.memory.ConversationBufferMemory langchain.embeddings.openai.OpenAIEmbeddings langchain.chat_models.ChatOpenAI langchain.agents.Tool langchain.vectorstores.Chroma
[((893, 908), 'langchain.tools.HumanInputRun', 'HumanInputRun', ([], {}), '()\n', (906, 908), False, 'from langchain.tools import HumanInputRun\n'), ((917, 955), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'temperature': '(0)', 'model': 'MODEL'}), '(temperature=0, model=MODEL)\n', (927, 955), False, 'from langchain.chat_models import ChatOpenAI\n'), ((969, 987), 'langchain.embeddings.openai.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (985, 987), False, 'from langchain.embeddings.openai import OpenAIEmbeddings\n'), ((1002, 1063), 'langchain.vectorstores.Chroma', 'Chroma', ([], {'persist_directory': '"""./"""', 'embedding_function': 'embeddings'}), "(persist_directory='./', embedding_function=embeddings)\n", (1008, 1063), False, 'from langchain.vectorstores import Chroma\n'), ((1073, 1146), 'langchain.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {'memory_key': '"""chat_history"""', 'return_messages': '(True)'}), "(memory_key='chat_history', return_messages=True)\n", (1097, 1146), False, 'from langchain.memory import ConversationBufferMemory\n'), ((1160, 1233), 'langchain.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {'memory_key': '"""chat_history"""', 'return_messages': '(True)'}), "(memory_key='chat_history', return_messages=True)\n", (1184, 1233), False, 'from langchain.memory import ConversationBufferMemory\n'), ((1246, 1319), 'langchain.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {'memory_key': '"""chat_history"""', 'return_messages': '(True)'}), "(memory_key='chat_history', return_messages=True)\n", (1270, 1319), False, 'from langchain.memory import ConversationBufferMemory\n'), ((1477, 2069), 'langchain.agents.initialize_agent', 'initialize_agent', ([], {'tools': 'cloud_tools', 'llm': 'llm', 'agent': 'AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION', 'memory': 'kubememory', 'verbose': '(False)', 'agent_kwargs': '{\'prefix\':\n    """\nYou are a Kubernetes Command line tool (kubectl) expert. \nGiven an input question, first create a syntactically correct kubectl command to run, then look at the results of the command and return the answer to the input question.\nIf there is no namespace name given please use the "default" namespace.\nOnly return the command. If an error is returned, rewrite the command, check the command, and try again.\n\n"""\n    }'}), '(tools=cloud_tools, llm=llm, agent=AgentType.\n    CHAT_CONVERSATIONAL_REACT_DESCRIPTION, memory=kubememory, verbose=False,\n    agent_kwargs={\'prefix\':\n    """\nYou are a Kubernetes Command line tool (kubectl) expert. \nGiven an input question, first create a syntactically correct kubectl command to run, then look at the results of the command and return the answer to the input question.\nIf there is no namespace name given please use the "default" namespace.\nOnly return the command. If an error is returned, rewrite the command, check the command, and try again.\n\n"""\n    })\n', (1493, 2069), False, 'from langchain.agents import Tool, AgentType, initialize_agent\n'), ((2113, 2682), 'langchain.agents.initialize_agent', 'initialize_agent', (['cloud_tools', 'llm'], {'agent': 'AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION', 'memory': 'awsmemory', 'verbose': '(False)', 'agent_kwargs': '{\'prefix\':\n    """\nYou are a AWS Command line tool (aws cli) expert. \nGiven an input question, first create a syntactically correct aws cli command to run, then look at the results of the query and return the answer to the input question.\nYou must generate the correct aws cli command to answer he question. \nOnly return the command. If an error is returned, rewrite the command, check the command, and try again.\n"""\n    }'}), '(cloud_tools, llm, agent=AgentType.\n    CHAT_CONVERSATIONAL_REACT_DESCRIPTION, memory=awsmemory, verbose=False,\n    agent_kwargs={\'prefix\':\n    """\nYou are a AWS Command line tool (aws cli) expert. \nGiven an input question, first create a syntactically correct aws cli command to run, then look at the results of the query and return the answer to the input question.\nYou must generate the correct aws cli command to answer he question. \nOnly return the command. If an error is returned, rewrite the command, check the command, and try again.\n"""\n    })\n', (2129, 2682), False, 'from langchain.agents import Tool, AgentType, initialize_agent\n'), ((3401, 3519), 'langchain.agents.initialize_agent', 'initialize_agent', (['tools', 'llm'], {'agent': 'AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION', 'verbose': '(False)', 'memory': 'memory'}), '(tools, llm, agent=AgentType.\n    CHAT_CONVERSATIONAL_REACT_DESCRIPTION, verbose=False, memory=memory)\n', (3417, 3519), False, 'from langchain.agents import Tool, AgentType, initialize_agent\n'), ((2723, 2970), 'langchain.agents.Tool', 'Tool', ([], {'name': '"""Kubernetes QA System"""', 'func': 'docs.run', 'description': '"""useful for when you need to answer questions about kubernetes or cloud native and from the kubernetes or cloud native documentation. input should be a fully formed question."""'}), "(name='Kubernetes QA System', func=docs.run, description=\n    'useful for when you need to answer questions about kubernetes or cloud native and from the kubernetes or cloud native documentation. input should be a fully formed question.'\n    )\n", (2727, 2970), False, 'from langchain.agents import Tool, AgentType, initialize_agent\n'), ((2997, 3166), 'langchain.agents.Tool', 'Tool', ([], {'name': '"""Kubectl"""', 'func': 'kubectl_agent_chain.run', 'description': '"""useful for when you need to use kubectl to look up, change or update your kubernetes cluster."""'}), "(name='Kubectl', func=kubectl_agent_chain.run, description=\n    'useful for when you need to use kubectl to look up, change or update your kubernetes cluster.'\n    )\n", (3001, 3166), False, 'from langchain.agents import Tool, AgentType, initialize_agent\n'), ((3193, 3349), 'langchain.agents.Tool', 'Tool', ([], {'name': '"""Aws CLI"""', 'func': 'aws_agent_chain.run', 'description': '"""useful for when you need to use aws cli to look up, change or update your AWS setup."""'}), "(name='Aws CLI', func=aws_agent_chain.run, description=\n    'useful for when you need to use aws cli to look up, change or update your AWS setup.'\n    )\n", (3197, 3349), False, 'from langchain.agents import Tool, AgentType, initialize_agent\n'), ((3613, 3702), 'termcolor.colored', 'colored', (['"""Welcome, i am Your AI cloud consultant. How can i help You today?"""', '"""green"""'], {}), "('Welcome, i am Your AI cloud consultant. How can i help You today?',\n    'green')\n", (3620, 3702), False, 'from termcolor import colored\n'), ((736, 761), 'approval.ApprovalCallBackHandler', 'ApprovalCallBackHandler', ([], {}), '()\n', (759, 761), False, 'from approval import ApprovalCallBackHandler\n'), ((3782, 3823), 'termcolor.colored', 'colored', (['"""You: """', '"""white"""'], {'attrs': "['bold']"}), "('You: ', 'white', attrs=['bold'])\n", (3789, 3823), False, 'from termcolor import colored\n'), ((3910, 3954), 'termcolor.colored', 'colored', (['"""Answer: """', '"""green"""'], {'attrs': "['bold']"}), "('Answer: ', 'green', attrs=['bold'])\n", (3917, 3954), False, 'from termcolor import colored\n'), ((3972, 4002), 'termcolor.colored', 'colored', (['result', '"""light_green"""'], {}), "(result, 'light_green')\n", (3979, 4002), False, 'from termcolor import colored\n')]