### Set up environment variables:
#
# export IMUN_URL="https://cognitivewudev.azure-api.net/computervision/imageanalysis:analyze"
# export IMUN_SUBSCRIPTION_KEY=a*
# export OPENAI_API_KEY=sk-*

from langchain import ConversationChain, LLMChain

from langchain.agents import load_tools, initialize_agent
from langchain.chains.conversation.memory import ConversationBufferMemory
from langchain.llms import OpenAI

MAX_TOKENS = 512

llm = OpenAI(temperature=0, max_tokens=MAX_TOKENS)

tool_names = ['pal-math', 'imun']
tools = load_tools(tool_names, llm=llm)

memory = ConversationBufferMemory(memory_key="chat_history")

chain = initialize_agent(tools, llm, agent="conversational-react-description", verbose=True, memory=memory)

text = "what is the meaning of life"
output = chain.run(input=text)

text = "if I have two red balls and a blue ball, with blue balls half as heavy as the red balls. How many more blue balls do I need to have equal weight blue and red balls"
output = chain.run(input=text)

text = "summarize what you see in this image https://upload.wikimedia.org/wikipedia/commons/thumb/a/ad/Football_in_Bloomington%2C_Indiana%2C_1996.jpg/1920px-Football_in_Bloomington%2C_Indiana%2C_1996.jpg"
output = chain.run(input=text)


# To run imun as a tool
from langchain.utilities import ImunAPIWrapper
imun = ImunAPIWrapper()
print(imun.run("https://upload.wikimedia.org/wikipedia/commons/thumb/a/ad/Football_in_Bloomington%2C_Indiana%2C_1996.jpg/1920px-Football_in_Bloomington%2C_Indiana%2C_1996.jpg"))