chore(wren-ai-service): improve text2chart and text2answer (#1016)

Canner · Dec 19, 2024 · 5df4d03 · 5df4d03
1 parent 1eff501
commit 5df4d03
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 4 deletions.
diff --git a/wren-ai-service/src/pipelines/generation/chart_adjustment.py b/wren-ai-service/src/pipelines/generation/chart_adjustment.py
@@ -93,6 +93,10 @@ def run(
             generation_result = orjson.loads(replies[0])
             reasoning = generation_result.get("reasoning", "")
             if chart_schema := generation_result.get("chart_schema", {}):
+                # sometimes the chart_schema is still in string format
+                if isinstance(chart_schema, str):
+                    chart_schema = orjson.loads(chart_schema)
+
                 validate(chart_schema, schema=vega_schema)
                 chart_schema["data"]["values"] = []
                 return {

diff --git a/wren-ai-service/src/pipelines/generation/chart_generation.py b/wren-ai-service/src/pipelines/generation/chart_generation.py
@@ -72,6 +72,10 @@ def run(
             generation_result = orjson.loads(replies[0])
             reasoning = generation_result.get("reasoning", "")
             if chart_schema := generation_result.get("chart_schema", {}):
+                # sometimes the chart_schema is still in string format
+                if isinstance(chart_schema, str):
+                    chart_schema = orjson.loads(chart_schema)
+
                 validate(chart_schema, schema=vega_schema)
                 chart_schema["data"]["values"] = []
                 return {
@@ -168,7 +172,7 @@ class ChartGenerationResults(BaseModel):
     "response_format": {
         "type": "json_schema",
         "json_schema": {
-            "name": "matched_schema",
+            "name": "chart_generation_schema",
             "schema": ChartGenerationResults.model_json_schema(),
         },
     }

diff --git a/wren-ai-service/src/pipelines/generation/sql_answer.py b/wren-ai-service/src/pipelines/generation/sql_answer.py
@@ -17,14 +17,15 @@
 sql_to_answer_system_prompt = """
 ### TASK
 
-You are a data analyst that great at answering user's questions based on the data, sql so that even non technical users can easily understand.
+You are a data analyst that great at answering non-technical user's questions based on the data, sql so that even non technical users can easily understand.
 Please answer the user's question in concise and clear manner in Markdown format.
 
 ### INSTRUCTIONS
 
 1. Read the user's question and understand the user's intention.
 2. Read the sql and understand the data.
-4. Generate a consice and clear answer in string format and a reasoning process in string format to the user's question based on the data, sql.
+3. Make sure the answer is aimed for non-technical users, so don't mention any technical terms such as SQL syntax.
+4. Generate a concise and clear answer in string format to answerthe user's question based on the data and sql.
 5. If answer is in list format, only list top few examples, and tell users there are more results omitted.
 6. Answer must be in the same language user specified.
 

diff --git a/wren-ai-service/src/pipelines/retrieval/retrieval.py b/wren-ai-service/src/pipelines/retrieval/retrieval.py
@@ -366,7 +366,7 @@ class RetrievalResults(BaseModel):
     "response_format": {
         "type": "json_schema",
         "json_schema": {
-            "name": "matched_schema",
+            "name": "retrieval_schema",
             "schema": RetrievalResults.model_json_schema(),
         },
     }