Spaces:
Build error
Build error
| import os | |
| from langchain import OpenAI, SQLDatabase, SQLDatabaseChain | |
| import emoji | |
| import pandas as pd | |
| import sqlite3 | |
| import streamlit as st | |
| OPENAI_API_KEY = os.getenv('OPENAI_API_KEY') | |
| TABLES = ['Album', 'Artist', 'Track'] | |
| QUERIES = ("How many albums are there?" | |
| , "Which album has the most tracks?" | |
| , "What artist has the album with the most tracks?" | |
| ) | |
| st.title(emoji.emojize(":robot_face: Text to SQL via LangChain :robot_face:")) | |
| st.subheader("table metadata") | |
| def _get_metadata(): | |
| 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] | |
| ts.append(table) | |
| cs.append(cols) | |
| con.close() | |
| return pd.DataFrame({'table': ts, 'columns': cs}) | |
| st.write(_get_metadata()) | |
| llm = OpenAI(temperature=0.0, openai_api_key=OPENAI_API_KEY) | |
| 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) | |
| for i, query in enumerate(QUERIES): | |
| output = db_chain(query) | |
| sql, result = output["intermediate_steps"][1], output["result"] | |
| st.subheader("query %s" % str(1 + i)) | |
| st.write("Q: %s" % query) | |
| st.code(sql) | |
| st.write("A: %s" % result) |