From 68635ff4aaac7af6abcf324a95baeb28cbd38cc2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Thu, 30 Nov 2023 15:43:13 +0800
Subject: [PATCH 01/19] add typing-extensions-4.8.0 for nbclient

---
 requirements.txt | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/requirements.txt b/requirements.txt
index c72260c04..1d1bc95a1 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -50,4 +50,5 @@ nbclient==0.9.0
 nbformat==5.9.2
 ipython==8.17.2
 ipykernel==6.27.0
-scikit_learn==1.3.2
\ No newline at end of file
+scikit_learn==1.3.2
+typing-extensions==4.8.0
\ No newline at end of file

From b28111ab3476f6ce51beffa60819ab82f1fc28b5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Thu, 30 Nov 2023 16:05:56 +0800
Subject: [PATCH 02/19] fix: "image/png" not in output["data"].

---
 metagpt/actions/execute_code.py            |   9 +-
 tests/metagpt/actions/test_execute_code.py | 114 +++++++++++++--------
 2 files changed, 76 insertions(+), 47 deletions(-)

diff --git a/metagpt/actions/execute_code.py b/metagpt/actions/execute_code.py
index 7b16d559a..981aa894c 100644
--- a/metagpt/actions/execute_code.py
+++ b/metagpt/actions/execute_code.py
@@ -17,6 +17,7 @@ from rich.syntax import Syntax
 
 from metagpt.actions import Action
 from metagpt.schema import Message
+from metagpt.logs import logger
 
 
 class ExecuteCode(ABC):
@@ -90,11 +91,14 @@ class ExecutePyCode(ExecuteCode, Action):
         if not outputs:
             return parsed_output
 
-        for output in outputs:
+        for i, output in enumerate(outputs):
             if output["output_type"] == "stream":
                 parsed_output += output["text"]
             elif output["output_type"] == "display_data":
-                self.show_bytes_figure(output["data"]["image/png"], self.interaction)
+                if "image/png" in output["data"]:
+                    self.show_bytes_figure(output["data"]["image/png"], self.interaction)
+                else:
+                    logger.info(f"{i}th output['data'] from nbclient outputs dont have image/png, continue next output ...")
             elif output["output_type"] == "execute_result":
                 parsed_output += output["data"]["text/plain"]
         return parsed_output
@@ -136,7 +140,6 @@ class ExecutePyCode(ExecuteCode, Action):
 
         if isinstance(code, str):
             return code, language
-
         if isinstance(code, dict):
             assert "code" in code
             if "language" not in code:
diff --git a/tests/metagpt/actions/test_execute_code.py b/tests/metagpt/actions/test_execute_code.py
index 88c5adf18..8894f2cb9 100644
--- a/tests/metagpt/actions/test_execute_code.py
+++ b/tests/metagpt/actions/test_execute_code.py
@@ -1,57 +1,83 @@
 import pytest
 
-from metagpt.actions import ExecutePyCode
+from metagpt.actions.execute_code import ExecutePyCode
 from metagpt.schema import Message
 
 
-@pytest.mark.asyncio
-async def test_code_running():
-    pi = ExecutePyCode()
-    output = await pi.run("print('hello world!')")
-    assert output.state == "done"
-    output = await pi.run({"code": "print('hello world!')", "language": "python"})
-    assert output.state == "done"
-    code_msg = Message("print('hello world!')")
-    output = await pi.run(code_msg)
-    assert output.state == "done"
+# @pytest.mark.asyncio
+# async def test_code_running():
+#     pi = ExecutePyCode()
+#     output = await pi.run("print('hello world!')")
+#     assert output[1] is True
+#     output = await pi.run({"code": "print('hello world!')", "language": "python"})
+#     assert output[1] is True
+#     code_msg = Message("print('hello world!')")
+#     output = await pi.run(code_msg)
+#     assert output[1] is True
+
+
+# @pytest.mark.asyncio
+# async def test_split_code_running():
+#     pi = ExecutePyCode()
+#     output = await pi.run("x=1\ny=2")
+#     output = await pi.run("z=x+y")
+#     output = await pi.run("assert z==3")
+#     assert output[1] is True
+
+
+# @pytest.mark.asyncio
+# async def test_execute_error():
+#     pi = ExecutePyCode()
+#     output = await pi.run("z=1/0")
+#     assert output[1] is False
+
+
+# @pytest.mark.asyncio
+# async def test_plotting_code():
+#     pi = ExecutePyCode()
+#     code = """
+#     import numpy as np
+#     import matplotlib.pyplot as plt
+
+#     # 生成随机数据
+#     random_data = np.random.randn(1000)  # 生成1000个符合标准正态分布的随机数
+
+#     # 绘制直方图
+#     plt.hist(random_data, bins=30, density=True, alpha=0.7, color='blue', edgecolor='black')
+
+#     # 添加标题和标签
+#     plt.title('Histogram of Random Data')
+#     plt.xlabel('Value')
+#     plt.ylabel('Frequency')
+
+#     # 显示图形
+#     plt.show()
+#     """
+#     output = await pi.run(code)
+#     assert output[1] is True
 
 
 @pytest.mark.asyncio
-async def test_split_code_running():
-    pi = ExecutePyCode()
-    output = await pi.run("x=1\ny=2")
-    output = await pi.run("z=x+y")
-    output = await pi.run("assert z==3")
-    assert output.state == "done"
-
-
-@pytest.mark.asyncio
-async def test_execute_error():
-    pi = ExecutePyCode()
-    output = await pi.run("z=1/0")
-    assert output.state == "error"
-
-
-@pytest.mark.asyncio
-async def test_plotting_code():
-    pi = ExecutePyCode()
+async def test_plotting_bug():
     code = """
-    import numpy as np
     import matplotlib.pyplot as plt
-
-    # 生成随机数据
-    random_data = np.random.randn(1000)  # 生成1000个符合标准正态分布的随机数
-
-    # 绘制直方图
-    plt.hist(random_data, bins=30, density=True, alpha=0.7, color='blue', edgecolor='black')
-
-    # 添加标题和标签
-    plt.title('Histogram of Random Data')
-    plt.xlabel('Value')
-    plt.ylabel('Frequency')
-
-    # 显示图形
+    import seaborn as sns
+    import pandas as pd
+    from sklearn.datasets import load_iris
+    # Load the Iris dataset
+    iris_data = load_iris()
+    # Convert the loaded Iris dataset into a DataFrame for easier manipulation
+    iris_df = pd.DataFrame(iris_data['data'], columns=iris_data['feature_names'])
+    # Add a column for the target
+    iris_df['species'] = pd.Categorical.from_codes(iris_data['target'], iris_data['target_names'])
+    # Set the style of seaborn
+    sns.set(style='whitegrid')
+    # Create a pairplot of the iris dataset
+    plt.figure(figsize=(10, 8))
+    pairplot = sns.pairplot(iris_df, hue='species')
+    # Show the plot
     plt.show()
     """
+    pi = ExecutePyCode()
     output = await pi.run(code)
-    assert output.state == "done"
+    assert output[1] is True

From 8aa096a33469fb5c3730d5c9b413772c1c7f2f8c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Thu, 30 Nov 2023 16:07:12 +0800
Subject: [PATCH 03/19] fix: remove escape and color codes for output of
 nbclient.

---
 metagpt/roles/ml_engineer.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index 1e4367372..5120a9011 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -3,6 +3,7 @@ import json
 import subprocess
 
 import fire
+import re
 
 from metagpt.roles import Role
 from metagpt.actions import Action
@@ -35,6 +36,13 @@ def truncate(result: str, keep_len: int = 1000) -> str:
     return desc
 
 
