From e5b68b42f289e98e5c5f8bb487564d7f4e967b5c Mon Sep 17 00:00:00 2001
From: Chih-Yu Yeh <chihyu.jimmy.yeh@gmail.com>
Date: Thu, 5 Sep 2024 15:33:40 +0800
Subject: [PATCH] chore(wren-ai-service): refine multi-language prompt (#648)

* refine prompt

* remove test deps
---
 wren-ai-service/docker/Dockerfile                |  2 +-
 wren-ai-service/src/pipelines/ask/sql_summary.py | 16 ++++++++++++----
 .../pipelines/ask_details/components/prompts.py  |  5 +++--
 .../src/pipelines/ask_details/generation.py      | 11 +++++++++--
 4 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/wren-ai-service/docker/Dockerfile b/wren-ai-service/docker/Dockerfile
index 1f0cf4080..a2bda1dd1 100644
--- a/wren-ai-service/docker/Dockerfile
+++ b/wren-ai-service/docker/Dockerfile
@@ -12,7 +12,7 @@ WORKDIR /app
 
 COPY pyproject.toml ./
 
-RUN poetry install --without dev,eval,demo --no-root && rm -rf $POETRY_CACHE_DIR
+RUN poetry install --without dev,eval,demo,test --no-root && rm -rf $POETRY_CACHE_DIR
 
 FROM python:3.12.0-slim-bookworm as runtime
 
diff --git a/wren-ai-service/src/pipelines/ask/sql_summary.py b/wren-ai-service/src/pipelines/ask/sql_summary.py
index b5beb1e9f..2ba66c6be 100644
--- a/wren-ai-service/src/pipelines/ask/sql_summary.py
+++ b/wren-ai-service/src/pipelines/ask/sql_summary.py
@@ -23,19 +23,27 @@ You are a great data analyst. You are now given a task to summarize a list SQL q
 You will be given a list of SQL queries and a user's question.
 
 ### INSTRUCTIONS ###
+- SQL query summary must be in the same language as the user's question.
+- PUT special care on whether it's traditional Chinese or simplified Chinese.
 - 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:
 
 {
+    "language": <LANGUAGE_IN_USER_QUESTION_STRING>,
     "sql_summary_results": [
-        {"summary": <SQL_QUERY_SUMMARY_1>},
-        {"summary": <SQL_QUERY_SUMMARY_2>},
-        {"summary": <SQL_QUERY_SUMMARY_3>}
+        {
+            "summary": <SQL_QUERY_SUMMARY_USING_SAME_LANGUAGE_AS_USER_QUESTION_1>
+        {
+            "summary": <SQL_QUERY_SUMMARY_USING_SAME_LANGUAGE_AS_USER_QUESTION_2>
+        },
+        {
+            "summary": <SQL_QUERY_SUMMARY_USING_SAME_LANGUAGE_AS_USER_QUESTION_3>
+        },
+        ...
     ]
 }
 """
diff --git a/wren-ai-service/src/pipelines/ask_details/components/prompts.py b/wren-ai-service/src/pipelines/ask_details/components/prompts.py
index c14ee452f..10117b544 100644
--- a/wren-ai-service/src/pipelines/ask_details/components/prompts.py
+++ b/wren-ai-service/src/pipelines/ask_details/components/prompts.py
@@ -13,8 +13,9 @@ The final step intentionally lacks a CTE name to simulate a final execution with
 - 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.
+### SUMMARY AND DESCRIPTION GENERATION INSTRUCTIONS ###
+- MUST BE in the same language as the user's question
+- PUT special care on whether it's traditional Chinese or simplified Chinese.
 - MUST BE human-readable and easy to understand.
 - MUST BE concise and to the point.
 
diff --git a/wren-ai-service/src/pipelines/ask_details/generation.py b/wren-ai-service/src/pipelines/ask_details/generation.py
index 43993a013..7cb2b5e73 100644
--- a/wren-ai-service/src/pipelines/ask_details/generation.py
+++ b/wren-ai-service/src/pipelines/ask_details/generation.py
@@ -33,13 +33,20 @@ SQL query: {{ sql }}
 The final answer must be a valid JSON format as following:
 
 {
+    "language": <LANGUAGE_IN_USER_QUESTION_STRING>,
     "description": <SHORT_SQL_QUERY_DESCRIPTION>,
     "steps: [
         {
             "sql": <SQL_QUERY_STRING_1>,
-            "summary": <SUMMARY_STRING_1>,
+            "summary": <SUMMARY_STRING_USING_SAME_LANGUAGE_AS_USER_QUESTION_1>,
             "cte_name": <CTE_NAME_STRING_1>
-        }
+        },
+        {
+            "sql": <SQL_QUERY_STRING_2>,
+            "summary": <SUMMARY_STRING_USING_SAME_LANGUAGE_AS_USER_QUESTION_2>,
+            "cte_name": <CTE_NAME_STRING_2>
+        },
+        ...
     ] # a list of steps
 }