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

from sklearn.datasets import fetch_openml

gpt4o_azure = AzureChatOpenAI(
    model_name='gpt-4o-mini',
    api_key=os.environ["AZURE_OPENAI_KEY"],
    azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
    api_version="2024-02-01",
    temperature=0
)

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

pandas_agent = create_pandas_dataframe_agent(
    llm=gpt4o_azure,
    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)

schema = """
The schema of the table you can query on is presented below:
| Column    | Type   | Description|
|-----------|--------|------------|
| age       | int    |Age of the customer|
| job       | str    |type of job (categorical: "admin.","unknown","unemployed","management","housemaid","entrepreneur","student", "blue-collar","self-employed","retired","technician","services") |
| marital   | str    |marital status (categorical: "married","divorced","single"; note: "divorced" means divorced or widowed)|
| education | str    |(categorical: "unknown","secondary","primary","tertiary")|
| default   | yes/no |has credit in default? (binary: "yes","no")|
| balance   | int    |average yearly balance, in euros (numeric)|
| housing   | yes/no |has housing loan? (binary: "yes","no")|
| loan      | yes/no |has personal loan? (binary: "yes","no")|
| contact   | str    |last contact communication type (categorical: "unknown","telephone","cellular")|
| day       | int    |last contact day of the month (numeric)|
| month     | str    |last contact month of year (categorical: "jan", "feb", "mar", ..., "nov", "dec")|
| duration  | int    |last contact duration, in seconds (numeric)|
| campaign  | int    |number of contacts performed during this campaign and for this client (numeric, includes last contact)|
| pdays     | int    |number of days that passed by after the client was last contacted from a previous campaign (numeric, -1 means client was not previously contacted)|
| poutcome  | str    |outcome of the previous marketing campaign (categorical: "unknown","other","failure","success")|
| deposit   | yes/no |has the client subscribed a term deposit? (binary: "yes","no")|
"""

demo = 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?", ""],
              ["For how many customers was our last marketing campaign a success?", ""],
              ["What is the most common marital status of our customers?"]
             ],
    cache_examples=False,
    theme="origin",
#    theme=gr.themes.Base(),
    concurrency_limit=8
)


demo.queue()
demo.launch(auth=("demouser", os.getenv('PASSWD')), ssr_mode=False)