import os
from langchain.agents import initialize_agent
from langchain.llms import OpenAI
from gradio_tools import (StableDiffusionTool, ImageCaptioningTool, StableDiffusionPromptGeneratorTool,
                          TextToVideoTool)
from langchain.memory import ConversationBufferMemory

# Set the API key
os.environ["OPENAI_API_KEY"] = "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"

llm = OpenAI(api_key=os.getenv("OPENAI_API_KEY"), temperature=0)
memory = ConversationBufferMemory(memory_key="chat_history")
tools = [StableDiffusionTool().langchain, ImageCaptioningTool().langchain,
         StableDiffusionPromptGeneratorTool().langchain, TextToVideoTool().langchain]


agent = initialize_agent(tools, llm, memory=memory, agent="conversational-react-description", verbose=True)
output = agent.run(input=("Please create a photo of a dog riding a skateboard "
                          "but improve my prompt prior to using an image generator."
                          "Please caption the generated image and create a video for it using the improved prompt."))
print(output)