chore(wren-ai-service): improve text2sql

wren-ai-service/eval/pipelines.py

@@ -30,9 +30,10 @@
 from src.pipelines import generation, indexing, retrieval
 
 
-def deploy_model(mdl: str, pipe: indexing.DBSchema) -> None:
+def deploy_model(mdl: str, pipes: list) -> None:
     async def wrapper():
-        await pipe.run(orjson.dumps(mdl).decode())
+        tasks = [pipe.run(orjson.dumps(mdl).decode()) for pipe in pipes]
+        await asyncio.gather(*tasks)
 
     asyncio.run(wrapper())
 
@@ -112,7 +113,8 @@ async def wrapper(batch: list):
         return [prediction for batch in batches for prediction in batch]
 
     @abstractmethod
-    def _process(self, prediction: dict, **_) -> dict: ...
+    def _process(self, prediction: dict, **_) -> dict:
+        ...
 
     async def _flat(self, prediction: dict, **_) -> dict:
         """
@@ -190,7 +192,9 @@ def __init__(
 
     async def _process(self, prediction: dict, **_) -> dict:
         result = await self._retrieval.run(query=prediction["input"])
-        documents = result.get("construct_retrieval_results", [])
+        documents = result.get("construct_retrieval_results", {}).get(
+            "retrieval_results", []
+        )
         prediction["retrieval_context"] = extract_units(documents)
 
         return prediction
@@ -241,7 +245,8 @@ async def _process(self, prediction: dict, document: list, **_) -> dict:
             query=prediction["input"],
             contexts=documents,
             samples=prediction["samples"],
-            exclude=[],
+            has_calculated_field=prediction.get("has_calculated_field", False),
+            has_metric=prediction.get("has_metric", False),
         )
 
         prediction["actual_output"] = actual_output
@@ -266,7 +271,10 @@ async def __call__(self, query: str, **_):
     def metrics(engine_info: dict, enable_semantics_comparison: bool) -> dict:
         return {
             "metrics": [
-                AccuracyMetric(engine_info=engine_info, enable_semantics_comparison=enable_semantics_comparison),
+                AccuracyMetric(
+                    engine_info=engine_info,
+                    enable_semantics_comparison=enable_semantics_comparison,
+                ),
                 AnswerRelevancyMetric(engine_info=engine_info),
                 FaithfulnessMetric(engine_info=engine_info),
                 # this is for spider dataset, rn we temporarily disable it
@@ -288,11 +296,14 @@ def __init__(
     ):
         super().__init__(meta)
 
-        _indexing = indexing.DBSchema(
+        _db_schema_indexing = indexing.DBSchema(
             **pipe_components["db_schema_indexing"],
             column_batch_size=settings.column_indexing_batch_size,
         )
-        deploy_model(mdl, _indexing)
+        _table_description_indexing = indexing.TableDescription(
+            **pipe_components["table_description_indexing"],
+        )
+        deploy_model(mdl, [_db_schema_indexing, _table_description_indexing])
 
         self._retrieval = retrieval.Retrieval(
             **pipe_components["db_schema_retrieval"],
@@ -315,15 +326,23 @@ async def _flat(self, prediction: dict, actual: str) -> dict:
 
     async def _process(self, prediction: dict, **_) -> dict:
         result = await self._retrieval.run(query=prediction["input"])
-        documents = result.get("construct_retrieval_results", [])
+        _retrieval_result = result.get("construct_retrieval_results", {})
+
+        documents = _retrieval_result.get("retrieval_results", [])
+        has_calculated_field = _retrieval_result.get("has_calculated_field", False)
+        has_metric = _retrieval_result.get("has_metric", False)
         actual_output = await self._generation.run(
             query=prediction["input"],
             contexts=documents,
             sql_samples=[],
+            has_calculated_field=has_calculated_field,
+            has_metric=has_metric,
         )
 
         prediction["actual_output"] = actual_output
         prediction["retrieval_context"] = extract_units(documents)
+        prediction["has_calculated_field"] = has_calculated_field
+        prediction["has_metric"] = has_metric
 
         return prediction
 

wren-ai-service/eval/prediction.py

@@ -85,7 +85,7 @@ def parse_args() -> Tuple[str, str]:
         "--file",
         "-F",
         type=str,
-        help="Eval dataset file name in the eval/dataset folder",
+        help="Eval dataset file path",
     )
     parser.add_argument(
         "--pipeline",

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

@@ -27,29 +27,6 @@ async def execute_sql(
         ...
 
 
-def clean_generation_result(result: str) -> str:
-    def _normalize_whitespace(s: str) -> str:
-        return re.sub(r"\s+", " ", s).strip()
-
-    def _remove_invalid_json_chars(s: str) -> str:
-        pattern = r'(\{\s*"results"\s*:\s*\[\s*\{\s*"sql"\s*:\s*".*?"\s*\}\s*\]\s*\})'
-        matches = re.findall(pattern, s, re.DOTALL)
-        if matches:
-            return matches[0]
-        return s
-
-    return (
-        _normalize_whitespace(_remove_invalid_json_chars(result))
-        .replace("\\n", " ")
-        .replace("```sql", "")
-        .replace("```json", "")
-        .replace('"""', "")
-        .replace("'''", "")
-        .replace("```", "")
-        .replace(";", "")
-    )
-
-
 def remove_limit_statement(sql: str) -> str:
     pattern = r"\s*LIMIT\s+\d+(\s*;?\s*--.*|\s*;?\s*)$"
     modified_sql = re.sub(pattern, "", sql, flags=re.IGNORECASE)

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

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
 from haystack import Document, component
 