+def remove_escape_and_color_codes(input_str):
+    # 使用正则表达式去除转义字符和颜色代码
+    pattern = re.compile(r'\x1b\[[0-9;]*[mK]')
+    result = pattern.sub('', input_str)
+    return result
+
+
 class AskReview(Action):
     async def run(self, context: List[Message], plan: Plan = None):
         logger.info("Current overall plan:")
@@ -137,8 +145,9 @@ class MLEngineer(Role):
             # truncated the result
             print(truncate(result))
             # print(result)
+            _result = truncate(remove_escape_and_color_codes(result))
             self.working_memory.add(
-                Message(content=result, role="user", cause_by=ExecutePyCode)
+                Message(content=_result, role="user", cause_by=ExecutePyCode)
             )
 
             if "!pip" in code:

From b81fefffa17233ff0654395841e8d5bdd604a225 Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Thu, 30 Nov 2023 16:28:02 +0800
Subject: [PATCH 04/19] avoid repetitive tool desc between steps

---
 metagpt/actions/write_analysis_code.py | 22 +++++++++++++++-------
 metagpt/prompts/ml_engineer.py         |  6 +++++-
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/metagpt/actions/write_analysis_code.py b/metagpt/actions/write_analysis_code.py
index 787fb8d3e..6fff1c66f 100644
--- a/metagpt/actions/write_analysis_code.py
+++ b/metagpt/actions/write_analysis_code.py
@@ -5,7 +5,7 @@
 @File    :   write_code_v2.py
 """
 import json
-from typing import Dict, List, Union
+from typing import Dict, List, Union, Tuple
 
 from metagpt.actions import Action
 from metagpt.prompts.ml_engineer import (
@@ -100,24 +100,31 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
     """Write code with help of local available tools. Choose tools first, then generate code to use the tools"""
 
     @staticmethod
-    def _parse_recommend_tools(module: str, recommend_tools: list) -> str:
+    def _parse_recommend_tools(module: str, recommend_tools: list) -> Tuple[Dict, List[Dict]]:
         """
-        Converts recommended tools to a JSON string and checks tool availability in the registry.
+        Parses and validates a list of recommended tools, and retrieves their schema from registry.
 
         Args:
             module (str): The module name for querying tools in the registry.
             recommend_tools (list): A list of lists of recommended tools for each step.
 
         Returns:
-            str: A JSON string with available tools and their schemas for each step.
+            Tuple[Dict, List[Dict]]:
+                - valid_tools: A dict of lists of valid tools for each step.
+                - tool_catalog: A list of dicts of unique tool schemas.
         """
         valid_tools = {}
         available_tools = registry.get_all_by_module(module).keys()
         for index, tools in enumerate(recommend_tools):
             key = f"Step {index + 1}"
             tools = [tool for tool in tools if tool in available_tools]
-            valid_tools[key] = registry.get_schemas(module, tools)
-        return json.dumps(valid_tools)
+            valid_tools[key] = tools
+
+        unique_tools = set()
+        for tools in valid_tools.values():
+            unique_tools.update(tools)
+        tool_catalog = registry.get_schemas(module, unique_tools)
+        return valid_tools, tool_catalog
 
     async def _tool_recommendation(
         self, task: str, data_desc: str, code_steps: str, available_tools: list
@@ -166,7 +173,7 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
         recommend_tools = await self._tool_recommendation(
             task, task_guide, available_tools
         )
-        recommend_tools = self._parse_recommend_tools(task_type, recommend_tools)
+        recommend_tools, tool_catalog = self._parse_recommend_tools(task_type, recommend_tools)
 
         special_prompt = ML_SPECIFIC_PROMPT.get(task_type, "")
         module_name = ML_MODULE_MAP[task_type]
@@ -191,6 +198,7 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
             module_name=module_name,
             output_desc=output_desc,
             available_tools=recommend_tools,
+            tool_catalog=tool_catalog,
         )
         tool_config = create_func_config(CODE_GENERATOR_WITH_TOOLS)
         rsp = await self.llm.aask_code(prompt, **tool_config)
diff --git a/metagpt/prompts/ml_engineer.py b/metagpt/prompts/ml_engineer.py
index 55ac27d82..70a40ef34 100644
--- a/metagpt/prompts/ml_engineer.py
+++ b/metagpt/prompts/ml_engineer.py
@@ -95,9 +95,13 @@ from metagpt.tools.functions.libs.feature_engineering import fill_missing_value
 ```
 
 ## Available Functions for Each Step:
-Each function is described in JSON format, including the function name and parameters. {output_desc}
+Here's a list of all available functions for each step. You can find more details about each function in [## Function Catalog]
 {available_tools}
 
+## Function Catalog:
+Each function is described in JSON format, including the function name and parameters. {output_desc}
+{function_catalog}
+
 ## Your Output Format:
 Generate the complete code for every step, listing any used function tools at the beginning of the step:
 ```python

From 2dd754d97740243602edec17a4611bbaa8a0c0dd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Thu, 30 Nov 2023 16:36:35 +0800
Subject: [PATCH 05/19] fix: reuse variables.

---
 metagpt/actions/write_analysis_code.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/metagpt/actions/write_analysis_code.py b/metagpt/actions/write_analysis_code.py
index 66e2137fe..ee4555ee1 100644
--- a/metagpt/actions/write_analysis_code.py
+++ b/metagpt/actions/write_analysis_code.py
@@ -40,8 +40,8 @@ class BaseWriteAnalysisCode(Action):
 
 class WriteCodeByGenerate(BaseWriteAnalysisCode):
     """Write code fully by generation"""
-    DEFAULT_SYSTEM_MSG = """You are Code Interpreter, a world-class programmer that can complete any goal by executing code. Strictly follow the plan and generate code step by step. Each step of the code will be executed on the user's machine, and the user will provide the code execution results to you.**Notice: Use !pip install in a standalone block to install missing packages.**""" # prompt reference: https://github.com/KillianLucas/open-interpreter/blob/v0.1.4/interpreter/system_message.txt
-    REUSE_CODE_INSTRUCTION = """ATTENTION: DONT include codes from previous tasks in your current code block, include new codes only, DONT repeat codes!"""
+    DEFAULT_SYSTEM_MSG = """You are Code Interpreter, a world-class programmer that can complete any goal by executing code. Strictly follow the plan and generate code step by step. Each step of the code will be executed on the user's machine, and the user will provide the code execution results to you.**Notice: The code for the next step depends on the code for the previous step. Reuse variables in other code directly. Use !pip install in a standalone block to install missing packages.**""" # prompt reference: https://github.com/KillianLucas/open-interpreter/blob/v0.1.4/interpreter/system_message.txt
+    # REUSE_CODE_INSTRUCTION = """ATTENTION: DONT include codes from previous tasks in your current code block, include new codes only, DONT repeat codes!"""
 
     def __init__(self, name: str = "", context=None, llm=None) -> str:
         super().__init__(name, context, llm)
@@ -89,7 +89,7 @@ class WriteCodeByGenerate(BaseWriteAnalysisCode):
         system_msg: str = None,
         **kwargs,
     ) -> str:
-        context.append(Message(content=self.REUSE_CODE_INSTRUCTION, role="user"))
+        # context.append(Message(content=self.REUSE_CODE_INSTRUCTION, role="user"))
         prompt = self.process_msg(context, system_msg)
         code_content = await self.llm.aask_code(prompt, **kwargs)
         return code_content["code"]

From 9d49caa8cc8566aeee5a8f8a7ad0c22d1271dae6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Thu, 30 Nov 2023 17:09:43 +0800
Subject: [PATCH 06/19] test: set temperature=0.0

