|
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 |
|
|
|
|
|
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 answer_question(query): |
|
query_check = check_query(query) |
|
if isinstance(query_check, dict): |
|
return("BIG TABLE") |
|
if query_check == 'small': |
|
return('SMALL TABLE') |
|
if query_check == 'error': |
|
return('ERROR: Wrong format for getting the big db schema') |
|
|
|
|
|
|
|
return("DONE: " + query) |
|
|
|
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() |