feat: update prompt of data analysis agent

app/agent/data_analysis.py

@@ -5,8 +5,8 @@ from app.config import config
 from app.prompt.visualization import NEXT_STEP_PROMPT, SYSTEM_PROMPT
 from app.tool import Terminate, ToolCollection
 from app.tool.chart_visualization.chart_prepare import VisualizationPrepare
-from app.tool.chart_visualization.chart_visualization import ChartVisualization
-from app.tool.chart_visualization.normal_python_execute import NormalPythonExecute
+from app.tool.chart_visualization.data_visualization import DataVisualization
+from app.tool.chart_visualization.python_execute import NormalPythonExecute
 
 
 class DataAnalysis(ToolCallAgent):
@@ -19,21 +19,21 @@ class DataAnalysis(ToolCallAgent):
 
     name: str = "DataAnalysis"
     description: str = (
-        "An analytical agent that utilizes multiple tools to solve diverse data tasks"
+        "An analytical agent that utilizes multiple tools to solve diverse data analysis tasks"
     )
 
     system_prompt: str = SYSTEM_PROMPT.format(directory=config.workspace_root)
     next_step_prompt: str = NEXT_STEP_PROMPT
 
-    max_observe: int = 10000
-    max_steps: int = 20
+    max_observe: int = 15000
+    max_steps: int = 30
 
     # Add general-purpose tools to the tool collection
     available_tools: ToolCollection = Field(
         default_factory=lambda: ToolCollection(
             NormalPythonExecute(),
-            # VisualizationPrepare(),
-            # ChartVisualization(),
+            VisualizationPrepare(),
+            DataVisualization(),
             Terminate(),
         )
     )

app/prompt/visualization.py

@@ -1,6 +1,10 @@
 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 Chart with insights may better for user's task"""
+2. Generate analysis conclusion report in the end"""
 
-NEXT_STEP_PROMPT = """Based on user needs, break down the problem and use different tools step by step to solve it. Each step select the most appropriate tool proactively (ONLY ONE). After using each tool, clearly explain the execution results and suggest the next steps."""
+NEXT_STEP_PROMPT = """Based on user needs, break down the problem and use different tools step by step to solve it.
+# Note
+1. Each step select the most appropriate tool proactively (ONLY ONE).
+2. After using each tool, clearly explain the execution results and suggest the next steps.
+3. When observation with Error, review and fix it."""

app/tool/chart_visualization/__init__.py

@@ -1,5 +1,5 @@
-from app.tool.chart_visualization.chart_visualization import ChartVisualization
+from app.tool.chart_visualization.data_visualization import DataVisualization
 from app.tool.chart_visualization.chart_prepare import VisualizationPrepare
-from app.tool.chart_visualization.normal_python_execute import NormalPythonExecute
+from app.tool.chart_visualization.python_execute import NormalPythonExecute
 
-__all__ = ["ChartVisualization", "VisualizationPrepare", "NormalPythonExecute"]
+__all__ = ["DataVisualization", "VisualizationPrepare", "NormalPythonExecute"]

app/tool/chart_visualization/chart_prepare.py

@@ -1,4 +1,4 @@
-from app.tool.chart_visualization.normal_python_execute import NormalPythonExecute
+from app.tool.chart_visualization.python_execute import NormalPythonExecute
 
 
 class VisualizationPrepare(NormalPythonExecute):

app/tool/chart_visualization/data_visualization.py

@@ -11,16 +11,15 @@ from app.logger import logger
 from app.config import config
 
 
-class ChartVisualization(BaseTool):
+class DataVisualization(BaseTool):
     name: str = "data_visualization"
     description: str = (
-        """Visualize statistical chart or Add insights in chart with JSON info from visualization_preparation tool.
-For each chart, you can do steps as follows:
+        """Visualize statistical chart or Add insights in chart with JSON info from visualization_preparation tool. You can do steps as follows:
 1. Visualize statistical chart
-2. Choose insights into chart based on step 1 (Optional.)
+2. Choose insights into chart based on step 1 (Optional)
 Outputs:
 1. Charts (png/html)
-2. Charts Insights (.md)(Optional)."""
+2. Charts Insights (.md)(Optional)"""
     )
     parameters: dict = {
         "type": "object",

app/tool/chart_visualization/normal_python_execute.py

@@ -1,24 +0,0 @@
-from app.tool.python_execute import PythonExecute
-
-
-class NormalPythonExecute(PythonExecute):
-    """A tool for executing Python code with timeout and safety restrictions."""
-
-    name: str = "common_python_execute"
-    description: str = (
-        "Execute Python code for in-depth data analysis without direct visualization. "
-        "The code should generate a comprehensive text-based report containing dataset overview, "
-        "column details, basic statistics, derived metrics, day-of-week comparisons, outliers, and key insights. "
-        "Use print() for all outputs so the analysis (including sections like 'Dataset Overview' or 'Preprocessing Results') "
-        "is clearly visible, and save any final report or processed files to config.workspace. "
-        "Include try-except blocks for error handling, and provide the code as a single string with '\\n' for line breaks."
-    )
-    parameters: dict = {
-        "type": "object",
-        "properties": {
-            "code": {
-                "type": "string",
-            },
-        },
-        "required": ["code"],
-    }

app/tool/chart_visualization/python_execute.py

@@ -0,0 +1,36 @@
+from app.tool.python_execute import PythonExecute
+from app.config import config
+
+
+class NormalPythonExecute(PythonExecute):
+    """A tool for executing Python code with timeout and safety restrictions."""
+
+    name: str = "python_execute"
+    description: str = (
+        """Execute Python code for in-depth data analysis / data report(task conclusion) / other normal task without direct visualization.
+# Note
+1. The code should generate a comprehensive text-based report containing dataset overview, column details, basic statistics, derived metrics, timeseries comparisons, outliers, and key insights.
+2. Use print() for all outputs so the analysis (including sections like 'Dataset Overview' or 'Preprocessing Results') is clearly visible and save it also
+3. Save any report / processed files / each analysis result in worksapce directory: {directory}
+4. You can invode this tool step-by-step to do data analysis from summary to in-depth""".format(
+            directory=config.workspace_root
+        )
+    )
+    parameters: dict = {
+        "type": "object",
+        "properties": {
+            "code_type": {
+                "description": "code type, data process / data report / others",
+                "type": "string",
+                "default": "process",
+                "enum": ["process", "report", "others"],
+            },
+            "code": {
+                "type": "string",
+            },
+        },
+        "required": ["code"],
+    }
+
+    async def execute(self, code: str, code_type: str | None = None, timeout=5):
+        return await super().execute(code, timeout)

app/tool/chart_visualization/test/tool_test.py

@@ -1,5 +1,5 @@
 import asyncio
-from app.tool.chart_visualization import ChartVisualization
+from app.tool.chart_visualization import DataVisualization
 
 
 async def mock_request(delay, value):
@@ -21,7 +21,7 @@ async def main():
 
 
 async def test_chart():
-    chartTool = ChartVisualization()
+    chartTool = DataVisualization()
     print(
         await chartTool.execute(
             "workspace/insight_metadata.json", "html", "insight", "en"