import os
from langchain.agents import Tool, load_tools, initialize_agent
from langchain.memory import ConversationBufferMemory
from langchain.chat_models import ChatOpenAI
from langchain.utilities import GoogleSearchAPIWrapper
from langchain.agents import initialize_agent
from langchain.utilities import WikipediaAPIWrapper
from langchain.utilities.wolfram_alpha import WolframAlphaAPIWrapper
import gradio as gr

description = """
👍 If you liked the app, sign up for my AI newsletter where I'll be sharing prompts, news, and tools like this one: [https://slothio.com/ai-newsletter/](https://slothio.com/ai-newsletter/)
"""

def initialize_keys(openai_api_key, google_api_key, google_cse_id, wolfram_alpha_appid):
    os.environ['OPENAI_API_KEY'] = openai_api_key
    os.environ['GOOGLE_API_KEY'] = google_api_key
    os.environ['GOOGLE_CSE_ID'] = google_cse_id
    os.environ['WOLFRAM_ALPHA_APPID'] = wolfram_alpha_appid

def chat_response(input_text, openai_api_key, google_api_key, google_cse_id, wolfram_alpha_appid):
    initialize_keys(openai_api_key, google_api_key, google_cse_id, wolfram_alpha_appid)

    llm=ChatOpenAI(temperature=1)
    search = GoogleSearchAPIWrapper()
    tools = load_tools(["google-search", "wolfram-alpha", "wikipedia"], llm=llm)
    memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

    agent_chain = initialize_agent(tools, llm, agent="zero-shot-react-description",
                                   verbose=True, memory=memory)
    response = agent_chain.run(input=input_text)
    return response

input_text = gr.inputs.Textbox(lines=2, label='Input Text')
openai_api_key = gr.inputs.Textbox(lines=1, label='OpenAI API Key', type='password')
google_api_key = gr.inputs.Textbox(lines=1, label='Google API Key', type='password')
google_cse_id = gr.inputs.Textbox(lines=1, label='Google CSE ID', type='password')
wolfram_alpha_appid = gr.inputs.Textbox(lines=1, label='Wolfram Alpha AppID', type='password')

interface = gr.Interface(fn=chat_response, inputs=[input_text, openai_api_key, google_api_key, google_cse_id, wolfram_alpha_appid], outputs="text", description=description)

interface.launch(share=False)