import gradio as gr
from flask import jsonify, request
import pandas as pd
from langchain.utilities import SerpAPIWrapper
from langchain.agents import Tool
from langchain.memory import ConversationBufferMemory
from langchain.llms import Cohere
from langchain.agents import initialize_agent
from langchain.agents.agent_types import AgentType

COHERE_API_KEY = "CggzsdnWH6QXtnGJvKYe4IRZyGZ8UkTSykpmAigW"
SERPAPI_API_KEY = "dbc53dd88c7b0957548a81fa162e2d547e03cc19267162a9166e52d4e882f361"


def get_chat_output(prompt: str):
    try:
        search = SerpAPIWrapper(serpapi_api_key=SERPAPI_API_KEY)
        tools = [
            Tool(
                name="Current search",
                func=search.run,
                description="useful for when you need to answer questions about current events or the current state of the world",
            ),
        ]
        memory = ConversationBufferMemory(memory_key="chat_history")
        input = prompt 
        llm = Cohere(cohere_api_key=COHERE_API_KEY, model="command-xlarge-nightly")
        agent_chain = initialize_agent(
            tools,
            llm,
            agent=AgentType.CONVERSATIONAL_REACT_DESCRIPTION,
            verbose=True,
            memory=memory,
            handle_parsing_errors=True,
        )
        response = agent_chain.run(input=input)
        return {"response": response}

    except:
        return jsonify({"error": "Facing errors, please try again..."})


iface = gr.Interface(fn=get_chat_output, inputs="text", outputs="json")
iface.launch()