chore(wren-ai-service): improve chart generation (#1101)

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

@@ -25,20 +25,20 @@
 chart_adjustment_system_prompt = f"""
 ### TASK ###
 
-You are a data analyst great at visualizing data using vega-lite! Given the data using the 'columns' formatted JSON from pandas.DataFrame APIs,
-original question, SQL query, vega-lite schema and the adjustment options, you need to regenerate vega-lite schema in JSON and provide suitable chart type according to the adjustment options;
-Besides, you need to give a concise and easy-to-understand reasoning to describe why you provide such vega-lite schema.
+You are a data analyst great at visualizing data using vega-lite! Given the user's question, SQL, data, vega-lite schema and adjustment options, 
+you need to re-generate vega-lite schema in JSON and provide suitable chart type.
+Besides, you need to give a concise and easy-to-understand reasoning within 20 words to describe why you provide such vega-lite schema.
 
 {chart_generation_instructions}
 - If you think the adjustment options are not suitable for the data, you can return an empty string for the schema and chart type and give reasoning to explain why.
 
 ### OUTPUT FORMAT ###
 
-Please provide your chain of thought reasoning and the vega-lite schema in JSON format.
+Please provide your chain of thought reasoning, chart type and the vega-lite schema in JSON format.
 
 {{
     "reasoning": <REASON_TO_CHOOSE_THE_SCHEMA_IN_STRING_FORMATTED_IN_LANGUAGE_PROVIDED_BY_USER>,
-    "chart_type": "line" | "bar" | "pie" | "grouped_bar" | "stacked_bar" | "area" | "",
+    "chart_type": "line" | "multi_line" | "bar" | "pie" | "grouped_bar" | "stacked_bar" | "area" | "",
     "chart_schema": <VEGA_LITE_JSON_SCHEMA>
 }}
 """

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

@@ -24,19 +24,18 @@
 chart_generation_system_prompt = f"""
 ### TASK ###
 
-You are a data analyst great at visualizing data using vega-lite! Given the data using the 'columns' formatted JSON from pandas.DataFrame APIs, 
-you need to generate vega-lite schema in JSON and provide suitable chart type; we will also give you the question and sql to understand the data.
-Besides, you need to give a concise and easy-to-understand reasoning to describe why you provide such vega-lite schema and a within 20 words description of the chart.
+You are a data analyst great at visualizing data using vega-lite! Given the user's question, SQL and data, you need to generate vega-lite schema in JSON and provide suitable chart type.
+Besides, you need to give a concise and easy-to-understand reasoning within 20 words to describe why you provide such vega-lite schema.
 
 {chart_generation_instructions}
 
 ### OUTPUT FORMAT ###
 
-Please provide your chain of thought reasoning, the vega-lite schema in JSON format and the chart type.
+Please provide your chain of thought reasoning, chart type and the vega-lite schema in JSON format.
 
 {{
     "reasoning": <REASON_TO_CHOOSE_THE_SCHEMA_IN_STRING_FORMATTED_IN_LANGUAGE_PROVIDED_BY_USER>,
-    "chart_type": "line" | "bar" | "pie" | "grouped_bar" | "stacked_bar" | "area" | "",
+    "chart_type": "line" | "multi_line" | "bar" | "pie" | "grouped_bar" | "stacked_bar" | "area" | "",
     "chart_schema": <VEGA_LITE_JSON_SCHEMA>
 }}
 """

wren-ai-service/src/pipelines/generation/utils/chart.py

@@ -8,7 +8,7 @@
 ### INSTRUCTIONS ###
 
 - Please generate vega-lite schema using v5 version, which is https://vega.github.io/schema/vega-lite/v5.json
-- Chart types: Bar chart, Line chart, Area chart, Pie chart, Stacked bar chart, Grouped bar chart
+- Chart types: Bar chart, Line chart, Multi line chart, Area chart, Pie chart, Stacked bar chart, Grouped bar chart
 - You can only use the chart types provided in the instructions
 - If the sample data is not suitable for visualization, you must return an empty string for the schema and chart type
 - If the sample data is empty, you must return an empty string for the schema and chart type
@@ -30,6 +30,7 @@
     - For daily question, the time unit should be "yearmonthdate".
     - Default time unit is "yearmonth".
 - For each axis, generate the corresponding human-readable title based on the language provided by the user.
+- Make sure all of the fields(x, y, xOffset, color, etc.) in the encoding section of the Vega-Lite schema are present in the column names of the data.
 
 ### GUIDELINES TO PLOT CHART ###
 
@@ -57,6 +58,15 @@
         - One temporal or ordinal variable (x-axis).
         - One quantitative variable (y-axis).
     - Example: Tracking monthly revenue over a year.
+- Multi Line Chart
+    - Use When: Displaying trends over continuous data, especially time.
+    - Data Requirements:
+        - One temporal or ordinal variable (x-axis).
+        - Two or more quantitative variables (y-axis and color).
+    - Implementation Notes:
+        - Uses `transform` with `fold` to combine multiple metrics into a single series
+        - The folded metrics are distinguished using the color encoding
+    - Example: Tracking monthly click rate and read rate over a year.
 - Area Chart
     - Use When: Similar to line charts but emphasizing the volume of change over time.
     - Data Requirements:
@@ -212,6 +222,32 @@
         "color": {"field": "Product", "type": "nominal", "title": "<TITLE_IN_LANGUAGE_PROVIDED_BY_USER>"}
     }
 }
