Minor fixes n notebook example on using configured guardrails; set version to v0.2.0

Author: pramitchoudhary

.DS_Store

@@ -158,3 +158,11 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 .idea/
+
+# osx generated files
+.DS_Store
+.DS_Store?
+.Trashes
+ehthumbs.db
+Thumbs.db
+.idea

.gitignore

@@ -4,7 +4,7 @@ title = "SQL-Sidekick"
 description = "QnA with tabular data using NLQ"
 LongDescription = "about.md"
 Tags = ["DATA_SCIENCE", "MACHINE_LEARNING", "NLP"]
-Version = "0.1.9"
+Version = "0.2.0"
 
 [Runtime]
 MemoryLimit = "64Gi"

app.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sql-sidekick"
-version = "0.1.9  "
+version = "0.2.0"
 license = "Apache-2.0 license"
 description = "An AI assistant for SQL generation"
 authors = [

examples/notebooks/Guardrails_SQL_injection.ipynb

@@ -23,7 +23,7 @@
                             execute_query_pd, extract_table_names,
                             generate_suggestions, save_query, setup_dir)
 
-__version__ = "0.1.9"
+__version__ = "0.2.0"
 
 # Load the config file and initialize required paths
 app_base_path = (Path(__file__).parent / "../").resolve()
@@ -692,15 +692,19 @@ def ask(
                             json.dump(table_metadata, outfile, indent=4, sort_keys=False)
                     try:
                         q_res = execute_query_pd(query=_val, tables_path=tables_path, n_rows=100)
-                        click.echo(f"The query results are:\n {q_res}")
                     except sqldf.PandaSQLException as e:
                         logger.error(f"Error in executing the query: {e}")
                         click.echo("Error in executing the query. Validate generated SQL and try again.")
                         click.echo("No result to display.")
 
                 results.append("**Result:** \n")
                 if q_res:
-                    click.echo(f"The query results are:\n {q_res}")
+                    # Check shape of the final result to avoid blowing up memory
+                    # Logging a quick preview of the result
+                    if isinstance(q_res, list) and len(q_res) > 10:
+                        click.echo(f"Preview of the result:\n {pd.DataFrame(q_res).head(2)}")
+                    else:
+                        click.echo(f"The results are:\n {q_res}")
                     results.extend([str(q_res), "\n"])
                 else:
                     click.echo(f"While executing query:\n {err}")

pyproject.toml

@@ -65,12 +65,13 @@ def test_no_error():
 
 def test_blind_select_injection():
     input_q = """
-    SELECT * FROM sleep_health_and_lifestyle WHERE person_id = 105 OR 1=1;"
+    SELECT * FROM sleep_health_and_lifestyle WHERE person_id = 105 OR 1=1;
     """
 
+    # 1. When guardrails are disabled
     result = None
     question = f"Execute SQL:\n{input_q}"
-    #1. Self correction is disabled
+    # Self correction is disabled
     result, _, error = ask(
         question=question,
         table_info_path=table_info_path,
@@ -83,6 +84,31 @@ def test_blind_select_injection():
         execute_query=True,
         local_base_path=base_path,
         debug_mode=True,
+        guardrails=False,
+        self_correction=False
+    )
+
+    assert 'malicious patterns' not in str(result)
+
+
+    # 2. When guardrails are enabled
+    result = None
+    question = f"Execute SQL:\n{input_q}"
+
+    # Self correction is disabled
+    result, _, error = ask(
+        question=question,
+        table_info_path=table_info_path,
+        sample_queries_path=None,
+        table_name=table_name,
+        is_command=False,
+        model_name="h2ogpt-sql-nsql-llama-2-7B",
+        is_regenerate=False,
+        is_regen_with_options=False,
+        execute_query=True,
+        local_base_path=base_path,
+        debug_mode=True,
+        guardrails=True,
         self_correction=False
     )
 
@@ -128,7 +154,7 @@ def test_stacked_queries():
     result = None
     question = f"Execute SQL:\n{input_q}"
 
-    result, _ar, error = ask(
+    result, _, _ = ask(
         question=question,
         table_info_path=table_info_path,
         sample_queries_path=None,