chore(wren-ai-service): minor-updates

deployment/kustomizations/base/cm.yaml

@@ -131,7 +131,6 @@ data:
         llm: litellm_llm.gpt-4o-mini-2024-07-18
       - name: sql_answer
         llm: litellm_llm.gpt-4o-mini-2024-07-18
-        engine: wren_ui
       - name: sql_breakdown
         llm: litellm_llm.gpt-4o-mini-2024-07-18
         engine: wren_ui
@@ -188,6 +187,7 @@ data:
 
     ---
     settings:
+      engine_timeout: 30
       column_indexing_batch_size: 50
       table_retrieval_size: 10
       table_column_retrieval_size: 100

docker/config.example.yaml

@@ -83,7 +83,6 @@ pipes:
     llm: litellm_llm.gpt-4o-mini-2024-07-18
   - name: sql_answer
     llm: litellm_llm.gpt-4o-mini-2024-07-18
-    engine: wren_ui
   - name: sql_breakdown
     llm: litellm_llm.gpt-4o-mini-2024-07-18
     engine: wren_ui
@@ -140,6 +139,7 @@ pipes:
 
 ---
 settings:
+  engine_timeout: 30
   column_indexing_batch_size: 50
   table_retrieval_size: 10
   table_column_retrieval_size: 100

wren-ai-service/src/config.py

@@ -34,6 +34,9 @@ class Settings(BaseSettings):
     # generation config
     allow_sql_generation_reasoning: bool = Field(default=True)
 
