Update script.py

script.py

@@ -7,38 +7,46 @@ import streamlit as st
 
 API_KEY = os.getenv('OPENAI_API_KEY')
 
-st.title("English to SQL via LangChain")
+TABLES = ['Album', 'Artist', 'Track']
 
-tables = ['Album', 'Artist', 'Track']
-db = SQLDatabase.from_uri("sqlite:///Chinook.db", include_tables=tables)
+QUERIES = ("How many albums are there?"
+    , "Which album has the most tracks?"
+    , "What artist has the album with the 2nd most tracks?"
+)
 
-con = sqlite3.connect("Chinook.db")
-cur = con.cursor()
+def _get_metadata()
 
-st.subheader("table metadata")
-for table in tables:
+    con = sqlite3.connect("Chinook.db")
+    cur = con.cursor()
+
+    ts, cs = list(), list()
+    for table in tables:
 
-    rows = cur.execute("select * from %s limit 1" % table)
-    cols = [k[0] for k in rows.description]
-    st.text("%s: %s" % (table, cols))
+        rows = cur.execute("select * from %s limit 1" % table)
+        cols = [k[0] for k in rows.description]
 
-con.close()
+        ts.append(table)
+        cs.append(cols)
 
-# st.dataframe(pd.DataFrame(metadata, columns=['table', 'columns']))
+    con.close()
+    return {'table': ts, 'columns': cs}
+
+st.title("English to SQL via LangChain")
+st.subheader("table metadata")
+st.dataframe(_get_metadata())
 
 llm = OpenAI(temperature=0.0, openai_api_key=API_KEY)
-db_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True, use_query_checker=True)
+db = SQLDatabase.from_uri("sqlite:///Chinook.db", include_tables=TABLES)
+db_chain = SQLDatabaseChain.from_llm(llm, db, use_query_checker=True, return_intermediate_steps=True)
 
-queries = ("How many albums are there?"
-    , "Which album has the most tracks?"
-    , "What artist has the album with the most tracks?"
-)
+# db_chain = SQLDatabaseChain.from_llm(llm, db, use_query_checker=True)
 
-for query in queries:
+for query in QUERIES:
 
     result = db_chain.run(query)
     # print(result)
     
     st.text("Q: %s" % query)
+    st.text(result["intermediate_steps"])
     st.text("A: %s" % result)