|
from langchain import SQLDatabaseChain |
|
from langchain.sql_database import SQLDatabase |
|
from langchain.llms.openai import OpenAI |
|
from langchain.chat_models import ChatOpenAI |
|
from langchain.prompts.prompt import PromptTemplate |
|
|
|
llm = ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo", verbose=True) |
|
|
|
DEFAULT_TABLES = [ |
|
'Active Players', |
|
'Team_Per_Game_Statistics_2022_23', |
|
"Team_Totals_Statistics_2022_23", |
|
"Player_Total_Statistics_2022_23", |
|
"Player_Per_Game_Statistics_2022_23" |
|
] |
|
|
|
def get_prompt(): |
|
_DEFAULT_TEMPLATE = """Given an input question, first create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer. |
|
Use the following format: |
|
|
|
Question: "Question here" |
|
SQLQuery: "SQL Query to run" |
|
SQLResult: "Result of the SQLQuery" |
|
|
|
Answer: "Final answer here" |
|
|
|
Only use the following tables: |
|
|
|
{table_info} |
|
|
|
Question: {input}""" |
|
|
|
PROMPT = PromptTemplate( |
|
input_variables=["input", "table_info", "dialect"], template=_DEFAULT_TEMPLATE |
|
) |
|
return PROMPT |
|
|
|
def check_query(query): |
|
if query.startswith("### Query"): |
|
split = query.split('\n\n') |
|
q_text = split[0] |
|
t_text = split[1] |
|
|
|
if t_text.startswith("### Tables"): |
|
query_params = dict() |
|
tables = t_text.split('\n') |
|
query_params['tables'] = tables[1:] |
|
query_params['q'] = q_text.split('\n')[1] |
|
print(query_params) |
|
return query_params |
|
else: |
|
return 'error' |
|
return 'small' |
|
|
|
def get_db(q, tables): |
|
if len(tables) == 0: |
|
db = SQLDatabase.from_uri("sqlite:///nba_small.db", |
|
sample_rows_in_table_info=2) |
|
else: |
|
tables.extend(DEFAULT_TABLES) |
|
db = SQLDatabase.from_uri("sqlite:///nba.db", |
|
include_tables = tables, |
|
sample_rows_in_table_info=2) |
|
return db |
|
def answer_question(query): |
|
PROMPT = get_prompt() |
|
query_check = check_query(query) |
|
if query_check == 'error': |
|
return('ERROR: Wrong format for getting the big db schema') |
|
if isinstance(query_check, dict): |
|
q = query_check['q'] |
|
tables = query_check['tables'] |
|
if query_check == 'small': |
|
q = query |
|
tables = [] |
|
db = get_db(q, tables) |
|
|
|
db_chain = SQLDatabaseChain.from_llm(llm, db, |
|
prompt=PROMPT, |
|
verbose=True, |
|
return_intermediate_steps=True, |
|
|
|
) |
|
result = db_chain(q) |
|
return result['result'] |
|
|
|
if __name__ == "__main__": |
|
import gradio as gr |
|
|
|
|
|
gr.Interface( |
|
answer_question, |
|
[ |
|
gr.inputs.Textbox(lines=10, label="Query"), |
|
], |
|
gr.outputs.Textbox(label="Response"), |
|
title="Ask NBA Stats", |
|
description=""" Ask NBA Stats is a tool that let's you ask a question with |
|
the NBA SQL tables as a reference |
|
|
|
Ask a simple question to use the small database |
|
|
|
If you would like to access the large DB use format |
|
|
|
### Query |
|
single line query |
|
|
|
### Tables |
|
tables to access line by line |
|
table1 |
|
table2""" |
|
).launch() |
