Parameterize model name n minor prompt adjustments

Author: pramitchoudhary

sidekick/configs/.env.toml

@@ -1,5 +1,6 @@
 [OPENAI]
 OPENAI_API_KEY = ""
+MODEL_NAME = "gpt-3.5-turbo-0301" # Others: e.g. text-davinci-003
 
 [LOCAL_DB_CONFIG]
 HOST_NAME = "localhost"

sidekick/configs/prompt_template.py

@@ -52,7 +52,7 @@
 DEBUGGING_PROMPT = {
     "system_prompt": "Act as a SQL expert for PostgreSQL code",
     "user_prompt": """
-                ### Fix syntax errors for provided SQL Query.
+                ### Fix syntax errors for provided incorrect SQL Query.
                 # Add ``` as prefix and ``` as suffix to generated SQL
                 # Error: {ex_traceback}
                 # Add explanation and reasoning for each SQL query

sidekick/prompter.py

@@ -26,7 +26,7 @@ def color(fore="", back="", text=None):
     return f"{fore}{back}{text}{Style.RESET_ALL}"
 
 
-msg = """Welcome to the SQL Sidekick!\nI am AI assistant that helps you with SQL queries.
+msg = """Welcome to the SQL Sidekick!\nI am an AI assistant that helps you with SQL queries.
 I can help you with the following:\n
 1. Configure a local database(for schema validation and syntax checking): `sql-sidekick configure db-setup`.\n
 2. Learn contextual query/answer pairs: `sql-sidekick learn add-samples`.\n
@@ -307,7 +307,9 @@ def query(question: str, table_info_path: str, sample_queries: str):
             click.echo("Skipping edit...")
     if updated_tasks is not None:
         sql_g._tasks = updated_tasks
-    res = sql_g.generate_sql(table_names, question)
+
+    model_name = env_settings["OPENAI"]["MODEL_NAME"]
+    res = sql_g.generate_sql(table_names, question, model_name=model_name)
     logger.info(f"Input query: {question}")
     logger.info(f"Generated response:\n\n{res}")
 
@@ -324,7 +326,7 @@ def query(question: str, table_info_path: str, sample_queries: str):
                 click.echo(f"Updated SQL:\n {updated_sql}")
             elif res_val.lower() == "r" or res_val.lower() == "regenerate":
                 click.echo("Attempting to regenerate...")
-                res = sql_g.generate_sql(table_names, question)
+                res = sql_g.generate_sql(table_names, question, model_name=model_name)
                 logger.info(f"Input query: {question}")
                 logger.info(f"Generated response:\n\n{res}")
 

sidekick/query.py

@@ -138,7 +138,7 @@ def generate_response(self, context_container, sql_index, input_prompt, attempt_
                 try:
                     # Attempt to heal with simple feedback
                     # Reference: Teaching Large Language Models to Self-Debug, https://arxiv.org/abs/2304.05128
-                    logger.info(f"Attempting to heal ...,\n {se}")
+                    logger.info(f"Attempting to fix syntax error ...,\n {se}")
                     system_prompt = DEBUGGING_PROMPT["system_prompt"]
                     user_prompt = DEBUGGING_PROMPT["user_prompt"].format(ex_traceback=ex_traceback, qry_txt=qry_txt)
                     # Role and content
@@ -193,7 +193,7 @@ def generate_tasks(self, table_names: list, input_question: str):
         except Exception as se:
             raise se
 
-    def generate_sql(self, table_name: list, input_question: str, _dialect: str = "postgres"):
+    def generate_sql(self, table_name: list, input_question: str, _dialect: str = "postgres", model_name: str = 'gpt-3.5-turbo-0301'):
         _tasks = self.task_formatter(self._tasks)
         context_file = f"{self.path}/var/lib/tmp/data/context.json"
         additional_context = json.load(open(context_file, "r")) if Path(context_file).exists() else {}
@@ -217,7 +217,7 @@ def generate_sql(self, table_name: list, input_question: str, _dialect: str = "p
         context_container = self.context_builder.build_context_container()
 
         # Reference: https://github.com/jerryjliu/llama_index/issues/987
-        llm_predictor_gpt3 = LLMPredictor(llm=OpenAI(temperature=0.7, model_name="text-davinci-003"))
+        llm_predictor_gpt3 = LLMPredictor(llm=OpenAI(temperature=0.5, model_name=model_name))
         service_context_gpt3 = ServiceContext.from_defaults(llm_predictor=llm_predictor_gpt3, chunk_size_limit=512)
 
         index = GPTSQLStructStoreIndex(