Spaces:
Build error
Build error
File size: 1,057 Bytes
9225820 d472eeb 2c74cef 9225820 67928b0 2c74cef 9225820 2c74cef 9225820 2c74cef 03abb31 e3a559e 67928b0 2c74cef e3a559e 67928b0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
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'])
|