import os 
import gradio as gr

from langchain.chat_models import ChatOpenAI
from langchain.agents import load_tools, initialize_agent
from langchain.agents import AgentType
from langchain.tools import AIPluginTool



def run(prompt, plugin_json, openai_api_key):
    os.environ["OPENAI_API_KEY"] = openai_api_key
    tool = AIPluginTool.from_plugin_url(plugin_json)
    llm = ChatOpenAI(temperature=0)
    tools = load_tools(["requests_all"] )
    tools += [tool]
    agent_chain = initialize_agent( tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True)
    return agent_chain.run(prompt)

title="""
<div style="text-align:center;">
<h1>LangChain + ChatGPT Plugins playground</h1>
<p>
This is a demo for the <a href="https://python.langchain.com/en/latest/modules/agents/tools/examples/chatgpt_plugins.html", target="_blank">ChatGPT Plugins LangChain</a> usecase<br />
Be aware that it currently only works with plugins that do not require auth.
</p>
</div>
"""
with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.HTML(title)
        prompt = gr.Textbox(label="Prompt", value="what t shirts are available in klarna?")
        plugin = gr.Textbox(label="Plugin json", info="You need the .json plugin manifesto file of the plugin you want to use", value="https://www.klarna.com/.well-known/ai-plugin.json")
        openai_api_key = gr.Textbox(label="OpenAI API Key", type="password")
        run_btn = gr.Button("Run")
        response = gr.Textbox(label="Response")
    run_btn.click(fn=run, 
                 inputs=[prompt, plugin, openai_api_key],
                 outputs=[response]
                 )
demo.launch()