@@ -7,11 +7,15 @@
 
 def build_table_ddl(
     content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None
-) -> str:
+) -> Tuple[str, bool]:
     columns_ddl = []
+    has_calculated_field = False
+
     for column in content["columns"]:
         if column["type"] == "COLUMN":
             if not columns or (columns and column["name"] in columns):
+                if "This column is a Calculated Field" in column["comment"]:
+                    has_calculated_field = True
                 column_ddl = (
                     f"{column['comment']}{column['name']} {column['data_type']}"
                 )
@@ -26,7 +30,7 @@ def build_table_ddl(
         f"{content['comment']}CREATE TABLE {content['name']} (\n  "
         + ",\n  ".join(columns_ddl)
         + "\n);"
-    )
+    ), has_calculated_field
 
 
 def dry_run_pipeline(pipeline_cls: BasicPipeline, pipeline_name: str, **kwargs):

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

@@ -6,13 +6,12 @@
 from hamilton.async_driver import AsyncDriver
 from haystack.components.builders.prompt_builder import PromptBuilder
 from langfuse.decorators import observe
-from pydantic import BaseModel
 
 from src.core.engine import Engine
 from src.core.pipeline import BasicPipeline
 from src.core.provider import LLMProvider
 from src.pipelines.generation.utils.sql import (
-    TEXT_TO_SQL_RULES,
+    SqlGenerationResult,
     SQLGenPostProcessor,
     construct_instructions,
     sql_generation_system_prompt,
@@ -33,74 +32,9 @@
     {{ document }}
 {% endfor %}
 
-### EXAMPLES ###
-
-Example 1
-[INPUT]
-Previous SQL Summary: A query to find the number of employees in each department.
-Previous SQL Query: SELECT department, COUNT(*) as employee_count FROM employees GROUP BY department;
-User's Question: How do I modify this to only show departments with more than 10 employees?
-
-[OUTPUT]
-{
-    "results": [
-        {
-            "sql": "SELECT department, COUNT() as employee_count FROM employees GROUP BY department HAVING COUNT() > 10"
-        }
-    ]
-}
-
-Example 2
-[INPUT]
-Previous SQL Summary: A query to retrieve the total sales per product.
-Previous SQL Query: SELECT product_id, SUM(sales) as total_sales FROM sales GROUP BY product_id;
-User's Question: Can you adjust this to include the product name as well?
-
-[OUTPUT]
-{
-    "results": [
-        {
-            "sql": "SELECT products.name, SUM(sales.sales) as total_sales FROM sales JOIN products ON sales.product_id = products.id GROUP BY products.name"
-        }
-    ]
-}
-
-Example 3
-[INPUT]
-Previous SQL Summary: Query to find the highest salary in each department.
-Previous SQL Query: SELECT department_id, MAX(salary) as highest_salary FROM employees GROUP BY department_id;
-User's Question: What if I want to see the employee names with the highest salary in each department?
-
-[OUTPUT]
-{
-    "results": [
-        {
-            "sql": "SELECT department_id, employee_name, salary FROM employees WHERE (department_id, salary) IN (SELECT department_id, MAX(salary) FROM employees GROUP BY department_id)"
-        }
-    ]
-}
-
-### FINAL ANSWER FORMAT ###
-The final answer must be the JSON format like following:
-
-{
-    "results": [
-        {"sql": <SQL_QUERY_STRING>}
-    ]
-}
-
-{{ alert }}
-
-### CONTEXT ###
-Previous SQL Summary:
-{% for summary in previous_query_summaries %}
-    {{ summary }}
-{% endfor %}
-Previous SQL Query: {{ history.sql }}
-Current Time: {{ current_time }}
-
 {% if instructions %}
-Instructions: {{ instructions }}
+### INSTRUCTIONS ###
+{{ instructions }}
 {% endif %}
 
 {% if sql_samples %}
@@ -113,8 +47,16 @@
 {% endfor %}
 {% endif %}
 
-### INPUT ###
+### CONTEXT ###
+Previous SQL Summary:
+{% for summary in previous_query_summaries %}
+    {{ summary }}
+{% endfor %}
+Previous SQL Query: {{ history.sql }}
+
+### QUESTION ###
 User's Follow-up Question: {{ query }}
+Current Time: {{ current_time }}
 
 Let's think step by step.
 """
@@ -126,10 +68,11 @@ def prompt(
     query: str,
     documents: List[str],
     history: AskHistory,
-    alert: str,
     configuration: Configuration,
     prompt_builder: PromptBuilder,
     sql_samples: List[Dict] | None = None,
+    has_calculated_field: bool = False,
+    has_metric: bool = False,
 ) -> dict:
     previous_query_summaries = [step.summary for step in history.steps if step.summary]
 
@@ -138,8 +81,12 @@ def prompt(
         documents=documents,
         history=history,
         previous_query_summaries=previous_query_summaries,
-        alert=alert,
-        instructions=construct_instructions(configuration),
+        instructions=construct_instructions(
+            configuration,
+            has_calculated_field,
+            has_metric,
+            sql_samples,
+        ),
         current_time=configuration.show_current_time(),
         sql_samples=sql_samples,
     )
@@ -164,22 +111,8 @@ async def post_process(
 ## End of Pipeline
 
 
-class SQLResult(BaseModel):
-    sql: str
-
-
-class GenerationResults(BaseModel):
-    results: list[SQLResult]
-
-
 FOLLOWUP_SQL_GENERATION_MODEL_KWARGS = {
-    "response_format": {
-        "type": "json_schema",
-        "json_schema": {
-            "name": "sql_results",
-            "schema": GenerationResults.model_json_schema(),
-        },
-    }
+    "response_format": SqlGenerationResult,
 }
 
 
@@ -201,10 +134,6 @@ def __init__(
             "post_processor": SQLGenPostProcessor(engine=engine),
         }
 
-        self._configs = {
-            "alert": TEXT_TO_SQL_RULES,
-        }
-
         super().__init__(
             AsyncDriver({}, sys.modules[__name__], result_builder=base.DictResult())
         )
@@ -218,6 +147,8 @@ async def run(
         configuration: Configuration = Configuration(),
         sql_samples: List[Dict] | None = None,
         project_id: str | None = None,
+        has_calculated_field: bool = False,
+        has_metric: bool = False,
     ):
         logger.info("Follow-Up SQL Generation pipeline is running...")
         return await self._pipe.execute(
@@ -229,8 +160,9 @@ async def run(
                 "project_id": project_id,
                 "configuration": configuration,
                 "sql_samples": sql_samples,
+                "has_calculated_field": has_calculated_field,
+                "has_metric": has_metric,
                 **self._components,
-                **self._configs,
             },
         )
 

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

@@ -207,11 +207,8 @@ def construct_db_schemas(dbschema_retrieval: list[Document]) -> list[str]:
     db_schemas_in_ddl = []
     for table_schema in list(db_schemas.values()):
         if table_schema["type"] == "TABLE":
-            db_schemas_in_ddl.append(
-                build_table_ddl(
-                    table_schema,
-                )
-            )
+            ddl, _ = build_table_ddl(table_schema)
+            db_schemas_in_ddl.append(ddl)
 
     return db_schemas_in_ddl