Enable alternate options in chat #4

Author: pramitchoudhary

sidekick/prompter.py

@@ -13,7 +13,8 @@
 from sidekick.db_config import DBConfig
 from sidekick.memory import EntityMemory
 from sidekick.query import SQLGenerator
-from sidekick.utils import execute_query_pd, extract_table_names, save_query, setup_dir
+from sidekick.utils import (execute_query_pd, extract_table_names, save_query,
+                            setup_dir)
 
 # Load the config file and initialize required paths
 base_path = (Path(__file__).parent / "../").resolve()
@@ -410,7 +411,7 @@ def query_api(
         if updated_tasks is not None:
             sql_g._tasks = updated_tasks
 
-    res = sql_g.generate_sql(
+    res, alt_res = sql_g.generate_sql(
         table_names, question, model_name=model_name, _dialect=db_dialect, is_regenerate=is_regenerate
     )
     logger.info(f"Input query: {question}")
@@ -430,13 +431,14 @@ def query_api(
                     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, model_name=model_name, _dialect=db_dialect, is_regenerate=True
+                    res, alt_res = sql_g.generate_sql(
+                        table_names, question, model_name=model_name, _dialect=db_dialect, is_regenerate=is_regenerate
                     )
                     logger.info(f"Input query: {question}")
                     logger.info(f"Generated response:\n\n{res}")
 
-        results.extend(["Generated Query:\n", res, "\n"])
+        results.extend(["**Generated Query:**\n", res, "\n"])
+        logger.info(f"Alternate responses:\n\n{alt_res}")
 
         exe_sql = click.prompt("Would you like to execute the generated SQL (y/n)?") if is_command else "y"
         if exe_sql.lower() == "y" or exe_sql.lower() == "yes":
@@ -484,14 +486,14 @@ def query_api(
                     click.echo("Error in executing the query. Validate generated SQL and try again.")
                     click.echo("No result to display.")
 
-            results.append("Query Results: \n")
+            results.append("**Query Results:** \n")
             if q_res:
                 click.echo(f"The query results are:\n {q_res}")
                 results.extend([str(q_res), "\n"])
             else:
                 click.echo(f"While executing query:\n {err}")
                 results.extend([str(err), "\n"])
-            # results.extend(["Query Results:", q_res])
+
         save_sql = click.prompt("Would you like to save the generated SQL (y/n)?") if is_command else "n"
         if save_sql.lower() == "y" or save_sql.lower() == "yes":
             # Persist for future use
@@ -500,7 +502,7 @@ def query_api(
         else:
             click.echo("Exiting...")
 
-    return results, err
+    return results, alt_res, err
 
 
 if __name__ == "__main__":

sidekick/query.py

@@ -421,6 +421,8 @@ def generate_sql(
             # Greedy search for quick response
             self.model.eval()
             device_type = "cuda" if torch.cuda.is_available() else "cpu"
+
+            alternate_queries = []
             if not is_regenerate:
                 output = self.model.generate(
                     **inputs.to(device_type),
@@ -476,7 +478,9 @@ def generate_sql(
                         res = "SELECT " + result.strip() + " LIMIT 100;"
                     else:
                         res = "SELECT " + result.strip() + ";"
-                    logger.info(f"Option: {idx+1}:\n{res}\nprobability: {probabilities_scores[idx]}")
+                    alt_res = f"Option {idx+1}: (_probability_: {probabilities_scores[idx]})\n{res}"
+                    alternate_queries.append(alt_res)
+                    logger.info(alt_res)
 
             _res = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
             # Below is a pre-caution in-case of an error in table name during generation
@@ -500,7 +504,7 @@ def generate_sql(
             except (sqlglot.errors.ParseError, ValueError, RuntimeError) as e:
                 logger.info("We did the best we could, there might be still be some error:\n")
                 logger.info(f"Realized query so far:\n {res}")
-        return result
+        return result, alternate_queries
 
     def task_formatter(self, input_task: str):
         # Generated format

ui/app.py

@@ -7,7 +7,7 @@
 import toml
 from h2o_wave import Q, app, data, handle_on, main, on, ui
 from sidekick.prompter import db_setup_api, query_api
-from sidekick.utils import setup_dir, update_tables, get_table_keys
+from sidekick.utils import get_table_keys, setup_dir, update_tables
 
 # Load the config file and initialize required paths
 base_path = (Path(__file__).parent / "../").resolve()
@@ -130,20 +130,27 @@ async def chatbot(q: Q):
         )
     elif q.args.chatbot.lower() == "regenerate":
         if q.client.query is not None and q.client.query.strip() != "":
-            llm_response, err = query_api(
+            llm_response, alt_response, err = query_api(
                 question=q.client.query,
                 sample_queries_path=q.user.sample_qna_path,
                 table_info_path=q.user.table_info_path,
                 table_name=q.user.table_name,
                 is_regenerate=True,
             )
-            llm_response = "\n".join(llm_response)
+            response = "\n".join(llm_response)
+            if alt_response:
+                llm_response = response + "\n\n" + "**Alternate options:**\n" + "\n".join(alt_response)
+                logging.info(f"Regenerate response: {llm_response}")
+            else:
+                llm_response = response
         else:
-            llm_response, err = ("Sure, I can generate a new response for you. However, in order to assist you "
-                                 "effectively could you please provide me with your question?"), None
+            llm_response, err = (
+                "Sure, I can generate a new response for you. However, in order to assist you "
+                "effectively could you please provide me with your question?"
+            ), None
     else:
         q.client.query = question
-        llm_response, err = query_api(
+        llm_response, alt_response, err = query_api(
             question=question,
             sample_queries_path=q.user.sample_qna_path,
             table_info_path=q.user.table_info_path,