+    # engine config
+    engine_timeout: float = Field(default=30.0)
+
     # service config
     query_cache_ttl: int = Field(default=3600)  # unit: seconds
     query_cache_maxsize: int = Field(

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

@@ -50,12 +50,14 @@ def remove_limit_statement(sql: str) -> str:
     return modified_sql
 
 
-def add_quotes(sql: str) -> Tuple[str, bool]:
+def add_quotes(sql: str) -> Tuple[str, str]:
     try:
-        quoted_sql = sqlglot.transpile(sql, read="trino", identify=True)[0]
+        quoted_sql = sqlglot.transpile(
+            sql, read="trino", identify=True, error_level=sqlglot.ErrorLevel.RAISE
+        )[0]
     except Exception as e:
         logger.exception(f"Error in sqlglot.transpile to {sql}: {e}")
 
-        return "", False
+        return "", str(e)
 
-    return quoted_sql, True
+    return quoted_sql, ""

wren-ai-service/src/globals.py

@@ -107,15 +107,18 @@ def create_service_container(
                 ),
                 "sql_generation": generation.SQLGeneration(
                     **pipe_components["sql_generation"],
+                    engine_timeout=settings.engine_timeout,
                 ),
                 "sql_generation_reasoning": generation.SQLGenerationReasoning(
                     **pipe_components["sql_generation_reasoning"],
                 ),
                 "sql_correction": generation.SQLCorrection(
                     **pipe_components["sql_correction"],
+                    engine_timeout=settings.engine_timeout,
                 ),
                 "followup_sql_generation": generation.FollowUpSQLGeneration(
                     **pipe_components["followup_sql_generation"],
+                    engine_timeout=settings.engine_timeout,
                 ),
                 "sql_summary": generation.SQLSummary(
                     **pipe_components["sql_summary"],
@@ -128,6 +131,7 @@ def create_service_container(
             pipelines={
                 "sql_executor": retrieval.SQLExecutor(
                     **pipe_components["sql_executor"],
+                    engine_timeout=settings.engine_timeout,
                 ),
                 "chart_generation": generation.ChartGeneration(
                     **pipe_components["chart_generation"],
@@ -139,6 +143,7 @@ def create_service_container(
             pipelines={
                 "sql_executor": retrieval.SQLExecutor(
                     **pipe_components["sql_executor"],
+                    engine_timeout=settings.engine_timeout,
                 ),
                 "chart_adjustment": generation.ChartAdjustment(
                     **pipe_components["chart_adjustment"],
@@ -153,6 +158,7 @@ def create_service_container(
                 ),
                 "sql_answer": generation.SQLAnswer(
                     **pipe_components["sql_answer"],
+                    engine_timeout=settings.engine_timeout,
                 ),
             },
             **query_cache,
@@ -161,6 +167,7 @@ def create_service_container(
             pipelines={
                 "sql_breakdown": generation.SQLBreakdown(
                     **pipe_components["sql_breakdown"],
+                    engine_timeout=settings.engine_timeout,
                 ),
                 "sql_summary": generation.SQLSummary(
                     **pipe_components["sql_summary"],
@@ -177,9 +184,11 @@ def create_service_container(
                 ),
                 "sql_expansion": generation.SQLExpansion(
                     **pipe_components["sql_expansion"],
+                    engine_timeout=settings.engine_timeout,
                 ),
                 "sql_correction": generation.SQLCorrection(
                     **pipe_components["sql_correction"],
+                    engine_timeout=settings.engine_timeout,
                 ),
                 "sql_summary": generation.SQLSummary(
                     **pipe_components["sql_summary"],
@@ -207,6 +216,7 @@ def create_service_container(
             pipelines={
                 "relationship_recommendation": generation.RelationshipRecommendation(
                     **pipe_components["relationship_recommendation"],
+                    engine_timeout=settings.engine_timeout,
                 )
             },
             **query_cache,
@@ -224,6 +234,7 @@ def create_service_container(
                 ),
                 "sql_generation": generation.SQLGeneration(
                     **pipe_components["question_recommendation_sql_generation"],
+                    engine_timeout=settings.engine_timeout,
                 ),
                 "sql_generation_reasoning": generation.SQLGenerationReasoning(
                     **pipe_components["sql_generation_reasoning"],

wren-ai-service/src/pipelines/generation/followup_sql_generation.py

@@ -1,6 +1,6 @@
 import logging
 import sys
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from hamilton import base
 from hamilton.async_driver import AsyncDriver
@@ -106,10 +106,13 @@ async def generate_sql_in_followup(prompt: dict, generator: Any) -> dict:
 async def post_process(
     generate_sql_in_followup: dict,
     post_processor: SQLGenPostProcessor,
+    engine_timeout: float,
     project_id: str | None = None,
 ) -> dict:
     return await post_processor.run(
-        generate_sql_in_followup.get("replies"), project_id=project_id
+        generate_sql_in_followup.get("replies"),
+        timeout=engine_timeout,
+        project_id=project_id,
     )
 
 
@@ -132,6 +135,7 @@ def __init__(
         self,
         llm_provider: LLMProvider,
         engine: Engine,
+        engine_timeout: Optional[float] = 30.0,
         **kwargs,
     ):
         self._components = {
@@ -145,6 +149,10 @@ def __init__(
             "post_processor": SQLGenPostProcessor(engine=engine),
         }
 
+        self._configs = {
+            "engine_timeout": engine_timeout,
+        }
+
         super().__init__(
             AsyncDriver({}, sys.modules[__name__], result_builder=base.DictResult())
         )
@@ -176,6 +184,7 @@ async def run(
                 "has_calculated_field": has_calculated_field,
                 "has_metric": has_metric,
                 **self._components,
+                **self._configs,
             },
         )
 

wren-ai-service/src/pipelines/generation/relationship_recommendation.py

@@ -1,7 +1,7 @@
 import logging
 import sys
 from enum import Enum
-from typing import Any
+from typing import Any, Optional
 
 import orjson
 from hamilton import base
@@ -170,6 +170,7 @@ def __init__(
         self,
         llm_provider: LLMProvider,
         engine: Engine,
+        engine_timeout: Optional[float] = 30.0,
         **_,
     ):
         self._components = {
@@ -181,6 +182,10 @@ def __init__(
             "engine": engine,
         }
 
+        self._configs = {
+            "engine_timeout": engine_timeout,
+        }
+
         self._final = "validated"
 
         super().__init__(
@@ -200,6 +205,7 @@ async def run(
                 "mdl": mdl,
                 "language": language,
                 **self._components,
+                **self._configs,
             },
         )
 

wren-ai-service/src/pipelines/generation/sql_breakdown.py

@@ -1,6 +1,6 @@
 import logging
 import sys
-from typing import Any
+from typing import Any, Optional
 
 from hamilton import base
 from hamilton.async_driver import AsyncDriver
@@ -135,10 +135,13 @@ async def generate_sql_details(prompt: dict, generator: Any) -> dict:
 async def post_process(
     generate_sql_details: dict,
     post_processor: SQLBreakdownGenPostProcessor,
+    engine_timeout: float,
     project_id: str | None = None,
 ) -> dict:
     return await post_processor.run(
-        generate_sql_details.get("replies"), project_id=project_id
+        generate_sql_details.get("replies"),
+        timeout=engine_timeout,
+        project_id=project_id,
     )
 
 
@@ -170,6 +173,7 @@ def __init__(
         self,
         llm_provider: LLMProvider,
         engine: Engine,
+        engine_timeout: Optional[float] = 30.0,
         **kwargs,
     ):
         self._components = {
@@ -185,6 +189,7 @@ def __init__(
 
         self._configs = {
             "text_to_sql_rules": TEXT_TO_SQL_RULES,
+            "engine_timeout": engine_timeout,
         }
 
         super().__init__(

wren-ai-service/src/pipelines/generation/sql_correction.py

@@ -1,7 +1,7 @@
 import asyncio
 import logging
 import sys
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from hamilton import base
 from hamilton.async_driver import AsyncDriver
@@ -82,9 +82,14 @@ async def generate_sql_corrections(prompts: list[dict], generator: Any) -> list[
 async def post_process(
     generate_sql_corrections: list[dict],
     post_processor: SQLGenPostProcessor,
+    engine_timeout: float,
     project_id: str | None = None,
 ) -> list[dict]:
-    return await post_processor.run(generate_sql_corrections, project_id=project_id)
+    return await post_processor.run(
+        generate_sql_corrections,
+        timeout=engine_timeout,
+        project_id=project_id,
+    )
 
 
 ## End of Pipeline
@@ -106,6 +111,7 @@ def __init__(
         self,
         llm_provider: LLMProvider,
         engine: Engine,
+        engine_timeout: Optional[float] = 30.0,
         **kwargs,
     ):
         self._components = {
@@ -119,6 +125,10 @@ def __init__(
             "post_processor": SQLGenPostProcessor(engine=engine),
         }
 
+        self._configs = {
+            "engine_timeout": engine_timeout,
+        }
+
         super().__init__(
             AsyncDriver({}, sys.modules[__name__], result_builder=base.DictResult())
         )
@@ -138,6 +148,7 @@ async def run(
                 "documents": contexts,
                 "project_id": project_id,
                 **self._components,
+                **self._configs,
             },
         )
 

wren-ai-service/src/pipelines/generation/sql_expansion.py

@@ -1,6 +1,6 @@
 import logging
 import sys
-from typing import Any, List
+from typing import Any, List, Optional
 
 from hamilton import base
 from hamilton.async_driver import AsyncDriver
@@ -75,10 +75,13 @@ async def generate_sql_expansion(prompt: dict, generator: Any) -> dict:
 async def post_process(
     generate_sql_expansion: dict,
     post_processor: SQLGenPostProcessor,
+    engine_timeout: float,
     project_id: str | None = None,
 ) -> dict:
     return await post_processor.run(
-        generate_sql_expansion.get("replies"), project_id=project_id
+        generate_sql_expansion.get("replies"),
+        timeout=engine_timeout,
+        project_id=project_id,
     )
 
 
@@ -105,6 +108,7 @@ def __init__(
         self,
         llm_provider: LLMProvider,
         engine: Engine,
+        engine_timeout: Optional[float] = 30.0,
         **kwargs,
     ):
         self._components = {
@@ -118,6 +122,10 @@ def __init__(
             "post_processor": SQLGenPostProcessor(engine=engine),
         }
 
+        self._configs = {
+            "engine_timeout": engine_timeout,
+        }
+
         super().__init__(
             AsyncDriver({}, sys.modules[__name__], result_builder=base.DictResult())
         )
@@ -141,6 +149,7 @@ async def run(
                 "project_id": project_id,
                 "configuration": configuration,
                 **self._components,
+                **self._configs,
             },
         )