Update script.py

script.py

@@ -4,6 +4,8 @@ from langchain import OpenAI, SQLDatabase, SQLDatabaseChain
 import sqlite3
 import streamlit as st
 
+API_KEY = os.getenv('OPENAI_API_KEY')
+
 st.title("English to SQL via LangChain")
 
 tables = ['Album', 'Artist', 'Track']
@@ -12,8 +14,30 @@ 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 5" % table).fetchall()
-    st.table(rows)
 
-con.close()
+    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'])