chore(wren-ai-service): handle add_quotes error (#595)

wren-ai-service/demo/utils.py

@@ -5,7 +5,7 @@ import os
 import re
 import time
 from pathlib import Path
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple
 
 import orjson
 import pandas as pd
@@ -25,8 +25,12 @@ LLM_MODELS = ["gpt-4o-mini", "gpt-4o", "gpt-4-turbo", "gpt-3.5-turbo"]
 load_dotenv()
 
 
-def add_quotes(sql: str) -> str:
-    return sqlglot.transpile(sql, read="trino", identify=True)[0]
+def add_quotes(sql: str) -> Tuple[str, bool]:
+    try:
+        quoted_sql = sqlglot.transpile(sql, read="trino", identify=True)[0]
+        return quoted_sql, True
+    except Exception:
+        return sql, False
 
 
 def get_connection_info(data_source: str):
@@ -135,10 +139,13 @@ def get_data_from_wren_engine(
     manifest: dict,
 ):
     if dataset_type == "duckdb":
+        quoted_sql, no_error = add_quotes(sql)
+        assert no_error, f"Error in adding quotes to SQL: {sql}"
+
         response = requests.get(
             f"{WREN_ENGINE_API_URL}/v1/mdl/preview",
             json={
-                "sql": add_quotes(sql),
+                "sql": quoted_sql,
                 "manifest": manifest,
                 "limit": 100,
             },
@@ -151,10 +158,12 @@ def get_data_from_wren_engine(
 
         return pd.DataFrame(data["data"], columns=column_names)
     else:
+        quoted_sql, no_error = add_quotes(sql)
+        assert no_error, f"Error in adding quotes to SQL: {sql}"
         response = requests.post(
             f"{WREN_IBIS_API_URL}/v2/connector/{dataset_type}/query?limit=100",
             json={
-                "sql": add_quotes(sql),
+                "sql": quoted_sql,
                 "manifestStr": base64.b64encode(orjson.dumps(manifest)).decode(),
                 "connectionInfo": get_connection_info(dataset_type),
                 "limit": 100,

wren-ai-service/eval/data_curation/utils.py

@@ -47,11 +47,14 @@ async def is_sql_valid(
     timeout: float = TIMEOUT_SECONDS,
 ) -> Tuple[bool, str]:
     sql = sql.rstrip(";") if sql.endswith(";") else sql
+    quoted_sql, no_error = add_quotes(sql)
+    assert no_error, f"Error in quoting SQL: {sql}"
+
     async with aiohttp.request(
         "POST",
         f"{api_endpoint}/v2/connector/{data_source}/query?dryRun=true",
         json={
-            "sql": add_quotes(sql),
+            "sql": quoted_sql,
             "manifestStr": base64.b64encode(orjson.dumps(mdl_json)).decode(),
             "connectionInfo": connection_info,
         },

wren-ai-service/eval/utils.py

@@ -1,6 +1,6 @@
 import base64
 import os
-from typing import Any, Dict, List, Literal, Optional, get_args
+from typing import Any, Dict, List, Literal, Optional, Tuple, get_args
 
 import aiohttp
 import orjson
@@ -8,8 +8,12 @@ import sqlglot
 from tomlkit import parse
 
 
-def add_quotes(sql: str) -> str:
-    return sqlglot.transpile(sql, read="trino", identify=True)[0]
+def add_quotes(sql: str) -> Tuple[str, bool]:
+    try:
+        quoted_sql = sqlglot.transpile(sql, read="trino", identify=True)[0]
+        return quoted_sql, True
+    except Exception:
+        return sql, False
 
 
 async def get_data_from_wren_engine(
@@ -25,11 +29,14 @@ async def get_data_from_wren_engine(
     if limit is not None:
         url += f"?limit={limit}"
 
+    quoted_sql, no_error = add_quotes(sql)
+    assert no_error, f"Error in quoting SQL: {sql}"
+
     async with aiohttp.request(
         "POST",
         url,
         json={
-            "sql": add_quotes(sql),
+            "sql": quoted_sql,
             "manifestStr": base64.b64encode(orjson.dumps(mdl_json)).decode(),
             "connectionInfo": connection_info,
         },
@@ -123,11 +130,14 @@ async def get_contexts_from_sql(
         timeout: float,
     ) -> List[dict]:
         sql = sql.rstrip(";") if sql.endswith(";") else sql
+        quoted_sql, no_error = add_quotes(sql)
+        assert no_error, f"Error in quoting SQL: {sql}"
+
         async with aiohttp.request(
             "GET",
             f"{api_endpoint}/v1/analysis/sql",
             json={
-                "sql": add_quotes(sql),
+                "sql": quoted_sql,
                 "manifest": mdl_json,
             },
             timeout=aiohttp.ClientTimeout(total=timeout),

wren-ai-service/src/core/engine.py

@@ -48,11 +48,16 @@ def remove_limit_statement(sql: str) -> str:
     return modified_sql
 
 
-def add_quotes(sql: str) -> str:
-    logger.debug(f"Original SQL: {sql}")
+def add_quotes(sql: str) -> Tuple[str, bool]:
+    try:
+        logger.debug(f"Original SQL: {sql}")
 
-    quoted_sql = sqlglot.transpile(sql, read="trino", identify=True)[0]
+        quoted_sql = sqlglot.transpile(sql, read="trino", identify=True)[0]
 
-    logger.debug(f"Quoted SQL: {quoted_sql}")
+        logger.debug(f"Quoted SQL: {quoted_sql}")
+    except Exception as e:
+        logger.exception(f"Error in sqlglot.transpile to {sql}: {e}")
 
-    return quoted_sql
+        return "", False
+
+    return quoted_sql, True

wren-ai-service/src/pipelines/ask/components/post_processors.py

@@ -63,25 +63,36 @@ class GenerationPostProcessor:
         invalid_generation_results = []
 
         async def _task(result: Dict[str, str]):
-            quoted_sql = add_quotes(result["sql"])
+            quoted_sql, no_error = add_quotes(result["sql"])
 
-            status, error = await self._engine.dry_run_sql(
-                quoted_sql, session, project_id=project_id
-            )
-
-            if status:
-                valid_generation_results.append(
-                    {
-                        "sql": quoted_sql,
-                        "summary": result["summary"],
-                    }
+            if no_error:
+                status, error = await self._engine.dry_run_sql(
+                    quoted_sql, session, project_id=project_id
                 )
+
+                if status:
+                    valid_generation_results.append(
+                        {
+                            "sql": quoted_sql,
+                            "summary": result["summary"],
+                        }
+                    )
+                else:
+                    invalid_generation_results.append(
+                        {
+                            "sql": quoted_sql,
+                            "summary": result["summary"],
+                            "type": "DRY_RUN",
+                            "error": error,
+                        }
+                    )
             else:
                 invalid_generation_results.append(
                     {
-                        "sql": quoted_sql,
+                        "sql": result["sql"],
                         "summary": result["summary"],
-                        "error": error,
+                        "type": "ADD_QUOTES",
+                        "error": "add_quotes failed",
                     }
                 )
 

wren-ai-service/src/pipelines/common.py

@@ -43,7 +43,14 @@ class GenerationPostProcessor:
         steps[-1]["cte_name"] = ""
 
         for step in steps:
-            step["sql"] = add_quotes(step["sql"])
+            step["sql"], no_error = add_quotes(step["sql"])
+            if not no_error:
+                return {
+                    "results": {
+                        "description": cleaned_generation_result["description"],
+                        "steps": [],
+                    },
+                }
 
         sql = self._build_cte_query(steps)
         logger.debug(f"GenerationPostProcessor: steps: {pformat(steps)}")

wren-ai-service/src/providers/engine/wren.py

@@ -6,7 +6,7 @@ from typing import Any, Dict, Optional, Tuple
 import aiohttp
 import orjson
 
-from src.core.engine import Engine, add_quotes, remove_limit_statement
+from src.core.engine import Engine, remove_limit_statement
 from src.providers.loader import provider
 
 logger = logging.getLogger("wren-ai-service")
@@ -33,7 +33,7 @@ class WrenUI(Engine):
                     "data": {
                         "dryRun": True,
                         "limit": 1,
-                        "sql": remove_limit_statement(add_quotes(sql)),
+                        "sql": remove_limit_statement(sql),
                         "projectId": project_id,
                     }
                 },
@@ -73,7 +73,7 @@ class WrenIbis(Engine):
         async with session.post(
             f"{self._endpoint}/v2/connector/{self._source}/query?dryRun=true&limit=1",
             json={
-                "sql": remove_limit_statement(add_quotes(sql)),
+                "sql": remove_limit_statement(sql),
                 "manifestStr": self._manifest,
                 "connectionInfo": self._connection_info,
             },
@@ -106,7 +106,7 @@ class WrenEngine(Engine):
                 "manifest": orjson.loads(base64.b64decode(properties.get("manifest")))
                 if properties.get("manifest")
                 else {},
-                "sql": remove_limit_statement(add_quotes(sql)),
+                "sql": remove_limit_statement(sql),
                 "limit": 1,
             },
         ) as response:

wren-ai-service/src/web/v1/services/ask.py

@@ -106,6 +106,11 @@ class AskService:
 
         return False
 
+    def _get_failed_dry_run_results(self, invalid_generation_results: list[dict]):
+        return list(
+            filter(lambda x: x["type"] == "DRY_RUN", invalid_generation_results)
+        )
+
     @async_timer
     @observe(name="Ask Question")
     @trace_metadata
@@ -208,16 +213,16 @@ class AskService:
                 logger.debug("Before sql correction:")
                 logger.debug(f"valid_generation_results: {valid_generation_results}")
 
-                if text_to_sql_generation_results["post_process"][
-                    "invalid_generation_results"
-                ]:
+                if failed_dry_run_results := self._get_failed_dry_run_results(
+                    text_to_sql_generation_results["post_process"][
+                        "invalid_generation_results"
+                    ]
+                ):
                     sql_correction_results = await self._pipelines[
                         "sql_correction"
                     ].run(
                         contexts=documents,
-                        invalid_generation_results=text_to_sql_generation_results[
-                            "post_process"
-                        ]["invalid_generation_results"],
+                        invalid_generation_results=failed_dry_run_results,
                         project_id=ask_request.project_id,
                     )
                     valid_generation_results += sql_correction_results["post_process"][
@@ -228,17 +233,15 @@ class AskService:
                         f'sql_correction_results: {sql_correction_results["post_process"]}'
                     )
 
-                    for invalid_generation_result in sql_correction_results[
-                        "post_process"
-                    ]["invalid_generation_results"]:
+                    for failed_dry_run_result in failed_dry_run_results:
                         logger.debug(
                             f"{sqlparse.format(
-                                invalid_generation_result['sql'],
+                                failed_dry_run_result['sql'],
                                 reindent=True,
                                 keyword_case='upper')
                             }"
                         )
-                        logger.debug(invalid_generation_result["error"])
+                        logger.debug(failed_dry_run_result["error"])
                         logger.debug("\n\n")
 
                 # remove duplicates of valid_generation_results, which consists of a sql and a summary

wren-ai-service/src/web/v1/services/ask_details.py

@@ -106,9 +106,10 @@ class AskDetailsService:
             ask_details_result = generation_result["post_process"]["results"]
 
             if not ask_details_result["steps"]:
+                quoted_sql, no_error = add_quotes(ask_details_request.sql)
                 ask_details_result["steps"] = [
                     {
-                        "sql": add_quotes(ask_details_request.sql),
+                        "sql": quoted_sql if no_error else ask_details_request.sql,
                         "summary": ask_details_request.summary,
                         "cte_name": "",
                     }