---
 tests/metagpt/actions/test_write_analysis_code.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/metagpt/actions/test_write_analysis_code.py b/tests/metagpt/actions/test_write_analysis_code.py
index 80d9438af..d4bccb552 100644
--- a/tests/metagpt/actions/test_write_analysis_code.py
+++ b/tests/metagpt/actions/test_write_analysis_code.py
@@ -159,7 +159,7 @@ async def test_write_code_reuse_code_long():
         Message(content=structural_context, role="user"),
     ]
     trials_num = 5
-    trials = [WriteCodeByGenerate().run(context=context) for _ in range(trials_num)]
+    trials = [WriteCodeByGenerate().run(context=context, temperature=0.0) for _ in range(trials_num)]
     trial_results = await asyncio.gather(*trials)
     print(*trial_results, sep="\n\n***\n\n")
     success = ["load_iris" not in result and "iris_data" in result \

From 870ece45b23dbdd27fb9407b8127865a21279d8b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Thu, 30 Nov 2023 17:10:36 +0800
Subject: [PATCH 07/19] fix: set temperature=0.0 for WriteCodeByGenerate.

---
 metagpt/roles/ml_engineer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index 5120a9011..f5bb559e1 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -128,7 +128,7 @@ class MLEngineer(Role):
             if not self.use_tools or self.plan.current_task.task_type == "":
                 # code = "print('abc')"
                 code = await WriteCodeByGenerate().run(
-                    context=context, plan=self.plan, task_guide=task_guide
+                    context=context, plan=self.plan, task_guide=task_guide, temperature=0.0
                 )
                 cause_by = WriteCodeByGenerate
             else:

From c2dba151fbe139291d8fd185aea87e15a04a093a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Thu, 30 Nov 2023 17:42:55 +0800
Subject: [PATCH 08/19] add unit test : write_code_reuse_code_long_for_wine.

---
 .../actions/test_write_analysis_code.py       | 274 +++++++++++-------
 1 file changed, 173 insertions(+), 101 deletions(-)

diff --git a/tests/metagpt/actions/test_write_analysis_code.py b/tests/metagpt/actions/test_write_analysis_code.py
index d4bccb552..1a727a9e4 100644
--- a/tests/metagpt/actions/test_write_analysis_code.py
+++ b/tests/metagpt/actions/test_write_analysis_code.py
@@ -6,110 +6,110 @@ from metagpt.actions.execute_code import ExecutePyCode
 from metagpt.schema import Message
 from metagpt.logs import logger
 
-@pytest.mark.asyncio
-async def test_write_code_by_list_plan():
-    write_code = WriteCodeByGenerate()
-    execute_code = ExecutePyCode()
-    messages = []
-    plan = ["随机生成一个pandas DataFrame时间序列", "绘制这个时间序列的直方图", "求均值"]
-    for task in plan:
-        print(f"\n任务: {task}\n\n")
-        messages.append(Message(task, role='assistant'))
-        code = await write_code.run(messages)
-        messages.append(Message(code, role='assistant'))
-        assert len(code) > 0
-        output = await execute_code.run(code)
-        print(f"\n[Output]: 任务{task}的执行结果是: \n{output}\n")
-        messages.append(output[0])
+# @pytest.mark.asyncio
+# async def test_write_code_by_list_plan():
+#     write_code = WriteCodeByGenerate()
+#     execute_code = ExecutePyCode()
+#     messages = []
+#     plan = ["随机生成一个pandas DataFrame时间序列", "绘制这个时间序列的直方图", "求均值"]
+#     for task in plan:
+#         print(f"\n任务: {task}\n\n")
+#         messages.append(Message(task, role='assistant'))
+#         code = await write_code.run(messages)
+#         messages.append(Message(code, role='assistant'))
+#         assert len(code) > 0
+#         output = await execute_code.run(code)
+#         print(f"\n[Output]: 任务{task}的执行结果是: \n{output}\n")
+#         messages.append(output[0])
 
-@pytest.mark.asyncio
-async def test_write_code_to_correct_error():
+# @pytest.mark.asyncio
+# async def test_write_code_to_correct_error():
 
-    structural_context = """
-    ## User Requirement
-    read a dataset test.csv and print its head
-    ## Current Plan
-    [
-        {
-            "task_id": "1",
-            "dependent_task_ids": [],
-            "instruction": "import pandas and load the dataset from 'test.csv'.",
-            "task_type": "",
-            "code": "",
-            "result": "",
-            "is_finished": false
-        },
-        {
-            "task_id": "2",
-            "dependent_task_ids": [
-                "1"
-            ],
-            "instruction": "Print the head of the dataset to display the first few rows.",
-            "task_type": "",
-            "code": "",
-            "result": "",
-            "is_finished": false
-        }
-    ]
-    ## Current Task
-    {"task_id": "1", "dependent_task_ids": [], "instruction": "import pandas and load the dataset from 'test.csv'.", "task_type": "", "code": "", "result": "", "is_finished": false}
-    """
-    wrong_code = """import pandas as pd\ndata = pd.read_excel('test.csv')\ndata"""  # use read_excel to read a csv
-    error = """
-    Traceback (most recent call last):
-        File "<stdin>", line 2, in <module>
-        File "/Users/gary/miniconda3/envs/py39_scratch/lib/python3.9/site-packages/pandas/io/excel/_base.py", line 478, in read_excel
-            io = ExcelFile(io, storage_options=storage_options, engine=engine)
-        File "/Users/gary/miniconda3/envs/py39_scratch/lib/python3.9/site-packages/pandas/io/excel/_base.py", line 1500, in __init__
-            raise ValueError(
-        ValueError: Excel file format cannot be determined, you must specify an engine manually.
-    """
-    context = [
-        Message(content=structural_context, role="user"),
-        Message(content=wrong_code, role="assistant"),
-        Message(content=error, role="user"),
-    ]
-    new_code = await WriteCodeByGenerate().run(context=context)
-    print(new_code)
-    assert "read_csv" in new_code # should correct read_excel to read_csv
+#     structural_context = """
+#     ## User Requirement
+#     read a dataset test.csv and print its head
+#     ## Current Plan
+#     [
+#         {
+#             "task_id": "1",
+#             "dependent_task_ids": [],
+#             "instruction": "import pandas and load the dataset from 'test.csv'.",
+#             "task_type": "",
+#             "code": "",
+#             "result": "",
+#             "is_finished": false
+#         },
+#         {
+#             "task_id": "2",
+#             "dependent_task_ids": [
+#                 "1"
+#             ],
+#             "instruction": "Print the head of the dataset to display the first few rows.",
+#             "task_type": "",
+#             "code": "",
+#             "result": "",
+#             "is_finished": false
+#         }
+#     ]
+#     ## Current Task
+#     {"task_id": "1", "dependent_task_ids": [], "instruction": "import pandas and load the dataset from 'test.csv'.", "task_type": "", "code": "", "result": "", "is_finished": false}
+#     """
+#     wrong_code = """import pandas as pd\ndata = pd.read_excel('test.csv')\ndata"""  # use read_excel to read a csv
+#     error = """
+#     Traceback (most recent call last):
+#         File "<stdin>", line 2, in <module>
+#         File "/Users/gary/miniconda3/envs/py39_scratch/lib/python3.9/site-packages/pandas/io/excel/_base.py", line 478, in read_excel
+#             io = ExcelFile(io, storage_options=storage_options, engine=engine)
+#         File "/Users/gary/miniconda3/envs/py39_scratch/lib/python3.9/site-packages/pandas/io/excel/_base.py", line 1500, in __init__
+#             raise ValueError(
+#         ValueError: Excel file format cannot be determined, you must specify an engine manually.
+#     """
+#     context = [
+#         Message(content=structural_context, role="user"),
+#         Message(content=wrong_code, role="assistant"),
+#         Message(content=error, role="user"),
+#     ]
+#     new_code = await WriteCodeByGenerate().run(context=context)
+#     print(new_code)
+#     assert "read_csv" in new_code # should correct read_excel to read_csv
 
-@pytest.mark.asyncio
-async def test_write_code_reuse_code_simple():
-    structural_context = """
-    ## User Requirement
-    read a dataset test.csv and print its head
-    ## Current Plan
-    [
-        {
-            "task_id": "1",
-            "dependent_task_ids": [],
-            "instruction": "import pandas and load the dataset from 'test.csv'.",
-            "task_type": "",
-            "code": "import pandas as pd\ndata = pd.read_csv('test.csv')",
-            "result": "",
-            "is_finished": true
-        },
-        {
-            "task_id": "2",
-            "dependent_task_ids": [
-                "1"
-            ],
-            "instruction": "Print the head of the dataset to display the first few rows.",
-            "task_type": "",
-            "code": "",
-            "result": "",
-            "is_finished": false
-        }
-    ]
-    ## Current Task
-    {"task_id": "2", "dependent_task_ids": ["1"], "instruction": "Print the head of the dataset to display the first few rows.", "task_type": "", "code": "", "result": "", "is_finished": false}
-    """
-    context = [
-        Message(content=structural_context, role="user"),
-    ]
-    code = await WriteCodeByGenerate().run(context=context)
-    print(code)
-    assert "pandas" not in code and "read_csv" not in code # should reuse import and read statement from previous one
+# @pytest.mark.asyncio
+# async def test_write_code_reuse_code_simple():
+#     structural_context = """
+#     ## User Requirement
+#     read a dataset test.csv and print its head
+#     ## Current Plan
+#     [
+#         {
+#             "task_id": "1",
+#             "dependent_task_ids": [],
+#             "instruction": "import pandas and load the dataset from 'test.csv'.",
+#             "task_type": "",
+#             "code": "import pandas as pd\ndata = pd.read_csv('test.csv')",
+#             "result": "",
+#             "is_finished": true
+#         },
+#         {
+#             "task_id": "2",
+#             "dependent_task_ids": [
+#                 "1"
+#             ],
+#             "instruction": "Print the head of the dataset to display the first few rows.",
+#             "task_type": "",
+#             "code": "",
+#             "result": "",
+#             "is_finished": false
+#         }
+#     ]
+#     ## Current Task
+#     {"task_id": "2", "dependent_task_ids": ["1"], "instruction": "Print the head of the dataset to display the first few rows.", "task_type": "", "code": "", "result": "", "is_finished": false}
+#     """
+#     context = [
+#         Message(content=structural_context, role="user"),
+#     ]
+#     code = await WriteCodeByGenerate().run(context=context)
+#     print(code)
+#     assert "pandas" not in code and "read_csv" not in code # should reuse import and read statement from previous one
 
 @pytest.mark.asyncio
 async def test_write_code_reuse_code_long():
@@ -167,3 +167,75 @@ async def test_write_code_reuse_code_long():
     success_rate = sum(success) / trials_num
     logger.info(f"success rate: {success_rate :.2f}")
     assert success_rate >= 0.8
+
+
+@pytest.mark.asyncio
+async def test_write_code_reuse_code_long_for_wine():
+    """test code reuse for long context"""
+
+    structural_context = """
+    ## User Requirement
+    Run data analysis on sklearn Wisconsin Breast Cancer dataset, include a plot, train a model to predict targets (20% as validation), and show validation accuracy
+    ## Current Plan
+    [
+        {
+            "task_id": "1",
+            "dependent_task_ids": [],
+            "instruction": "Load the sklearn Wine recognition dataset and perform exploratory data analysis."
+            "task_type": "",
+            "code": "from sklearn.datasets import load_wine\n# Load the Wine recognition dataset\nwine_data = load_wine()\n# Perform exploratory data analysis\nwine_data.keys()",
+            "result": "Truncated to show only the last 1000 characters\ndict_keys(['data', 'target', 'frame', 'target_names', 'DESCR', 'feature_names'])",
+            "is_finished": true
+        },
+        {
+            "task_id": "2",
+            "dependent_task_ids": ["1"],
+            "instruction": "Create a plot to visualize some aspect of the wine dataset."
+            "task_type": "",
+            "code": "",
+            "result": "",
+            "is_finished": false
+        },
+        {
+            "task_id": "3",
+            "dependent_task_ids": ["1"],
+            "instruction": "Split the dataset into training and validation sets with a 20% validation size.",
+            "task_type": "",
+            "code": "",
+            "result": "",
+            "is_finished": false
+        },
+        {
+            "task_id": "4",
+            "dependent_task_ids": ["3"],
+            "instruction": "Train a model on the training set to predict wine class.",
+            "task_type": "",
+            "code": "",
+            "result": "",
+            "is_finished": false
+        },
+        {
+            "task_id": "5",
+            "dependent_task_ids": ["4"],
+            "instruction": "Evaluate the model on the validation set and report the accuracy.",
+            "task_type": "",
+            "code": "",
+            "result": "",
+            "is_finished": false
+        }
+    ]
+    ## Current Task
+    {"task_id": "2", "dependent_task_ids": ["1"], "instruction": "Create a plot to visualize some aspect of the Wine dataset.", "task_type": "", "code": "", "result": "", "is_finished": false}
+    """
+    context = [
+        Message(content=structural_context, role="user"),
+    ]
+    trials_num = 5
+    trials = [WriteCodeByGenerate().run(context=context, temperature=0.0) for _ in range(trials_num)]
+    trial_results = await asyncio.gather(*trials)
+    print(*trial_results, sep="\n\n***\n\n")
+    success = ["load_wine" not in result\
+        for result in trial_results]  # should reuse iris_data from previous tasks
+    success_rate = sum(success) / trials_num
+    logger.info(f"success rate: {success_rate :.2f}")
+    assert success_rate >= 0.8

From 25c536f3e10f9f08584c07b23ceca16dab85dc0a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Thu, 30 Nov 2023 17:44:22 +0800
Subject: [PATCH 09/19] fix: reuse variables in code.

---
 metagpt/actions/write_analysis_code.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/metagpt/actions/write_analysis_code.py b/metagpt/actions/write_analysis_code.py
index ee4555ee1..2b56d6fc1 100644
--- a/metagpt/actions/write_analysis_code.py
+++ b/metagpt/actions/write_analysis_code.py
@@ -40,7 +40,7 @@ class BaseWriteAnalysisCode(Action):
 
 class WriteCodeByGenerate(BaseWriteAnalysisCode):
     """Write code fully by generation"""
-    DEFAULT_SYSTEM_MSG = """You are Code Interpreter, a world-class programmer that can complete any goal by executing code. Strictly follow the plan and generate code step by step. Each step of the code will be executed on the user's machine, and the user will provide the code execution results to you.**Notice: The code for the next step depends on the code for the previous step. Reuse variables in other code directly. Use !pip install in a standalone block to install missing packages.**""" # prompt reference: https://github.com/KillianLucas/open-interpreter/blob/v0.1.4/interpreter/system_message.txt
+    DEFAULT_SYSTEM_MSG = """You are Code Interpreter, a world-class programmer that can complete any goal by executing code. Strictly follow the plan and generate code step by step. Each step of the code will be executed on the user's machine, and the user will provide the code execution results to you.**Notice: The code for the next step depends on the code for the previous step. Must reuse variables in the lastest other code directly, dont creat it again, it is very import for you. Use !pip install in a standalone block to install missing packages.**""" # prompt reference: https://github.com/KillianLucas/open-interpreter/blob/v0.1.4/interpreter/system_message.txt
     # REUSE_CODE_INSTRUCTION = """ATTENTION: DONT include codes from previous tasks in your current code block, include new codes only, DONT repeat codes!"""
 
     def __init__(self, name: str = "", context=None, llm=None) -> str:

From 87acf9b4535f6269a1869d0e054e7c713c04b82d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Thu, 30 Nov 2023 17:48:24 +0800
Subject: [PATCH 10/19] chore

---
 tests/metagpt/actions/test_write_analysis_code.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/metagpt/actions/test_write_analysis_code.py b/tests/metagpt/actions/test_write_analysis_code.py
index 1a727a9e4..e0c3c5230 100644
--- a/tests/metagpt/actions/test_write_analysis_code.py
+++ b/tests/metagpt/actions/test_write_analysis_code.py
@@ -234,7 +234,7 @@ async def test_write_code_reuse_code_long_for_wine():
     trials = [WriteCodeByGenerate().run(context=context, temperature=0.0) for _ in range(trials_num)]
     trial_results = await asyncio.gather(*trials)
     print(*trial_results, sep="\n\n***\n\n")
-    success = ["load_wine" not in result\
+    success = ["load_wine" not in result and "wine_data" in result\
         for result in trial_results]  # should reuse iris_data from previous tasks
     success_rate = sum(success) / trials_num
     logger.info(f"success rate: {success_rate :.2f}")

From 897d1bf0d0c737d465679a38352659485d80e570 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Thu, 30 Nov 2023 17:49:38 +0800
Subject: [PATCH 11/19] chore

---
 .../actions/test_write_analysis_code.py       | 202 +++++++++---------
 1 file changed, 101 insertions(+), 101 deletions(-)

diff --git a/tests/metagpt/actions/test_write_analysis_code.py b/tests/metagpt/actions/test_write_analysis_code.py
index e0c3c5230..211c6ba13 100644
--- a/tests/metagpt/actions/test_write_analysis_code.py
+++ b/tests/metagpt/actions/test_write_analysis_code.py
@@ -6,110 +6,110 @@ from metagpt.actions.execute_code import ExecutePyCode
 from metagpt.schema import Message
 from metagpt.logs import logger
 
-# @pytest.mark.asyncio
-# async def test_write_code_by_list_plan():
-#     write_code = WriteCodeByGenerate()
-#     execute_code = ExecutePyCode()
-#     messages = []
-#     plan = ["随机生成一个pandas DataFrame时间序列", "绘制这个时间序列的直方图", "求均值"]
-#     for task in plan:
-#         print(f"\n任务: {task}\n\n")
-#         messages.append(Message(task, role='assistant'))
-#         code = await write_code.run(messages)
-#         messages.append(Message(code, role='assistant'))
-#         assert len(code) > 0
-#         output = await execute_code.run(code)
-#         print(f"\n[Output]: 任务{task}的执行结果是: \n{output}\n")
-#         messages.append(output[0])
+@pytest.mark.asyncio
+async def test_write_code_by_list_plan():
+    write_code = WriteCodeByGenerate()
+    execute_code = ExecutePyCode()
+    messages = []
+    plan = ["随机生成一个pandas DataFrame时间序列", "绘制这个时间序列的直方图", "求均值"]
+    for task in plan:
+        print(f"\n任务: {task}\n\n")
+        messages.append(Message(task, role='assistant'))
+        code = await write_code.run(messages)
+        messages.append(Message(code, role='assistant'))
+        assert len(code) > 0
+        output = await execute_code.run(code)
+        print(f"\n[Output]: 任务{task}的执行结果是: \n{output}\n")
+        messages.append(output[0])
 
-# @pytest.mark.asyncio
-# async def test_write_code_to_correct_error():
+@pytest.mark.asyncio
+async def test_write_code_to_correct_error():
 
-#     structural_context = """
-#     ## User Requirement
-#     read a dataset test.csv and print its head
-#     ## Current Plan
-#     [
-#         {
-#             "task_id": "1",
-#             "dependent_task_ids": [],
-#             "instruction": "import pandas and load the dataset from 'test.csv'.",
-#             "task_type": "",
-#             "code": "",
-#             "result": "",
-#             "is_finished": false
-#         },
-#         {
-#             "task_id": "2",
-#             "dependent_task_ids": [
-#                 "1"
-#             ],
-#             "instruction": "Print the head of the dataset to display the first few rows.",
-#             "task_type": "",
-#             "code": "",
-#             "result": "",
-#             "is_finished": false
-#         }
-#     ]
-#     ## Current Task
-#     {"task_id": "1", "dependent_task_ids": [], "instruction": "import pandas and load the dataset from 'test.csv'.", "task_type": "", "code": "", "result": "", "is_finished": false}
-#     """
-#     wrong_code = """import pandas as pd\ndata = pd.read_excel('test.csv')\ndata"""  # use read_excel to read a csv
-#     error = """
-#     Traceback (most recent call last):
-#         File "<stdin>", line 2, in <module>
-#         File "/Users/gary/miniconda3/envs/py39_scratch/lib/python3.9/site-packages/pandas/io/excel/_base.py", line 478, in read_excel
-#             io = ExcelFile(io, storage_options=storage_options, engine=engine)
-#         File "/Users/gary/miniconda3/envs/py39_scratch/lib/python3.9/site-packages/pandas/io/excel/_base.py", line 1500, in __init__
-#             raise ValueError(
-#         ValueError: Excel file format cannot be determined, you must specify an engine manually.
-#     """
-#     context = [
-#         Message(content=structural_context, role="user"),
-#         Message(content=wrong_code, role="assistant"),
-#         Message(content=error, role="user"),
-#     ]
-#     new_code = await WriteCodeByGenerate().run(context=context)
-#     print(new_code)
-#     assert "read_csv" in new_code # should correct read_excel to read_csv
+    structural_context = """
+    ## User Requirement
+    read a dataset test.csv and print its head
+    ## Current Plan
+    [
+        {
+            "task_id": "1",
+            "dependent_task_ids": [],
+            "instruction": "import pandas and load the dataset from 'test.csv'.",
+            "task_type": "",
+            "code": "",
+            "result": "",
+            "is_finished": false
+        },
+        {
+            "task_id": "2",
+            "dependent_task_ids": [
+                "1"
+            ],
+            "instruction": "Print the head of the dataset to display the first few rows.",
+            "task_type": "",
+            "code": "",
+            "result": "",
+            "is_finished": false
+        }
+    ]
+    ## Current Task
+    {"task_id": "1", "dependent_task_ids": [], "instruction": "import pandas and load the dataset from 'test.csv'.", "task_type": "", "code": "", "result": "", "is_finished": false}
+    """
+    wrong_code = """import pandas as pd\ndata = pd.read_excel('test.csv')\ndata"""  # use read_excel to read a csv
+    error = """
+    Traceback (most recent call last):
+        File "<stdin>", line 2, in <module>
+        File "/Users/gary/miniconda3/envs/py39_scratch/lib/python3.9/site-packages/pandas/io/excel/_base.py", line 478, in read_excel
+            io = ExcelFile(io, storage_options=storage_options, engine=engine)
+        File "/Users/gary/miniconda3/envs/py39_scratch/lib/python3.9/site-packages/pandas/io/excel/_base.py", line 1500, in __init__
+            raise ValueError(
+        ValueError: Excel file format cannot be determined, you must specify an engine manually.
+    """
+    context = [
+        Message(content=structural_context, role="user"),
+        Message(content=wrong_code, role="assistant"),
+        Message(content=error, role="user"),
+    ]
+    new_code = await WriteCodeByGenerate().run(context=context)
+    print(new_code)
+    assert "read_csv" in new_code # should correct read_excel to read_csv
 
-# @pytest.mark.asyncio
-# async def test_write_code_reuse_code_simple():
-#     structural_context = """
-#     ## User Requirement
-#     read a dataset test.csv and print its head
-#     ## Current Plan
-#     [
-#         {
-#             "task_id": "1",
-#             "dependent_task_ids": [],
-#             "instruction": "import pandas and load the dataset from 'test.csv'.",
-#             "task_type": "",
-#             "code": "import pandas as pd\ndata = pd.read_csv('test.csv')",
-#             "result": "",
-#             "is_finished": true
-#         },
-#         {
-#             "task_id": "2",
-#             "dependent_task_ids": [
-#                 "1"
-#             ],
-#             "instruction": "Print the head of the dataset to display the first few rows.",
-#             "task_type": "",
-#             "code": "",
-#             "result": "",
-#             "is_finished": false
-#         }
-#     ]
-#     ## Current Task
-#     {"task_id": "2", "dependent_task_ids": ["1"], "instruction": "Print the head of the dataset to display the first few rows.", "task_type": "", "code": "", "result": "", "is_finished": false}
-#     """
-#     context = [
-#         Message(content=structural_context, role="user"),
-#     ]
-#     code = await WriteCodeByGenerate().run(context=context)
-#     print(code)
-#     assert "pandas" not in code and "read_csv" not in code # should reuse import and read statement from previous one
+@pytest.mark.asyncio
+async def test_write_code_reuse_code_simple():
+    structural_context = """
+    ## User Requirement
+    read a dataset test.csv and print its head
+    ## Current Plan
+    [
+        {
+            "task_id": "1",
+            "dependent_task_ids": [],
+            "instruction": "import pandas and load the dataset from 'test.csv'.",
+            "task_type": "",
+            "code": "import pandas as pd\ndata = pd.read_csv('test.csv')",
+            "result": "",
+            "is_finished": true
+        },
+        {
+            "task_id": "2",
+            "dependent_task_ids": [
+                "1"
+            ],
+            "instruction": "Print the head of the dataset to display the first few rows.",
+            "task_type": "",
+            "code": "",
+            "result": "",
+            "is_finished": false
+        }
+    ]
+    ## Current Task
+    {"task_id": "2", "dependent_task_ids": ["1"], "instruction": "Print the head of the dataset to display the first few rows.", "task_type": "", "code": "", "result": "", "is_finished": false}
+    """
+    context = [
+        Message(content=structural_context, role="user"),
+    ]
+    code = await WriteCodeByGenerate().run(context=context)
+    print(code)
+    assert "pandas" not in code and "read_csv" not in code # should reuse import and read statement from previous one
 
 @pytest.mark.asyncio
 async def test_write_code_reuse_code_long():

From f440ff69d04768da1f8183cb4386d36bd9886456 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Thu, 30 Nov 2023 18:04:55 +0800
Subject: [PATCH 12/19] chore

---
 tests/metagpt/actions/test_execute_code.py | 82 +++++++++++-----------
 1 file changed, 41 insertions(+), 41 deletions(-)

diff --git a/tests/metagpt/actions/test_execute_code.py b/tests/metagpt/actions/test_execute_code.py
index 8894f2cb9..73b5886dc 100644
--- a/tests/metagpt/actions/test_execute_code.py
+++ b/tests/metagpt/actions/test_execute_code.py
@@ -4,57 +4,57 @@ from metagpt.actions.execute_code import ExecutePyCode
 from metagpt.schema import Message
 
 
-# @pytest.mark.asyncio
-# async def test_code_running():
-#     pi = ExecutePyCode()
-#     output = await pi.run("print('hello world!')")
-#     assert output[1] is True
-#     output = await pi.run({"code": "print('hello world!')", "language": "python"})
-#     assert output[1] is True
-#     code_msg = Message("print('hello world!')")
-#     output = await pi.run(code_msg)
-#     assert output[1] is True
+@pytest.mark.asyncio
+async def test_code_running():
+    pi = ExecutePyCode()
+    output = await pi.run("print('hello world!')")
+    assert output[1] is True
+    output = await pi.run({"code": "print('hello world!')", "language": "python"})
+    assert output[1] is True
+    code_msg = Message("print('hello world!')")
+    output = await pi.run(code_msg)
+    assert output[1] is True
 
 
-# @pytest.mark.asyncio
-# async def test_split_code_running():
-#     pi = ExecutePyCode()
-#     output = await pi.run("x=1\ny=2")
-#     output = await pi.run("z=x+y")
-#     output = await pi.run("assert z==3")
-#     assert output[1] is True
+@pytest.mark.asyncio
+async def test_split_code_running():
+    pi = ExecutePyCode()
+    output = await pi.run("x=1\ny=2")
+    output = await pi.run("z=x+y")
+    output = await pi.run("assert z==3")
+    assert output[1] is True
 
 
-# @pytest.mark.asyncio
-# async def test_execute_error():
-#     pi = ExecutePyCode()
-#     output = await pi.run("z=1/0")
-#     assert output[1] is False
+@pytest.mark.asyncio
+async def test_execute_error():
+    pi = ExecutePyCode()
+    output = await pi.run("z=1/0")
+    assert output[1] is False
 
 
-# @pytest.mark.asyncio
-# async def test_plotting_code():
-#     pi = ExecutePyCode()
-#     code = """
-#     import numpy as np
-#     import matplotlib.pyplot as plt
+@pytest.mark.asyncio
+async def test_plotting_code():
+    pi = ExecutePyCode()
+    code = """
+    import numpy as np
+    import matplotlib.pyplot as plt
 
-#     # 生成随机数据
-#     random_data = np.random.randn(1000)  # 生成1000个符合标准正态分布的随机数
+    # 生成随机数据
+    random_data = np.random.randn(1000)  # 生成1000个符合标准正态分布的随机数
 
-#     # 绘制直方图
-#     plt.hist(random_data, bins=30, density=True, alpha=0.7, color='blue', edgecolor='black')
+    # 绘制直方图
+    plt.hist(random_data, bins=30, density=True, alpha=0.7, color='blue', edgecolor='black')
 
-#     # 添加标题和标签
-#     plt.title('Histogram of Random Data')
-#     plt.xlabel('Value')
-#     plt.ylabel('Frequency')
+    # 添加标题和标签
+    plt.title('Histogram of Random Data')
+    plt.xlabel('Value')
+    plt.ylabel('Frequency')
 
-#     # 显示图形
-#     plt.show()
-#     """
-#     output = await pi.run(code)
-#     assert output[1] is True
+    # 显示图形
+    plt.show()
+    """
+    output = await pi.run(code)
+    assert output[1] is True
 
 
 @pytest.mark.asyncio

From aad201e06f288778ddd1fea20640a761d8afc62e Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Fri, 1 Dec 2023 11:57:58 +0800
Subject: [PATCH 13/19] assign task_type for task

---
 metagpt/actions/write_plan.py  | 33 +++++++++++++++++++++++-----
 metagpt/prompts/ml_engineer.py | 40 +++++++++++++++++++++++++++++-----
 2 files changed, 63 insertions(+), 10 deletions(-)

diff --git a/metagpt/actions/write_plan.py b/metagpt/actions/write_plan.py
index dcfa25d55..5e42de199 100644
--- a/metagpt/actions/write_plan.py
+++ b/metagpt/actions/write_plan.py
@@ -4,12 +4,14 @@
 @Author  :   orange-crow
 @File    :   plan.py
 """
-from typing import List
+from typing import List, Dict
 import json
 
 from metagpt.actions import Action
+from metagpt.prompts.ml_engineer import ASSIGN_TASK_TYPE_PROMPT, ASSIGN_TASK_TYPE
 from metagpt.schema import Message, Task
-from metagpt.utils.common import CodeParser
+from metagpt.utils.common import CodeParser, create_func_config
+
 
 class WritePlan(Action):
     PROMPT_TEMPLATE = """
@@ -30,7 +32,28 @@ class WritePlan(Action):
     ]
     ```
     """
-    async def run(self, context: List[Message], max_tasks: int = 5) -> str:
+
+    async def assign_task_type(self, tasks: List[Dict]) -> List[Dict]:
+        """Assign task type to each task in tasks
+
+        Args:
+            tasks (List[Dict]): tasks to be assigned task type
+
+        Returns:
+            List[Dict]: tasks with task type assigned
+        """
+        task_list = "\n".join(
+            [f"Task {task['task_id']}: {task['instruction']}" for task in tasks]
+        )
+        prompt = ASSIGN_TASK_TYPE_PROMPT.format(task_list=task_list)
+        tool_config = create_func_config(ASSIGN_TASK_TYPE)
+        rsp = await self.llm.aask_code(prompt, **tool_config)
+        task_type_list = rsp["task_type"]
+        for task, task_type in zip(tasks, task_type_list):
+            task["task_type"] = task_type
+        return tasks
+
+    async def run(self, context: List[Message], max_tasks: int = 5) -> List[Dict]:
         prompt = (
             self.PROMPT_TEMPLATE.replace("__context__", "\n".join([str(ct) for ct in context]))
             # .replace("__current_plan__", current_plan)
@@ -38,10 +61,10 @@ class WritePlan(Action):
         )
         rsp = await self._aask(prompt)
         rsp = CodeParser.parse_code(block=None, text=rsp)
+        rsp = await self.assign_task_type(json.loads(rsp))
         return rsp
 
     @staticmethod
-    def rsp_to_tasks(rsp: str) -> List[Task]:
-        rsp = json.loads(rsp)
+    def rsp_to_tasks(rsp: List[Dict]) -> List[Task]:
         tasks = [Task(**task_config) for task_config in rsp]
         return tasks
diff --git a/metagpt/prompts/ml_engineer.py b/metagpt/prompts/ml_engineer.py
index 70a40ef34..0c4d036fc 100644
--- a/metagpt/prompts/ml_engineer.py
+++ b/metagpt/prompts/ml_engineer.py
@@ -4,6 +4,35 @@
 # @Author  : lidanyang
 # @File    : ml_engineer
 # @Desc    :
+ASSIGN_TASK_TYPE_PROMPT = """
+## All Task Type:
+- **data_preprocess**: Only involve cleaning and preparing data through techniques like imputation, scaling, and encoding, not containing reading data, feature engineering, model training, etc.
+- **feature_engineering**: Involves enhancing data features through techniques like encoding, aggregation, time component analysis, and creating polynomial and interaction features, etc.
+- **other**: Any tasks that do not fit into the previous categories, such as visualization, summarizing findings, build model, etc.
+
+Please assign a task type to each task in the list below from the given categories:
+{task_list}
+"""
+
+ASSIGN_TASK_TYPE = {
+    "name": "assign_task_type",
+    "description": "assign task type to each task by order",
+    "parameters": {
+        "type": "object",
+        "properties": {
+            "task_type": {
+                "type": "array",
+                "description": "List of task type.",
+                "items": {
+                    "type": "string",
+                },
+            },
+        },
+        "required": ["task_type"],
+    },
+}
+
+
 TOOL_RECOMMENDATION_PROMPT = """
 ## Comprehensive Task Description:
 {task}
@@ -137,11 +166,12 @@ When performing feature engineering, please adhere to the following principles:
 - Importantly, provide detailed comments explaining the purpose of each feature and how it might enhance model performance, especially when the features are generated based on semantic understanding without clear user directives.
 """
 
-CLASSIFICATION_MODEL_PROMPT = """
+MODEL_TRAIN_PROMPT = """
+When selecting and training a model, please follow these guidelines to ensure optimal performance:
+- Keep in mind that your user prioritizes results and is highly focused on model performance. So, when needed, feel free to use models of any complexity to improve effectiveness, such as lightGBM, XGBoost, CatBoost, etc.
+— If user specifies a model, use that model. Otherwise, use the model you believe will best solve the problem.
 """
 
-REGRESSION_MODEL_PROMPT = """
-"""
 
 DATA_PREPROCESS_OUTPUT_DESC = "Please note that all functions uniformly output a processed pandas.DataFrame, facilitating seamless integration into the broader workflow."
 
@@ -155,8 +185,8 @@ REGRESSION_MODEL_OUTPUT_DESC = ""
 ML_SPECIFIC_PROMPT = {
     "data_preprocess": DATA_PREPROCESS_PROMPT,
     "feature_engineering": FEATURE_ENGINEERING_PROMPT,
-    "classification_model": CLASSIFICATION_MODEL_PROMPT,
-    "regression_model": REGRESSION_MODEL_PROMPT,
+    "classification_model": MODEL_TRAIN_PROMPT,
+    "regression_model": MODEL_TRAIN_PROMPT,
 }
 
 TOOL_OUTPUT_DESC = {

From 35e8a501c54762bd95bddebc1b3c8367a8993238 Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Fri, 1 Dec 2023 11:59:28 +0800
Subject: [PATCH 14/19] add log print

---
 metagpt/actions/write_analysis_code.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/metagpt/actions/write_analysis_code.py b/metagpt/actions/write_analysis_code.py
index 6fff1c66f..e81228109 100644
--- a/metagpt/actions/write_analysis_code.py
+++ b/metagpt/actions/write_analysis_code.py
@@ -4,10 +4,10 @@
 @Author  :   orange-crow
 @File    :   write_code_v2.py
 """
-import json
 from typing import Dict, List, Union, Tuple
 
 from metagpt.actions import Action
+from metagpt.logs import logger
 from metagpt.prompts.ml_engineer import (
     TOOL_RECOMMENDATION_PROMPT,
     SELECT_FUNCTION_TOOLS,
@@ -174,6 +174,7 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
             task, task_guide, available_tools
         )
         recommend_tools, tool_catalog = self._parse_recommend_tools(task_type, recommend_tools)
+        logger.info(f"Recommended tools for every steps: {recommend_tools}")
 
         special_prompt = ML_SPECIFIC_PROMPT.get(task_type, "")
         module_name = ML_MODULE_MAP[task_type]

From cb8a8ffd5cbf4e13c25bab7ea51ec6736a6a9bcc Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Fri, 1 Dec 2023 13:44:24 +0800
Subject: [PATCH 15/19] fix rsp type

---
 metagpt/actions/write_plan.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/metagpt/actions/write_plan.py b/metagpt/actions/write_plan.py
index 5e42de199..f7c096f2c 100644
--- a/metagpt/actions/write_plan.py
+++ b/metagpt/actions/write_plan.py
@@ -33,7 +33,7 @@ class WritePlan(Action):
     ```
     """
 
-    async def assign_task_type(self, tasks: List[Dict]) -> List[Dict]:
+    async def assign_task_type(self, tasks: List[Dict]) -> str:
         """Assign task type to each task in tasks
 
         Args:
@@ -51,9 +51,9 @@ class WritePlan(Action):
         task_type_list = rsp["task_type"]
         for task, task_type in zip(tasks, task_type_list):
             task["task_type"] = task_type
-        return tasks
+        return json.dumps(tasks)
 
-    async def run(self, context: List[Message], max_tasks: int = 5) -> List[Dict]:
+    async def run(self, context: List[Message], max_tasks: int = 5) -> str:
         prompt = (
             self.PROMPT_TEMPLATE.replace("__context__", "\n".join([str(ct) for ct in context]))
             # .replace("__current_plan__", current_plan)
@@ -65,6 +65,7 @@ class WritePlan(Action):
         return rsp
 
     @staticmethod
-    def rsp_to_tasks(rsp: List[Dict]) -> List[Task]:
+    def rsp_to_tasks(rsp: str) -> List[Task]:
+        rsp = json.loads(rsp)
         tasks = [Task(**task_config) for task_config in rsp]
         return tasks

From e4a17d122c9c115530375c4f095f5a6be46ec03a Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Fri, 1 Dec 2023 14:15:36 +0800
Subject: [PATCH 16/19] fill other task_type

---
 metagpt/roles/ml_engineer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index 08c5649d4..0ea73a045 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -123,7 +123,7 @@ class MLEngineer(Role):
             # print("*" * 10)
             # breakpoint()
 
-            if not self.use_tools or self.plan.current_task.task_type == "":
+            if not self.use_tools or self.plan.current_task.task_type == "other":
                 # code = "print('abc')"
                 code = await WriteCodeByGenerate().run(
                     context=context, plan=self.plan, task_guide=task_guide

From 2049f6cd01c66e5ed0402a18bebc20b1a9ceda5d Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Fri, 1 Dec 2023 14:29:51 +0800
Subject: [PATCH 17/19] only assign task_type when use_tools

---
 metagpt/actions/write_plan.py | 7 +++++--
 metagpt/roles/ml_engineer.py  | 4 +++-
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/metagpt/actions/write_plan.py b/metagpt/actions/write_plan.py
index f7c096f2c..5145ffd68 100644
--- a/metagpt/actions/write_plan.py
+++ b/metagpt/actions/write_plan.py
@@ -53,7 +53,9 @@ class WritePlan(Action):
             task["task_type"] = task_type
         return json.dumps(tasks)
 
-    async def run(self, context: List[Message], max_tasks: int = 5) -> str:
+    async def run(
+        self, context: List[Message], max_tasks: int = 5, use_tools: bool = False
+    ) -> str:
         prompt = (
             self.PROMPT_TEMPLATE.replace("__context__", "\n".join([str(ct) for ct in context]))
             # .replace("__current_plan__", current_plan)
@@ -61,7 +63,8 @@ class WritePlan(Action):
         )
         rsp = await self._aask(prompt)
         rsp = CodeParser.parse_code(block=None, text=rsp)
-        rsp = await self.assign_task_type(json.loads(rsp))
+        if use_tools:
+            rsp = await self.assign_task_type(json.loads(rsp))
         return rsp
 
     @staticmethod
diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index 0ea73a045..8e02e093b 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -173,7 +173,9 @@ class MLEngineer(Role):
         plan_confirmed = False
         while not plan_confirmed:
             context = self.get_useful_memories()
-            rsp = await WritePlan().run(context, max_tasks=max_tasks)
+            rsp = await WritePlan().run(
+                context, max_tasks=max_tasks, use_tools=self.use_tools
+            )
             self.working_memory.add(
                 Message(content=rsp, role="assistant", cause_by=WritePlan)
             )

From 59af6d96921fadbba22c57ea171ab0725d8e5b0d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Fri, 1 Dec 2023 15:21:40 +0800
Subject: [PATCH 18/19] chore: remove _result.

---
 metagpt/roles/ml_engineer.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index f5bb559e1..ae346579b 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -145,9 +145,8 @@ class MLEngineer(Role):
             # truncated the result
             print(truncate(result))
             # print(result)
-            _result = truncate(remove_escape_and_color_codes(result))
             self.working_memory.add(
-                Message(content=_result, role="user", cause_by=ExecutePyCode)
+                Message(content=truncate(remove_escape_and_color_codes(result)), role="user", cause_by=ExecutePyCode)
             )
 
             if "!pip" in code:

From f1cfbea7728084e14bd93cecbd0b8624c381cbb9 Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Fri, 1 Dec 2023 15:31:38 +0800
Subject: [PATCH 19/19] add test for write code with tools

---
 .../actions/test_write_analysis_code.py       | 74 ++++++++++++++++++-
 1 file changed, 72 insertions(+), 2 deletions(-)

diff --git a/tests/metagpt/actions/test_write_analysis_code.py b/tests/metagpt/actions/test_write_analysis_code.py
index cde5fa7ad..2319331d4 100644
--- a/tests/metagpt/actions/test_write_analysis_code.py
+++ b/tests/metagpt/actions/test_write_analysis_code.py
@@ -1,8 +1,8 @@
 import pytest
 
-from metagpt.actions.write_analysis_code import WriteCodeByGenerate
+from metagpt.actions.write_analysis_code import WriteCodeByGenerate, WriteCodeWithTools
 from metagpt.actions.execute_code import ExecutePyCode
-from metagpt.schema import Message
+from metagpt.schema import Message, Plan, Task
 
 
 # @pytest.mark.asyncio
@@ -37,3 +37,73 @@ async def test_write_code_by_list_plan():
         output = await execute_code.run(code)
         print(f"\n[Output]: 任务{task}的执行结果是: \n{output}\n")
         messages.append(output[0])
+
+
+@pytest.mark.asyncio
+async def test_tool_recommendation():
+    task = "对已经读取的数据集进行数据清洗"
+    code_steps = """
+    step 1: 对数据集进行去重
+    step 2: 对数据集进行缺失值处理
+    """
+    available_tools = [
+        {
+            "name": "fill_missing_value",
+            "description": "Completing missing values with simple strategies",
+        },
+        {
+            "name": "split_bins",
+            "description": "Bin continuous data into intervals and return the bin identifier encoded as an integer value",
+        },
+    ]
+    write_code = WriteCodeWithTools()
+    tools = await write_code._tool_recommendation(task, code_steps, available_tools)
+
+    assert len(tools) == 2
+    assert tools[0] == []
+    assert tools[1] == ["fill_missing_value"]
+
+
+@pytest.mark.asyncio
+async def test_write_code_with_tools():
+    write_code = WriteCodeWithTools()
+    messages = []
+    task_map = {
+        "1": Task(
+                task_id="1",
+                instruction="随机生成一个pandas DataFrame数据集",
+                task_type="unknown",
+                dependent_task_ids=[],
+                code="""
+                import pandas as pd
+                df = pd.DataFrame({
+                    'a': [1, 2, 3, 4, 5],
+                    'b': [1.1, 2.2, 3.3, 4.4, np.nan],
+                    'c': ['aa', 'bb', 'cc', 'dd', 'ee'],
+                    'd': [1, 2, 3, 4, 5]
+                })
+                """,
+                is_finished=True,
+            ),
+        "2": Task(
+                task_id="2",
+                instruction="对数据集进行数据清洗",
+                task_type="data_preprocess",
+                dependent_task_ids=["1"],
+            ),
+    }
+    plan = Plan(
+        goal="构造数据集并进行数据清洗",
+        tasks=list(task_map.values()),
+        task_map=task_map,
+        current_task_id="2",
+    )
+    task_guide = """
+    step 1: 对数据集进行去重
+    step 2: 对数据集进行缺失值处理
+    """
+    data_desc = "None"
+
+    code = await write_code.run(messages, plan, task_guide, data_desc)
+    assert len(code) > 0
+    print(code)