+7. Multi Line Chart
+- Vega-Lite Spec:
+{
+    "$schema": "https://vega.github.io/schema/vega-lite/v5.json",
+    "title": <TITLE_IN_LANGUAGE_PROVIDED_BY_USER>,
+    "data": {
+        "values": [
+            {"Date": "2022-01-01", "readCount": 100, "clickCount": 10},
+            {"Date": "2022-01-02", "readCount": 200, "clickCount": 30},
+            {"Date": "2022-01-03", "readCount": 300, "clickCount": 20},
+            {"Date": "2022-01-04", "readCount": 400, "clickCount": 40}
+        ]
+    },
+    "mark": {"type": "line"},
+    "transform": [
+        {
+        "fold": ["readCount", "clickCount"],
+        "as": ["Metric", "Value"]
+        }
+    ],
+    "encoding": {
+        "x": {"field": "Date", "type": "temporal", "title": <TITLE_IN_LANGUAGE_PROVIDED_BY_USER>},
+        "y": {"field": "Value", "type": "quantitative", "title": <TITLE_IN_LANGUAGE_PROVIDED_BY_USER>},
+        "color": {"field": "Metric", "type": "nominal", "title": <TITLE_IN_LANGUAGE_PROVIDED_BY_USER>}
+    }
+}
 """
 
 
@@ -279,6 +315,24 @@ class LineChartEncoding(BaseModel):
     encoding: LineChartEncoding
 
 
+class MultiLineChartSchema(ChartSchema):
+    class MultiLineChartMark(BaseModel):
+        type: Literal["line"] = Field(default="line")
+
+    class MultiLineChartTransform(BaseModel):
+        fold: list[str]
+        as_: list[str] = Field(alias="as")
+
+    class MultiLineChartEncoding(BaseModel):
+        x: TemporalChartEncoding | ChartSchema.ChartEncoding
+        y: ChartSchema.ChartEncoding
+        color: ChartSchema.ChartEncoding
+
+    mark: MultiLineChartMark
+    transform: list[MultiLineChartTransform]
+    encoding: MultiLineChartEncoding
+
+
 class BarChartSchema(ChartSchema):
     class BarChartMark(BaseModel):
         type: Literal["bar"] = Field(default="bar")
@@ -349,9 +403,12 @@ class AreaChartEncoding(BaseModel):
 
 class ChartGenerationResults(BaseModel):
     reasoning: str
-    chart_type: Literal["line", "bar", "pie", "grouped_bar", "stacked_bar", "area", ""]
+    chart_type: Literal[
+        "line", "multi_line", "bar", "pie", "grouped_bar", "stacked_bar", "area", ""
+    ]  # empty string for no chart
     chart_schema: (
         LineChartSchema
+        | MultiLineChartSchema
         | BarChartSchema
         | PieChartSchema
         | GroupedBarChartSchema

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

@@ -64,7 +64,9 @@ class ChartResultRequest(BaseModel):
 
 class ChartResult(BaseModel):
     reasoning: str
-    chart_type: Literal["line", "bar", "pie", "grouped_bar", "stacked_bar", "area", ""]
+    chart_type: Literal[
+        "line", "bar", "pie", "grouped_bar", "stacked_bar", "area", "multi_line", ""
+    ]  # empty string for no chart
     chart_schema: dict
 
 

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

@@ -14,7 +14,9 @@
 
 # POST /v1/chart-adjustments
 class ChartAdjustmentOption(BaseModel):
-    chart_type: Literal["bar", "grouped_bar", "line", "pie", "stacked_bar", "area"]
+    chart_type: Literal[
+        "bar", "grouped_bar", "line", "pie", "stacked_bar", "area", "multi_line"
+    ]
     x_axis: Optional[str] = None
     y_axis: Optional[str] = None
     x_offset: Optional[str] = None
@@ -75,7 +77,9 @@ class ChartAdjustmentResultRequest(BaseModel):
 
 class ChartAdjustmentResult(BaseModel):
     reasoning: str
-    chart_type: Literal["line", "bar", "pie", "grouped_bar", "stacked_bar", "area", ""]
+    chart_type: Literal[
+        "line", "bar", "pie", "grouped_bar", "stacked_bar", "area", "multi_line", ""
+    ]  # empty string for no chart
     chart_schema: dict