From 49333c1b336ce757718b71adfd7aa5c3e6c558a5 Mon Sep 17 00:00:00 2001 From: Chih-Yu Yeh Date: Wed, 14 Aug 2024 13:09:07 +0800 Subject: [PATCH] chore(wren-ai-service): handle add_quotes error (#595) --- wren-ai-service/demo/utils.py | 19 +++++++--- wren-ai-service/eval/data_curation/utils.py | 5 ++- wren-ai-service/eval/utils.py | 20 +++++++--- wren-ai-service/src/core/engine.py | 15 +++++--- .../ask/components/post_processors.py | 37 ++++++++++++------- wren-ai-service/src/pipelines/common.py | 9 ++++- wren-ai-service/src/providers/engine/wren.py | 8 ++-- wren-ai-service/src/web/v1/services/ask.py | 25 +++++++------ .../src/web/v1/services/ask_details.py | 3 +- 9 files changed, 95 insertions(+), 46 deletions(-) diff --git a/wren-ai-service/demo/utils.py b/wren-ai-service/demo/utils.py index 0f6f8d91b..55864c12d 100644 --- a/wren-ai-service/demo/utils.py +++ b/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, diff --git a/wren-ai-service/eval/data_curation/utils.py b/wren-ai-service/eval/data_curation/utils.py index 792b7ebf4..a78f25c65 100644 --- a/wren-ai-service/eval/data_curation/utils.py +++ b/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, }, diff --git a/wren-ai-service/eval/utils.py b/wren-ai-service/eval/utils.py index 3a6931215..ef273c072 100644 --- a/wren-ai-service/eval/utils.py +++ b/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), diff --git a/wren-ai-service/src/core/engine.py b/wren-ai-service/src/core/engine.py index 7e8acc3c7..8b7db7b9f 100644 --- a/wren-ai-service/src/core/engine.py +++ b/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 diff --git a/wren-ai-service/src/pipelines/ask/components/post_processors.py b/wren-ai-service/src/pipelines/ask/components/post_processors.py index 2b4b75370..edfe95d88 100644 --- a/wren-ai-service/src/pipelines/ask/components/post_processors.py +++ b/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", } ) diff --git a/wren-ai-service/src/pipelines/common.py b/wren-ai-service/src/pipelines/common.py index c8ae9af1a..c6f2f7b8c 100644 --- a/wren-ai-service/src/pipelines/common.py +++ b/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)}") diff --git a/wren-ai-service/src/providers/engine/wren.py b/wren-ai-service/src/providers/engine/wren.py index 5122feb05..09204ab56 100644 --- a/wren-ai-service/src/providers/engine/wren.py +++ b/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: diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index 5244fea04..ab7973428 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/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 diff --git a/wren-ai-service/src/web/v1/services/ask_details.py b/wren-ai-service/src/web/v1/services/ask_details.py index 2a1cd3632..bdf262d74 100644 --- a/wren-ai-service/src/web/v1/services/ask_details.py +++ b/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": "", }