import os
import gradio as gr

from langchain_community.utilities.sql_database import SQLDatabase
from langchain_community.agent_toolkits import create_sql_agent

from langchain_openai import ChatOpenAI

ccms_db_loc = 'ccms.db'

ccms_db = SQLDatabase.from_uri(f"sqlite:///{ccms_db_loc}")

gpt4o = ChatOpenAI(
    model_name='gpt-4o',
    api_key=os.environ["OPENAI_API_KEY"],
    temperature=0
)

gpt35 = ChatOpenAI(
    model_name='gpt-3.5-turbo',
    api_key=os.environ["OPENAI_API_KEY"],
    temperature=0
)

sqlite_agent = create_sql_agent(
    gpt4o,
    db=ccms_db,
    agent_type="openai-tools",
    agent_executor_kwargs={'handle_parsing_errors':True},
    verbose=True
)

def predict(user_input):

    try:
        response = sqlite_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 for the database is presented below: \n <img src="https://cdn-uploads.huggingface.co/production/uploads/64118e60756b9e455c7eddd6/S1alVt_D88qatd-N4Dkjd.png" > \n<img src="https://cdn-uploads.huggingface.co/production/uploads/64118e60756b9e455c7eddd6/81ggHEjrt6wFrMyXJtHVS.png" > (Source: https://github.com/shrivastavasatyam/Credit-Card-Management-System)'

demo = gr.Interface(
    inputs=textbox, fn=predict, outputs="text",
    title="Query a Credit Card Database",
    description="This web API presents an interface to ask questions on information stored in a credit card database.",
    article=schema,
    examples=[["Who are the top 5 merchants by total transactions?", ""]],
    cache_examples=False,
    concurrency_limit=8
)


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