feat(wren-ai-service): allow multi-language (#646)

* allow multi-language

* fix test

wren-ai-service/src/pipelines/ask/sql_summary.py

@@ -22,6 +22,11 @@ sql_summary_system_prompt = """
 You are a great data analyst. You are now given a task to summarize a list SQL queries in a human-readable format where each summary should be within 10-20 words.
 You will be given a list of SQL queries and a user's question.
 
+### INSTRUCTIONS ###
+- SQL query summary must be within 10-20 words.
+- SQL query summary must be human-readable and easy to understand.
+- SQL query summary must be concise and to the point.
+- SQL query summary must be in the same language as the user's question.
 
 ### OUTPUT FORMAT ###
 Please return the result in the following JSON format:

wren-ai-service/src/pipelines/ask_details/components/prompts.py

@@ -1,9 +1,23 @@
 ask_details_system_prompt = """
+### TASK ###
 You are a Trino SQL expert with exceptional logical thinking skills. 
 You are going to break a complex SQL query into 1 to 10 steps to make it easier to understand for end users.
 Each step should have a SQL query part, a summary explaining the purpose of that query, and a CTE name to link the queries. 
 The final step intentionally lacks a CTE name to simulate a final execution without a subsequent CTE.
 
+### SQL QUERY BREAKDOWN INSTRUCTIONS ###
+- YOU MUST BREAK DOWN any SQL query into small steps if there is JOIN operations or sub-queries.
+- ONLY USE the tables and columns mentioned in the original sql query.
+- ONLY CHOOSE columns belong to the tables mentioned in the database schema.
+- ONLY THE last step should not have a CTE name.
+- ALWAYS SHOW alias for columns and tables such as SELECT [column_name] AS [alias_column_name].
+- MUST USE alias from the original SQL query.
+
+### SUMMARY INSTRUCTIONS ###
+- MUST BE in the same language as the user's question.
+- MUST BE human-readable and easy to understand.
+- MUST BE concise and to the point.
+
 ### EXAMPLES ###
 
 Example 1:
@@ -42,12 +56,4 @@ Results:
     - sql: SELECT product_id FROM sales_data
     - summary: Selects product IDs from the sales_data table.
     - cte_name: <empty_string>
-
-### ALERT ###
-- YOU MUST BREAK DOWN any SQL query into small steps if there is JOIN operations or sub-queries.
-- ONLY USE the tables and columns mentioned in the original sql query.
-- ONLY CHOOSE columns belong to the tables mentioned in the database schema.
-- ONLY THE last step should not have a CTE name.
-- ALWAYS SHOW alias for columns and tables such as SELECT [column_name] AS [alias_column_name].
-- MUST USE alias from the original SQL query.
 """

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

@@ -26,6 +26,7 @@ logger = logging.getLogger("wren-ai-service")
 
 ask_details_user_prompt_template = """
 ### INPUT ###
+User's Question: {{ query }}
 SQL query: {{ sql }}
 
 ### FINAL ANSWER FORMAT ###
@@ -49,9 +50,10 @@ Let's think step by step.
 ## Start of Pipeline
 @timer
 @observe(capture_input=False)
-def prompt(sql: str, prompt_builder: PromptBuilder) -> dict:
+def prompt(query: str, sql: str, prompt_builder: PromptBuilder) -> dict:
+    logger.debug(f"query: {query}")
     logger.debug(f"sql: {sql}")
-    return prompt_builder.run(sql=sql)
+    return prompt_builder.run(query=query, sql=sql)
 
 
 @async_timer
@@ -95,7 +97,7 @@ class Generation(BasicPipeline):
             AsyncDriver({}, sys.modules[__name__], result_builder=base.DictResult())
         )
 
-    def visualize(self, sql: str, project_id: str | None = None) -> None:
+    def visualize(self, query: str, sql: str, project_id: str | None = None) -> None:
         destination = "outputs/pipelines/ask_details"
         if not Path(destination).exists():
             Path(destination).mkdir(parents=True, exist_ok=True)
@@ -104,6 +106,7 @@ class Generation(BasicPipeline):
             ["post_process"],
             output_file_path=f"{destination}/generation.dot",
             inputs={
+                "query": query,
                 "sql": sql,
                 "generator": self.generator,
                 "prompt_builder": self.prompt_builder,
@@ -116,11 +119,12 @@ class Generation(BasicPipeline):
 
     @async_timer
     @observe(name="Ask_Details Generation")
-    async def run(self, sql: str, project_id: str | None = None) -> dict:
+    async def run(self, query: str, sql: str, project_id: str | None = None) -> dict:
         logger.info("Ask_Details Generation pipeline is running...")
         return await self._pipe.execute(
             ["post_process"],
             inputs={
+                "query": query,
                 "sql": sql,
                 "generator": self.generator,
                 "prompt_builder": self.prompt_builder,
@@ -146,7 +150,7 @@ if __name__ == "__main__":
         engine=engine,
     )
 
-    pipeline.visualize("SELECT * FROM table_name")
-    async_validate(lambda: pipeline.run("SELECT * FROM table_name"))
+    pipeline.visualize("", "SELECT * FROM table_name")
+    async_validate(lambda: pipeline.run("", "SELECT * FROM table_name"))
 
     langfuse_context.flush()

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

@@ -101,6 +101,7 @@ class AskDetailsService:
             )
 
             generation_result = await self._pipelines["generation"].run(
+                query=ask_details_request.query,
                 sql=ask_details_request.sql,
                 project_id=ask_details_request.project_id,
             )

wren-ai-service/tests/pytest/pipelines/test_ask_details.py

@@ -28,5 +28,6 @@ async def test_generation_pipeline_producing_executable_sqls():
 
     for candidate_sql_query in candidate_sql_queries:
         assert await generation_pipeline.run(
+            "",
             candidate_sql_query,
         )