import os
import gradio as gr

from langchain.agents.agent_types import AgentType
from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent

from langchain_openai import AzureChatOpenAI

from sklearn.datasets import fetch_openml


gpt35 = AzureChatOpenAI(
    api_key=os.environ["AZURE_OPENAI_KEY"],
    azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
    api_version="2023-05-15",
    deployment_name="gpt-35-turbo"
)

bank_data, _ = fetch_openml(data_id=43718, return_X_y=True)

pandas_agent = create_pandas_dataframe_agent(
    llm=gpt35,
    df=bank_data,
    verbose=True,
    agent_type=AgentType.OPENAI_FUNCTIONS,
)


def predict(user_input):

    try:
        response = pandas_agent.invoke(user_input)

        prediction = response['output']

    except Exception as e:
        prediction = e
        
    return prediction


textbox = gr.Textbox(placeholder="Enter your query here", lines=6)

interface = gr.Interface(
    inputs=textbox, fn=predict, outputs="text",
    title="Query BFSI customer information",
    description="This web API presents an interface to ask questions on customer information stored in a database.",
    examples=[["What is the average balance maintained by the users?", ""],
              ["How many users have subscribed to a term deposit?", ""]
             ]
)

with gr.Blocks() as demo:
    interface.launch()

demo.queue(concurrency_count=16)
demo.launch()