feat(wren-ai-service): prompt enhancement for localization language c…

…ode and scoping the categories for questions generation (#955)
Canner · Nov 26, 2024 · 23061b5 · 23061b5
1 parent c6f8fdc
commit 23061b5
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 11 deletions.
diff --git a/wren-ai-service/src/pipelines/generation/question_recommendation.py b/wren-ai-service/src/pipelines/generation/question_recommendation.py
@@ -150,9 +150,9 @@ class QuestionResult(BaseModel):
 {% endif %}
 
 Current Date: {{current_date}}
-Language: {{language}}
+Localization Language: {{language}}
 
-Please generate {{max_questions}} insightful questions for each of the {{max_categories}} categories based on the provided data model{% if user_question %} and the user's question{% endif %}.
+Please generate {{max_questions}} insightful questions for each of the {{max_categories}} categories based on the provided data model, using the localization language provided{% if user_question %} and the user's question{% endif %}.
 """
 
 
@@ -181,7 +181,7 @@ def visualize(
         mdl: dict,
         previous_questions: list[str] = [],
         categories: list[str] = [],
-        language: str = "English",
+        language: str = "en",
         current_date: str = datetime.now(),
         max_questions: int = 5,
         max_categories: int = 3,
@@ -214,7 +214,7 @@ async def run(
         mdl: dict,
         previous_questions: list[str] = [],
         categories: list[str] = [],
-        language: str = "English",
+        language: str = "en",
         current_date: str = datetime.now().strftime("%Y-%m-%d %A %H:%M:%S"),
         max_questions: int = 5,
         max_categories: int = 3,
@@ -245,7 +245,7 @@ async def run(
         mdl={},
         previous_questions=[],
         categories=[],
-        language="English",
+        language="en",
         current_date=datetime.now().strftime("%Y-%m-%d %A %H:%M:%S"),
         max_questions=5,
         max_categories=3,

diff --git a/wren-ai-service/src/pipelines/generation/semantics_description.py b/wren-ai-service/src/pipelines/generation/semantics_description.py
@@ -185,7 +185,7 @@ class SemanticResult(BaseModel):
 ### Input:
 User's prompt: {{ user_prompt }}
 Picked models: {{ picked_models }}
-Language: {{ language }}
+Localization Language: {{ language }}
 
 Please provide a brief description for the model and each column based on the user's prompt.
 """
@@ -211,7 +211,7 @@ def visualize(
         user_prompt: str,
         selected_models: list[str],
         mdl: dict,
-        language: str = "English",
+        language: str = "en",
     ) -> None:
         destination = "outputs/pipelines/generation"
         if not Path(destination).exists():
@@ -237,7 +237,7 @@ async def run(
         user_prompt: str,
         selected_models: list[str],
         mdl: dict,
-        language: str = "English",
+        language: str = "en",
     ) -> dict:
         logger.info("Semantics Description Generation pipeline is running...")
         return await self._pipe.execute(
@@ -261,5 +261,5 @@ async def run(
         user_prompt="Track student enrollments, grades, and GPA calculations to monitor academic performance and identify areas for student support",
         selected_models=[],
         mdl={},
-        language="English",
+        language="en",
     )
diff --git a/wren-ai-service/src/web/v1/services/__init__.py b/wren-ai-service/src/web/v1/services/__init__.py
@@ -41,7 +41,7 @@ def show_current_time(self):
         return f'{current_time.strftime("%Y-%m-%d %A %H:%M:%S")}'  # YYYY-MM-DD weekday_name HH:MM:SS, ex: 2024-10-23 Wednesday 12:00:00
 
     fiscal_year: Optional[FiscalYear] = None
-    language: Optional[str] = "English"
+    language: Optional[str] = "en"
     timezone: Optional[Timezone] = Timezone()
 
 

diff --git a/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py b/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py
@@ -27,7 +27,7 @@ async def test_recommend_success(relationship_recommendation_service, mock_pipel
     assert response.id == "test_id"
     assert response.status == "finished"
     assert response.response == {"test": "data"}
-    mock_pipeline.run.assert_called_once_with(mdl={"key": "value"}, language="English")
+    mock_pipeline.run.assert_called_once_with(mdl={"key": "value"}, language="en")
 
 
 @pytest.mark.asyncio