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, ChatOpenAI

from sklearn.datasets import fetch_openml

#gpt35 = ChatOpenAI(
#    openai_api_key=os.environ["OPENAI_API_KEY"],
#    model="gpt-3.5-turbo",
#    temperature=0
#)

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",
    temperature=0
)

bank_data, _ = fetch_openml(data_id=43718, return_X_y=True, parser="auto")

pandas_agent = create_pandas_dataframe_agent(
    llm=gpt35,
    df=bank_data,
    verbose=False,
    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)

schema = """
The schema of the table you can query on is presented below:
| Column    | Type   |
|-----------|--------|
| age       | int    |
| job       | str    |
| marital   | str    |
| education | str    |
| default   | yes/no |
| balance   | int    |
| housing   | yes/no |
| loan      | yes/no |
| contact   | str    |
| day       | int    |
| month     | str    |
| duration  | int    |
| campaign  | int    |
| pdays     | int    |
| poutcome  | str    |
| deposit   | yes/no |
"""

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.",
    article=schema,
    examples=[["What is the average balance maintained by our customers?", ""],
              ["How many customers have subscribed to a term deposit?", ""],
              ["How many customers have defaulted on loans?", ""],
              ["Do customers who default maintain a low balance?", ""]
             ],
    concurrency_limit=16
)

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

demo.queue()
demo.launch()