First Chain Check

app.py

@@ -4,6 +4,36 @@ 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"):
@@ -22,18 +52,36 @@ def check_query(query):
             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=3)
+    else:
+        tables.extend(DEFAULT_TABLES)
+        db = SQLDatabase.from_uri("sqlite:///nba_small.db",
+                                  include_tables = tables
+                                  sample_rows_in_table_info=3)
+    return db
 def answer_question(query):
+    PROMPT = get_prompt()
     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')
-  
-
+    if isinstance(query_check, dict):
+        q = query_check['q']
+        tables = query_check
+    if query_check == 'small':
+        q = query
+        tables = []
+    db = get_db(q, tables)
 
-    return("DONE: " + query)
+    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