import os
from langchain import OpenAI, SQLDatabase, SQLDatabaseChain

import pandas as pd
import sqlite3
import streamlit as st

API_KEY = os.getenv('OPENAI_API_KEY')

st.title("English to SQL via LangChain")

tables = ['Album', 'Artist', 'Track']
db = SQLDatabase.from_uri("sqlite:///Chinook.db", include_tables=tables)

con = sqlite3.connect("Chinook.db")
cur = con.cursor()

metadata = dict()
for table in tables:

    rows = cur.execute("select * from %s limit 1" % table)
    cols = [k[0] for k in rows.description]

    metadata[table] = [cols]

con.close()

llm = OpenAI(temperature=0.0, openai_api_key=API_KEY)

db_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True, use_query_checker=True)

queries = ("How many albums are there?"
    , "Which album has the most tracks?"
    , "What artist has the album with the most tracks?"
)

for query in queries:
    result = db_chain.run(query)
    print(result)
    st.text(query)
    st.text(result)

st.subheader("table metadata")
st.dataframe(pd.DataFrame(metadata), columns=['table', 'columns'])