udpated error message and validation

app.py

@@ -1,13 +1,12 @@
 import streamlit as st
 import requests
-import sqlglot
-import tempfile
+import subprocess
 import re
-from validate_sql import validate_query
+import sys
 
 PROMPT_TEMPLATE = """### Instruction:\n{instruction}\n\n### Input:\n{input}\n### Question:\n{question}\n\n### Response (use duckdb shorthand if possible):\n"""
 INSTRUCTION_TEMPLATE = """Your task is to generate valid duckdb SQL to answer the following question{has_schema}"""  # noqa: E501
-ERROR_MESSAGE = ":red[ Quack! Much to our regret, SQL generation has gone a tad duck-side-down.\nThe model is currently not capable of crafting the desired SQL. \nSorry my duck friend. ]\n\n:red[ If the question is about your own database, make sure to set the correct schema. ]\n\n```sql\n{sql_query}\n```\n\n```sql\n{error_msg}\n```"
+ERROR_MESSAGE = ":red[ Quack! Much to our regret, SQL generation has gone a tad duck-side-down.\nThe model is currently not able to craft a correct SQL query for this request. \nSorry my duck friend. ]\n\n:red[ Try rephrasing the question/instruction. And if the question is about your own database, make sure to set the correct schema. ]\n\n```sql\n{sql_query}\n```\n\n```sql\n{error_msg}\n```"
 STOP_TOKENS = ["###", ";", "--", "```"]
 
 
@@ -54,17 +53,26 @@ def generate_sql(question, schema):
     with s.post(url, json=body, headers=headers) as resp:
         sql_query = resp.json()["choices"][0]["text"]
 
-    #for token in STOP_TOKENS:
-    #    sql_query = sql_query.split(token)[0]
-    #sql_query = sqlglot.parse_one(sql_query, read="duckdb").sql(
-    #    dialect="duckdb", pretty=True
-    #)
     return sql_query
 
 
 def validate_sql(query, schema):
-    valid, msg = validate_query(query, schema)
-    return valid, msg
+    try:
+        # Define subprocess
+        process = subprocess.Popen(
+            [sys.executable, './validate_sql.py', query, schema],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE
+        )
+        # Get output and potential parser, and binder error message
+        stdout, stderr = process.communicate(timeout=0.5)
+        if stderr:
+            return False, stderr.decode('utf8')
+        return True, ""
+    except subprocess.TimeoutExpired:
+        process.kill()
+        # timeout reached, so parsing and binding was very likely successful
+        return True, ""
 
 
 st.title("DuckDB-NSQL-7B Demo")

validate_sql.py

@@ -4,8 +4,6 @@ from duckdb import ParserException, SyntaxException, BinderException, CatalogExc
 
 
 def validate_query(query, schemas):
-    valid = True
-    msg = ""
     try:
         print("Running query: ", query)
         with duckdb.connect(
@@ -17,25 +15,17 @@ def validate_query(query, schemas):
             cursor = duckdb_conn.cursor()
             cursor.execute(query)
     except ParserException as e:
-        msg = str(e)
-        valid = False
+        raise e
     except SyntaxException as e:
-        msg = str(e)
-        valid = False
+        raise e
     except BinderException as e:
-        msg = str(e)
-        valid = False
+        raise e
     except CatalogException as e:
-        msg = str(e)
-        if "but it exists" in msg and "extension" in msg:
-            valid = True
-            msg = ""
-        else:
-            valid = False
+        if not ("but it exists" in str(e) and "extension" in str(e)):
+            raise e
     except Exception as e:
-        msg = str(e)
-        valid = True
-    return valid, msg
+        return True
+    return True
 
 
 if __name__ == "__main__":