From 445e4bbfeebd58d90d045d119ef1f87339277d6f Mon Sep 17 00:00:00 2001 From: Chih-Yu Yeh Date: Thu, 4 Jul 2024 10:53:48 +0800 Subject: [PATCH] add quotes for the fallback option when sql breakdown fails (#480) * add quotes * fix failed test --- wren-ai-service/src/web/v1/services/ask_details.py | 3 ++- wren-ai-service/tests/pytest/services/test_ask_details.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) 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 037fbeda8..8eba91b0c 100644 --- a/wren-ai-service/src/web/v1/services/ask_details.py +++ b/wren-ai-service/src/web/v1/services/ask_details.py @@ -4,6 +4,7 @@ from typing import List, Literal, Optional from haystack import Pipeline from pydantic import BaseModel +from src.core.engine import add_quotes from src.utils import async_timer logger = logging.getLogger("wren-ai-service") @@ -93,7 +94,7 @@ class AskDetailsService: if not ask_details_result["steps"]: ask_details_result["steps"] = [ { - "sql": ask_details_request.sql, + "sql": add_quotes(ask_details_request.sql), "summary": ask_details_request.summary, "cte_name": "", } diff --git a/wren-ai-service/tests/pytest/services/test_ask_details.py b/wren-ai-service/tests/pytest/services/test_ask_details.py index 3f1a04317..6d36b10a9 100644 --- a/wren-ai-service/tests/pytest/services/test_ask_details.py +++ b/wren-ai-service/tests/pytest/services/test_ask_details.py @@ -70,7 +70,7 @@ def test_ask_details_with_successful_sql(ask_details_service: AskDetailsService) def test_ask_details_with_failed_sql(ask_details_service: AskDetailsService): # asking details query_id = str(uuid.uuid4()) - sql = "SELECT * FROM xxx" + sql = 'SELECT * FROM "xxx"' summary = "This is a summary" ask_details_request = AskDetailsRequest( query="How many books are there?'",