import gradio as gr
from langchain.agents import initialize_agent
from langchain.llms import OpenAI
from gradio_tools.tools import (StableDiffusionTool, ImageCaptioningTool, StableDiffusionPromptGeneratorTool,
                                TextToVideoTool)

from langchain.memory import ConversationBufferMemory

llm = OpenAI(openai_api_key="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",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)

def run_text(text, state):
    output = agent.run(input=(text))
    print(memory)
    print(output)
    return [(text,output)]

with gr.Blocks(css="#chatbot {overflow:auto; height:500px;}") as demo:
    chatbot = gr.Chatbot(elem_id="chatbot",show_label=False)
    state = gr.State([])
    with gr.Row() as input_raws:
        with gr.Column(scale=0.6):
            txt = gr.Textbox(show_label=False).style(container=False)
        with gr.Column(scale=0.20, min_width=0):
            run = gr.Button("🏃‍♂️Run")
        with gr.Column(scale=0.20, min_width=0):
            clear = gr.Button("🔄Clear️")

    txt.submit(run_text, [txt, state], [chatbot])
    run.click(run_text, [txt, state], [chatbot])
 
demo.queue(concurrency_count=10).launch(server_name="0.0.0.0", server_port=7860)