From f1cad2f69e5dbcebaccc249f565c4cc6ed89b36d Mon Sep 17 00:00:00 2001
From: ZJU_czx <952370295@qq.com>
Date: Thu, 3 Apr 2025 18:48:41 +0800
Subject: [PATCH] feat: optimize chart generation file name

---
 app/prompt/visualization.py                    |  3 ++-
 app/tool/chart_visualization/chart_prepare.py  |  2 +-
 .../chart_visualization/data_visualization.py  |  6 +++++-
 .../chart_visualization/src/chartVisualize.ts  | 18 +++++++++++-------
 4 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/app/prompt/visualization.py b/app/prompt/visualization.py
index 8e4fecc..f8b1813 100644
--- a/app/prompt/visualization.py
+++ b/app/prompt/visualization.py
@@ -1,7 +1,8 @@
 SYSTEM_PROMPT = """You are an AI agent designed to data analysis / visualization task. You have various tools at your disposal that you can call upon to efficiently complete complex requests.
 # Note:
 1. The workspace directory is: {directory}; Read / write file in workspace
-2. Generate analysis conclusion report in the end"""
+2. Generate analysis conclusion report in the end
+3. Use at most One Tool in each step"""
 
 NEXT_STEP_PROMPT = """Based on user needs, break down the problem and use different tools step by step to solve it.
 # Note
diff --git a/app/tool/chart_visualization/chart_prepare.py b/app/tool/chart_visualization/chart_prepare.py
index 244535b..aa98e8b 100644
--- a/app/tool/chart_visualization/chart_prepare.py
+++ b/app/tool/chart_visualization/chart_prepare.py
@@ -28,7 +28,7 @@ class VisualizationPrepare(NormalPythonExecute):
 3. Save information in json file.( format: {"csvFilePath": string, "chartTitle": string}[])
 ## Insight Type
 1. Select the insights from the data_visualization results that you want to add to the chart.
-2. Save information in json file.( format: {"chartPath": string, "insights_id": number[]}[])
+2. Save information in json file.( format: {"chartPath": string, "insights_id": number[]}[]) (chartPath is output by data_visualization tool)
 # Note
 1. You can generate one or multiple csv data with different visualization needs.
 2. Make each chart data esay, clean and different.
diff --git a/app/tool/chart_visualization/data_visualization.py b/app/tool/chart_visualization/data_visualization.py
index 6bab11e..efd0c1a 100644
--- a/app/tool/chart_visualization/data_visualization.py
+++ b/app/tool/chart_visualization/data_visualization.py
@@ -195,7 +195,11 @@ Outputs:
             return {"observation": f"{success_template}"}
 
     async def execute(
-        self, json_path: str, output_type: str, tool_type: str, language: str
+        self,
+        json_path: str,
+        output_type: str | None = "html",
+        tool_type: str | None = "visualization",
+        language: str | None = "en",
     ) -> str:
         try:
             logger.info(f"📈 data_visualization with {json_path} in: {tool_type} ")
diff --git a/app/tool/chart_visualization/src/chartVisualize.ts b/app/tool/chart_visualization/src/chartVisualize.ts
index 9d83ec7..3153932 100644
--- a/app/tool/chart_visualization/src/chartVisualize.ts
+++ b/app/tool/chart_visualization/src/chartVisualize.ts
@@ -140,7 +140,10 @@ const setInsightTemplate = (
   }
   if (res) {
     fs.writeFileSync(path, res, "utf-8");
-    return { insight_path: path, insight_md: res };
+    return {
+      insight_path: path,
+      insight_md: res + "\n ## Next Step\nConsider add insights into chart",
+    };
   }
   return {};
 };
@@ -191,7 +194,7 @@ async function generateChart(
   } = {};
   const {
     dataset,
-    userPrompt,
+    userPrompt = "",
     directory,
     width,
     height,
@@ -199,11 +202,12 @@ async function generateChart(
     fileName,
     language,
   } = options;
+  const chartTitle = userPrompt.trim();
   try {
     // Get chart spec and save in local file
     const jsonDataset = isString(dataset) ? JSON.parse(dataset) : dataset;
     const { spec, error, chartType } = await vmind.generateChart(
-      userPrompt,
+      chartTitle,
       undefined,
       jsonDataset,
       {
@@ -218,18 +222,18 @@ async function generateChart(
     }
 
     spec.title = {
-      text: userPrompt,
+      text: chartTitle,
     };
     if (!fs.existsSync(path.join(directory, "visualization"))) {
       fs.mkdirSync(path.join(directory, "visualization"));
     }
-    const specPath = getSavedPathName(directory, fileName, "json");
+    const specPath = getSavedPathName(directory, chartTitle, "json");
     res.chart_path = await saveChartRes({
       directory,
       spec,
       width,
       height,
-      fileName,
+      fileName: chartTitle,
       outputType,
     });
 
@@ -274,7 +278,7 @@ async function generateChart(
     res = {
       ...res,
       ...setInsightTemplate(
-        getSavedPathName(directory, fileName, "md"),
+        getSavedPathName(directory, chartTitle, "md"),
         userPrompt,
         insightsText
       ),