import gradio as gr
import os

from langchain.agents import load_tools
from langchain.agents import initialize_agent
from langchain.agents import AgentType
from langchain.llms import OpenAI

def get_response(input, google_cse_id, google_api_key, openai_api_key):
    os.environ["GOOGLE_CSE_ID"] = google_cse_id
    os.environ["GOOGLE_API_KEY"] = google_api_key
    os.environ["OPENAI_API_KEY"] = openai_api_key

    llm = OpenAI(temperature=0)
    tools = load_tools(["google-search", "llm-math"], llm=llm)
    agent = initialize_agent(tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True, return_intermediate_steps=True)

    response = agent({"input": input})
    return response["output"], response["intermediate_steps"]

iface = gr.Interface(
    fn=get_response,
    inputs=[
        gr.Textbox(label="Goal definition:"),
        gr.Textbox(label="Google CSE - ID:", type="text"),
        gr.Textbox(label="Google API KEY:", type="password"),
        gr.Textbox(label="OpenAI API KEY:", placeholder="sk-...", type="password")
    ],
    outputs=[
        gr.Textbox(label="Goal output:"),
        gr.Json(label="Intermediate Steps")
    ],
    title="GPT Agents Demo",
    description="Demo application of gpt-based agents including two tools (google-search and llm-math). The result and intermediate steps are included."
)

# error capturing in integration as a component

error_message = ""

try:
    iface.queue(concurrency_count=20)
    iface.launch()
except Exception as e:
    error_message = "An error occurred: " + str(e)
    iface.outputs[1].value = error_message