From 962632cd15e76ba142d89ef086467be97f6ba7f0 Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Wed, 6 Dec 2023 14:16:48 +0800
Subject: [PATCH 01/49] add GenerateDataDesc action

---
 metagpt/roles/ml_engineer.py | 131 ++++++++++++++++++++++++++++++-----
 1 file changed, 112 insertions(+), 19 deletions(-)

diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index 65583638e..15edb2b06 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -1,25 +1,38 @@
-from typing import Dict, List, Union
+import glob
 import json
-import subprocess
+from typing import List
 
 import fire
+import pandas as pd
 import re
 
-from metagpt.roles import Role
 from metagpt.actions import Action
-from metagpt.schema import Message, Task, Plan
-from metagpt.logs import logger
-from metagpt.actions.write_plan import WritePlan
-from metagpt.actions.write_analysis_code import WriteCodeByGenerate, WriteCodeWithTools
 from metagpt.actions.execute_code import ExecutePyCode
+from metagpt.actions.write_analysis_code import WriteCodeByGenerate, WriteCodeWithTools
+from metagpt.actions.write_plan import WritePlan
+from metagpt.actions.write_task_guide import WriteTaskGuide
+from metagpt.logs import logger
+from metagpt.prompts.ml_engineer import GEN_DATA_DESC_PROMPT
+from metagpt.roles import Role
+from metagpt.schema import Message, Plan
+from metagpt.utils.common import CodeParser
 
 STRUCTURAL_CONTEXT = """
 ## User Requirement
 {user_requirement}
+## Dataset Description
+{data_desc}
 ## Current Plan
 {tasks}
 ## Current Task
 {current_task}
+## Packages Installed
+scikit-learn
+pandas
+numpy
+lightgbm
+xgboost
+catboost
 """
 
 
@@ -43,6 +56,50 @@ def remove_escape_and_color_codes(input_str):
     return result
 
 
+def read_data(file: str) -> pd.DataFrame:
+    if file.endswith(".csv"):
+        df = pd.read_csv(file, sep=",")
+        sep_list = [";", "\t", ":", " ", "|"]
+        for sep in sep_list:
+            if df.shape[1] == 1:
+                df = pd.read_csv(file, sep=sep)
+            else:
+                break
+    else:
+        raise ValueError(f"Unsupported file type: {file}")
+    return df
+
+
+def get_samples(df: pd.DataFrame) -> str:
+    data = []
+
+    if len(df) > 5:
+        df_ = df.sample(5, random_state=0)
+    else:
+        df_ = df
+
+    for i in list(df_):
+        nan_freq = float("%.2g" % (df[i].isna().mean() * 100))
+        n_unique = df[i].nunique()
+        s = df_[i].tolist()
+
+        if str(df[i].dtype) == "float64":
+            s = [round(sample, 2) if not pd.isna(sample) else None for sample in s]
+
+        data.append([df_[i].name, df[i].dtype, nan_freq, n_unique, s])
+    samples = pd.DataFrame(
+        data,
+        columns=[
+            "Column_name",
+            "Data_type",
+            "NaN_Frequency(%)",
+            "N_unique",
+            "Samples",
+        ],
+    )
+    return samples.to_string(index=False)
+
+
 class AskReview(Action):
     async def run(self, context: List[Message], plan: Plan = None):
         logger.info("Current overall plan:")
@@ -66,24 +123,47 @@ class AskReview(Action):
         return rsp, confirmed
 
 
-class WriteTaskGuide(Action):
-    async def run(self, task_instruction: str, data_desc: str = "") -> str:
-        return ""
+# class WriteTaskGuide(Action):
+#     async def run(self, task_instruction: str, data_desc: dict = None) -> str:
+#         return ""
+
+
+class GenerateDataDesc(Action):
+    async def run(self, files: list) -> dict:
+        data_desc = {}
+        for file in files:
+            df = read_data(file)
+            file_name = file.split("/")[-1]
+            data_head = df.head().to_dict(orient="list")
+            data_head = json.dumps(data_head, indent=4, ensure_ascii=False)
+            prompt = GEN_DATA_DESC_PROMPT.replace("{data_head}", data_head)
+            rsp = await self._aask(prompt)
+            rsp = CodeParser.parse_code(block=None, text=rsp)
+            data_desc[file_name] = {}
+            data_desc[file_name]["path"] = file
+            data_desc[file_name]["description"] = rsp
+            data_desc[file_name]["column_info"] = get_samples(df)
+        return data_desc
 
 
 class MLEngineer(Role):
     def __init__(
-        self, name="ABC", profile="MLEngineer", goal="", auto_run: bool = False
+        self, name="ABC", profile="MLEngineer", goal="", auto_run: bool = False, data_path: str = None
     ):
         super().__init__(name=name, profile=profile, goal=goal)
         self._set_react_mode(react_mode="plan_and_act")
         self.plan = Plan(goal=goal)
-        self.use_tools = False
-        self.use_task_guide = False
+        self.use_tools = True
+        self.use_task_guide = True
         self.execute_code = ExecutePyCode()
         self.auto_run = auto_run
+        self.data_path = data_path
+        self.data_desc = {}
 
     async def _plan_and_act(self):
+        if self.data_path:
+            self.data_desc = await self._generate_data_desc()
+
         # create initial plan and update until confirmation
         await self._update_plan()
 
@@ -108,9 +188,14 @@ class MLEngineer(Role):
                 # update plan according to user's feedback and to take on changed tasks
                 await self._update_plan()
 
+    async def _generate_data_desc(self):
+        files = glob.glob(self.data_path + "/*.csv")
+        data_desc = await GenerateDataDesc().run(files=files)
+        return data_desc
+
     async def _write_and_exec_code(self, max_retry: int = 3):
         task_guide = (
-            await WriteTaskGuide().run(self.plan.current_task.instruction)
+            await WriteTaskGuide().run(self.plan)
             if self.use_task_guide
             else ""
         )
@@ -126,14 +211,16 @@ class MLEngineer(Role):
             # breakpoint()
 
             if not self.use_tools or self.plan.current_task.task_type == "other":
+                logger.info("Write code with pure generation")
                 # code = "print('abc')"
                 code = await WriteCodeByGenerate().run(
                     context=context, plan=self.plan, task_guide=task_guide, temperature=0.0
                 )
                 cause_by = WriteCodeByGenerate
             else:
+                logger.info("Write code with tools")
                 code = await WriteCodeWithTools().run(
-                    context=context, plan=self.plan, task_guide=task_guide, data_desc=""
+                    context=context, plan=self.plan, task_guide=task_guide
                 )
                 cause_by = WriteCodeWithTools
 
@@ -192,7 +279,10 @@ class MLEngineer(Role):
         )
         current_task = self.plan.current_task.json() if self.plan.current_task else {}
         context = STRUCTURAL_CONTEXT.format(
-            user_requirement=user_requirement, tasks=tasks, current_task=current_task
+            user_requirement=user_requirement,
+            data_desc=self.data_desc,
+            tasks=tasks,
+            current_task=current_task
         )
         context_msg = [Message(content=context, role="user")]
 
@@ -204,14 +294,17 @@ class MLEngineer(Role):
 
 
 if __name__ == "__main__":
-    requirement = "Run data analysis on sklearn Iris dataset, include a plot"
+    # requirement = "Run data analysis on sklearn Iris dataset, include a plot.."
     # requirement = "Run data analysis on sklearn Diabetes dataset, include a plot"
     # requirement = "Run data analysis on sklearn Wine recognition dataset, include a plot, and train a model to predict wine class (20% as validation), and show validation accuracy"
     # 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"
     # requirement = "Run EDA and visualization on this dataset, train a model to predict survival, report metrics on validation set (20%), dataset: workspace/titanic/train.csv"
 
-    async def main(requirement: str = requirement, auto_run: bool = False):
-        role = MLEngineer(goal=requirement, auto_run=auto_run)
+    requirement = "Perform data analysis on the provided data. Train a model to predict the target variable Survived. Include data preprocessing, feature engineering, and modeling in your pipeline. The metric is accuracy."
+    data_path = "/data/lidanyang/tabular_data/titanic"
+
+    async def main(requirement: str = requirement, auto_run: bool = True, data_path: str = data_path):
+        role = MLEngineer(goal=requirement, auto_run=auto_run, data_path=data_path)
         await role.run(requirement)
 
     fire.Fire(main)

From 6edbed8fb6e9ea19c2fa37de8d7f74888b83b903 Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Wed, 6 Dec 2023 14:17:29 +0800
Subject: [PATCH 02/49] refine schema

---
 metagpt/tools/functions/schemas/feature_engineering.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/metagpt/tools/functions/schemas/feature_engineering.py b/metagpt/tools/functions/schemas/feature_engineering.py
index c14bb933e..df2eebff6 100644
--- a/metagpt/tools/functions/schemas/feature_engineering.py
+++ b/metagpt/tools/functions/schemas/feature_engineering.py
@@ -20,7 +20,7 @@ class PolynomialExpansion(ToolSchema):
 
 
 class OneHotEncoding(ToolSchema):
-    """Apply one-hot encoding to specified categorical columns in a DataFrame."""
+    """Apply one-hot encoding to specified categorical columns, the original columns will be dropped."""
 
     df: pd.DataFrame = tool_field(description="DataFrame to process.")
     cols: list = tool_field(description="Categorical columns to be one-hot encoded.")

From 0b918eb224e07621525d2518dba8e417de6fab8a Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Wed, 6 Dec 2023 14:18:38 +0800
Subject: [PATCH 03/49] Standardize the process with or without task guide

---
 metagpt/actions/write_analysis_code.py | 147 ++++++++++-------------
 metagpt/prompts/ml_engineer.py         | 159 +++++++++++--------------
 metagpt/tools/functions/__init__.py    |   1 +
 3 files changed, 136 insertions(+), 171 deletions(-)

diff --git a/metagpt/actions/write_analysis_code.py b/metagpt/actions/write_analysis_code.py
index db0df2f90..646b4f3f1 100644
--- a/metagpt/actions/write_analysis_code.py
+++ b/metagpt/actions/write_analysis_code.py
@@ -23,28 +23,8 @@ from metagpt.utils.common import create_func_config
 
 
 class BaseWriteAnalysisCode(Action):
-    async def run(
-        self, context: List[Message], plan: Plan = None, task_guide: str = ""
-    ) -> str:
-        """Run of a code writing action, used in data analysis or modeling
-
-        Args:
-            context (List[Message]): Action output history, source action denoted by Message.cause_by
-            plan (Plan, optional): Overall plan. Defaults to None.
-            task_guide (str, optional): suggested step breakdown for the current task. Defaults to "".
-
-        Returns:
-            str: The code string.
-        """
-
-
-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. 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:
-        super().__init__(name, context, llm)
+    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."""  # 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 process_msg(self, prompt: Union[str, List[Dict], Message, List[Message]], system_msg: str = None):
         default_system_msg = system_msg or self.DEFAULT_SYSTEM_MSG
@@ -81,6 +61,27 @@ class WriteCodeByGenerate(BaseWriteAnalysisCode):
             }
         return messages
 
+    async def run(
+        self, context: List[Message], plan: Plan = None, task_guide: str = ""
+    ) -> str:
+        """Run of a code writing action, used in data analysis or modeling
+
+        Args:
+            context (List[Message]): Action output history, source action denoted by Message.cause_by
+            plan (Plan, optional): Overall plan. Defaults to None.
+            task_guide (str, optional): suggested step breakdown for the current task. Defaults to "".
+
+        Returns:
+            str: The code string.
+        """
+
+
+class WriteCodeByGenerate(BaseWriteAnalysisCode):
+    """Write code fully by generation"""
+
+    def __init__(self, name: str = "", context=None, llm=None) -> str:
+        super().__init__(name, context, llm)
+
     async def run(
         self,
         context: [List[Message]],
@@ -89,7 +90,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"]
@@ -99,7 +100,7 @@ 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) -> Tuple[Dict, List[Dict]]:
+    def _parse_recommend_tools(module: str, recommend_tools: list) -> List[Dict]:
         """
         Parses and validates a list of recommended tools, and retrieves their schema from registry.
 
@@ -108,44 +109,40 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
             recommend_tools (list): A list of lists of recommended tools for each step.
 
         Returns:
-            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.
+            List[Dict]: A list of dicts of valid tool schemas.
         """
-        valid_tools = {}
+        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] = tools
+        for tool in recommend_tools:
+            if tool in available_tools:
+                valid_tools.append(tool)
 
-        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
+        tool_catalog = registry.get_schemas(module, valid_tools)
+        return tool_catalog
 
     async def _tool_recommendation(
-        self, task: str, data_desc: str, code_steps: str, available_tools: list
+        self,
+        context: [List[Message]],
+        code_steps: str,
+        available_tools: list
     ) -> list:
         """
-        Recommend tools for each step of the specified task
+        Recommend tools for the specified task.
 
         Args:
-            task (str): the task description
-            data_desc (str): the description of the dataset for the task
+            context (List[Message]): Action output history, source action denoted by Message.cause_by
             code_steps (str): the code steps to generate the full code for the task
             available_tools (list): the available tools for the task
 
         Returns:
-            list: recommended tools for each step of the specified task
+            list: recommended tools for the specified task
         """
-        prompt = TOOL_RECOMMENDATION_PROMPT.format(
-            task=task,
-            data_desc=data_desc,
+        system_prompt = TOOL_RECOMMENDATION_PROMPT.format(
             code_steps=code_steps,
             available_tools=available_tools,
         )
+        prompt = self.process_msg(context, system_prompt)
+
         tool_config = create_func_config(SELECT_FUNCTION_TOOLS)
         rsp = await self.llm.aask_code(prompt, **tool_config)
         recommend_tools = rsp["recommend_tools"]
@@ -156,50 +153,36 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
         context: List[Message],
         plan: Plan = None,
         task_guide: str = "",
-        data_desc: str = "",
     ) -> str:
         task_type = plan.current_task.task_type
-        task = plan.current_task.instruction
         available_tools = registry.get_all_schema_by_module(task_type)
-        available_tools = [
-            {k: tool[k] for k in ["name", "description"] if k in tool}
-            for tool in available_tools
-        ]
-        task_guide = "\n".join(
-            [f"Step {step.strip()}" for step in task_guide.split("\n")]
-        )
-
-        recommend_tools = await self._tool_recommendation(
-            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]
-        output_desc = TOOL_OUTPUT_DESC.get(task_type, "")
-        all_tasks = ""
-        completed_code = ""
 
-        for i, task in enumerate(plan.tasks):
-            stats = "DONE" if task.is_finished else "TODO"
-            all_tasks += f"Subtask {task.task_id}: {task.instruction}({stats})\n"
+        if len(available_tools) > 0:
+            available_tools = [
+                {k: tool[k] for k in ["name", "description"] if k in tool}
+                for tool in available_tools
+            ]
 
-        for task in plan.tasks:
-            if task.code:
-                completed_code += task.code + "\n"
+            recommend_tools = await self._tool_recommendation(context, task_guide, available_tools)
+            tool_catalog = self._parse_recommend_tools(task_type, recommend_tools)
+            logger.info(f"Recommended tools: \n{recommend_tools}")
 
-        prompt = TOO_ORGANIZATION_PROMPT.format(
-            all_tasks=all_tasks,
-            completed_code=completed_code,
-            data_desc=data_desc,
-            special_prompt=special_prompt,
-            code_steps=task_guide,
-            module_name=module_name,
-            output_desc=output_desc,
-            available_tools=recommend_tools,
-            tool_catalog=tool_catalog,
-        )
+            module_name = ML_MODULE_MAP[task_type]
+            output_desc = TOOL_OUTPUT_DESC.get(task_type, "")
+            prompt = TOO_ORGANIZATION_PROMPT.format(
+                special_prompt=special_prompt,
+                code_steps=task_guide,
+                module_name=module_name,
+                output_desc=output_desc,
+                function_catalog=tool_catalog,
+            )
+            context.append(Message(content=prompt, role="user"))
+        else:
+            context.append(Message(content=self.REUSE_CODE_INSTRUCTION, role="user"))
+            context.append(Message(content=special_prompt, role="user"))
+
+        prompt = self.process_msg(context)
         tool_config = create_func_config(CODE_GENERATOR_WITH_TOOLS)
         rsp = await self.llm.aask_code(prompt, **tool_config)
         return rsp["code"]
diff --git a/metagpt/prompts/ml_engineer.py b/metagpt/prompts/ml_engineer.py
index 0c4d036fc..d568bdd1f 100644
--- a/metagpt/prompts/ml_engineer.py
+++ b/metagpt/prompts/ml_engineer.py
@@ -4,25 +4,46 @@
 # @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.
+GEN_DATA_DESC_PROMPT = """
+Here is the head 5 rows of the dataset:
+{data_head}
 
+Please provide a brief one-sentence background of the dataset, and concise descriptions for each column. Keep descriptions short yet informative.
+
+Output the information in a JSON format, as shown in this example:
+```json
+{
+    "data_desc": "Brief dataset background.",
+    "column_desc": {
+        "column_name1": "Description of the first column.",
+        "column_name2": "Description of the second column.",
+        ...
+    }
+}
+```
+"""
+
+
+ASSIGN_TASK_TYPE_PROMPT = """
 Please assign a task type to each task in the list below from the given categories:
 {task_list}
+
+## All Task Type:
+- **feature_engineering**: Only for creating new columns for input data.
+- **data_preprocess**: Only for changing value inplace.
+- **model_train**: Only for training model.
+- **other**: Any tasks that do not fit into the previous categories, such as visualization, summarizing findings, build model, etc.
 """
 
 ASSIGN_TASK_TYPE = {
     "name": "assign_task_type",
-    "description": "assign task type to each task by order",
+    "description": "Assign task type to each task by order.",
     "parameters": {
         "type": "object",
         "properties": {
             "task_type": {
                 "type": "array",
-                "description": "List of task type.",
+                "description": "List of task type. The length should as long as task list",
                 "items": {
                     "type": "string",
                 },
@@ -34,43 +55,32 @@ ASSIGN_TASK_TYPE = {
 
 
 TOOL_RECOMMENDATION_PROMPT = """
-## Comprehensive Task Description:
-{task}
-
-## Dataset Description:
-Details about the dataset for the project:
-{data_desc}
-
-This task is divided into several steps, and you need to select the most suitable tools for each step. A tool means a function that can be used to help you solve the task.
-
-## Detailed Code Steps for the Task:
-{code_steps}
+Your are a tool recommender, the main goal is to recommend suitable tools for current task before coding. A tool means a function that can be used to help you solve the task.
 
 ## List of Available Tools:
 {available_tools}
 
+This is a task guide for the current task, including detailed code steps. You can refer to it when recommending tools.
+{code_steps}
+
 ## Tool Selection and Instructions:
-- For each code step listed above, choose up to five tools that are most likely to be useful in solving the task.
-- If you believe that no tools are suitable for a step, indicate with an empty list.
+- For the task, choose up to five tools that are most likely to be useful in solving the task.
+- If you believe that no tools are suitable, indicate with an empty list.
 - Only list the names of the tools, not the full schema of each tool.
 - The result should only contain tool names that are in the list of available tools.
-- The result list should be in the same order as the code steps.
 """
 
 SELECT_FUNCTION_TOOLS = {
     "name": "select_function_tools",
-    "description": "Given code steps to generate full code for a task, select suitable tools for each step by order.",
+    "description": "For current task, select suitable tools for it.",
     "parameters": {
         "type": "object",
         "properties": {
             "recommend_tools": {
                 "type": "array",
-                "description": "List of tool names for each code step. Empty list if no tool is suitable.",
+                "description": "List of tool names. Empty list if no tool is suitable.",
                 "items": {
-                    "type": "array",
-                    "items": {
-                        "type": "string",
-                    },
+                    "type": "string",
                 },
             },
         },
@@ -81,13 +91,13 @@ SELECT_FUNCTION_TOOLS = {
 
 CODE_GENERATOR_WITH_TOOLS = {
     "name": "add_subtask_code",
-    "description": "Add new code of current subtask to the end of an active Jupyter notebook.",
+    "description": "Add new code cell of current task to the end of an active Jupyter notebook.",
     "parameters": {
         "type": "object",
         "properties": {
             "code": {
                 "type": "string",
-                "description": "The code to be added.",
+                "description": "The code to be added to a new cell in jupyter.",
             },
         },
         "required": ["code"],
@@ -95,84 +105,60 @@ CODE_GENERATOR_WITH_TOOLS = {
 }
 
 TOO_ORGANIZATION_PROMPT = """
-As a senior data scientist, your role involves developing code for a specific sub-task within a larger project. This project is divided into several sub-tasks, which may either be new challenges or extensions of previous work.
+The previous conversation has provided all tasks step-by-step for the use goal and their statuses. 
+Now, begin writing code for the current task. This code should writen strictly on the basis of all previous completed tasks code, not a standalone code. And avoid writing duplicate code that has already been written in previous tasks, such as repeated import of packages, reading data, etc.
+Specifically, {special_prompt}
+You can utilize pre-defined tools in 'Available Tools' if the tools are sufficient. And you should combine the use of other public packages if necessary, like sklearn, numpy, pandas, etc..
 
-## Sub-tasks Overview
-Here's a list of all the sub-tasks, indicating their current status (DONE or TODO). Your responsibility is the first TODO task on this list.
-{all_tasks}
-
-## Historical Code (Previously Done Sub-tasks):
-This code, already executed in the Jupyter notebook, is critical for understanding the background and foundation for your current task.
-```python
-{completed_code}
-```
-
-## Dataset Description:
-Details about the dataset for the project:
-{data_desc}
-
-## Current Task Notion:
-{special_prompt}
-
-## Code Steps for Your Sub-task:
-Follow these steps to complete your current TODO task. You may use external Python functions or write custom code as needed. Ensure your code is self-contained.
+## Code Steps for Current Task:
+Follow steps below when you writing code if it's convenient.
 {code_steps}
 
-When you call a function, you should import the function from `{module_name}` first, e.g.:
-```python
-from metagpt.tools.functions.libs.feature_engineering import fill_missing_value
-```
-
-## Available Functions for Each Step:
-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:
+## Available Tools:
 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:
+When you call a function above, you should import the function from `{module_name}` first, e.g.:
 ```python
-# Step 1
-# Tools used: [function names or 'none']
-<your code for this step, without any comments>
+from metagpt.tools.functions.libs.data_preprocess import fill_missing_value
+```end
 
-# Step 2
+## Your Output Format:
+Generate the complete code for this task:
+```python
 # Tools used: [function names or 'none']
-<your code for this step, without any comments>
-
-# Continue with additional steps, following the same format...
+<your code for the current task, without any comments>
 ```end
 
 *** Important Rules ***
-- Use only the tools designated for each code step.
-- Your output should only include code for the current sub-task. Don't repeat historical code.
-- Only mention functions in comments if used in the code.
-- Ensure the output new code is executable in the current Jupyter notebook environment, with all historical code executed.
+- If you use tool not in the list, you should implement it by yourself.
+- Ensure the output new code is executable in the same Jupyter notebook environment with previous tasks code have been executed.
+- When write code for current task, remember the code should be coherent with previous tasks code.
+- Remember that don't process the columns have been processed in previous tasks and don't mock data yourself.
+- Prioritize using tools for the same functionality.
 """
 
-
 DATA_PREPROCESS_PROMPT = """
-In data preprocessing, closely monitor each column's data type. Apply suitable methods for various types (numerical, categorical, datetime, textual, etc.) to ensure the pandas.DataFrame is correctly formatted.
+The current task is about data preprocessing, closely monitor each column's data type. Apply suitable methods for various types (numerical, categorical, datetime, textual, etc.) to ensure the pandas.DataFrame is correctly formatted.
 Additionally, ensure that the columns being processed must be the ones that actually exist in the dataset.
+Don't write processed data to files.
 """
 
 FEATURE_ENGINEERING_PROMPT = """
-When performing feature engineering, please adhere to the following principles:
-- For specific user requests (such as removing a feature, creating a new feature based on existing data), directly generate the corresponding code.
-- In cases of unclear user requirements, write feature engineering code that you believe will most improve model performance. This may include feature transformation, combination, aggregation, etc., with a limit of five features at a time.
+The current task is about feature engineering. when performing it, please adhere to the following principles:
 - Ensure that the feature you're working with is indeed present in the dataset and consider the data type (numerical, categorical, etc.) and application scenario (classification, regression tasks, etc.).
-- 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.
+- When generate new features, you should combine real world knowledge and decide what features are useful for the task.
+- Generate as diverse features as possible to improve the model's performance.
+- Before generating a new feature, ensure the used features are already processed and ready to use.
 """
 
 MODEL_TRAIN_PROMPT = """
-When selecting and training a model, please follow these guidelines to ensure optimal performance:
+The current task is about training a model, please ensure high 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.
+- Before training, first check not is_numeric_dtype columns and use label encoding to convert them to numeric columns.
+- Use the data from previous task result directly, do not mock or reload data yourself.
 """
 
-
 DATA_PREPROCESS_OUTPUT_DESC = "Please note that all functions uniformly output a processed pandas.DataFrame, facilitating seamless integration into the broader workflow."
 
 FEATURE_ENGINEERING_OUTPUT_DESC = "Please note that all functions uniformly output updated pandas.DataFrame with feature engineering applied."
@@ -185,20 +171,15 @@ REGRESSION_MODEL_OUTPUT_DESC = ""
 ML_SPECIFIC_PROMPT = {
     "data_preprocess": DATA_PREPROCESS_PROMPT,
     "feature_engineering": FEATURE_ENGINEERING_PROMPT,
-    "classification_model": MODEL_TRAIN_PROMPT,
-    "regression_model": MODEL_TRAIN_PROMPT,
+    "model_train": MODEL_TRAIN_PROMPT,
 }
 
 TOOL_OUTPUT_DESC = {
     "data_preprocess": DATA_PREPROCESS_OUTPUT_DESC,
     "feature_engineering": FEATURE_ENGINEERING_OUTPUT_DESC,
-    "classification_model": CLASSIFICATION_MODEL_OUTPUT_DESC,
-    "regression_model": REGRESSION_MODEL_OUTPUT_DESC,
 }
 
 ML_MODULE_MAP = {
-    "data_preprocess": "metagpt.tools.functions.libs.machine_learning.data_preprocess",
-    "feature_engineering": "metagpt.tools.functions.libs.machine_learning.feature_engineering",
-    "classification_model": "metagpt.tools.functions.libs.machine_learning.ml_model",
-    "regression_model": "metagpt.tools.functions.libs.machine_learning.ml_model",
+    "data_preprocess": "metagpt.tools.functions.libs.data_preprocess",
+    "feature_engineering": "metagpt.tools.functions.libs.feature_engineering",
 }
diff --git a/metagpt/tools/functions/__init__.py b/metagpt/tools/functions/__init__.py
index b81e85833..30ee10827 100644
--- a/metagpt/tools/functions/__init__.py
+++ b/metagpt/tools/functions/__init__.py
@@ -6,3 +6,4 @@
 # @Desc    :
 from metagpt.tools.functions.register.register import registry
 import metagpt.tools.functions.libs.feature_engineering
+import metagpt.tools.functions.libs.data_preprocess

From 21d97a23bb65b92a0379ff101ecbd497bd6e8537 Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Wed, 6 Dec 2023 17:31:51 +0800
Subject: [PATCH 04/49] output code_steps to json

---
 metagpt/actions/write_analysis_code.py |  1 -
 metagpt/actions/write_code_steps.py    | 25 ++++++++++++++-----------
 metagpt/roles/ml_engineer.py           |  2 +-
 3 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/metagpt/actions/write_analysis_code.py b/metagpt/actions/write_analysis_code.py
index cfec95deb..71467edd0 100644
--- a/metagpt/actions/write_analysis_code.py
+++ b/metagpt/actions/write_analysis_code.py
@@ -153,7 +153,6 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
         context: List[Message],
         plan: Plan = None,
         code_steps: str = "",
-        data_desc: str = "",
     ) -> str:
         task_type = plan.current_task.task_type
         available_tools = registry.get_all_schema_by_module(task_type)
diff --git a/metagpt/actions/write_code_steps.py b/metagpt/actions/write_code_steps.py
index d3f6e5553..0bfb9c225 100644
--- a/metagpt/actions/write_code_steps.py
+++ b/metagpt/actions/write_code_steps.py
@@ -4,18 +4,12 @@ from typing import Dict, List, Union
 
 from metagpt.actions import Action
 from metagpt.schema import Message, Task, Plan
-
+from metagpt.utils.common import CodeParser
 
 CODE_STEPS_PROMPT_TEMPLATE = """
 # Context
 {context}
 
-##  Format example
-1.
-2.
-3.
-...
-
 -----
 Tasks are all code development tasks.
 You are a professional engineer, the main goal is to plan out concise solution steps for Current Task before coding.
@@ -25,7 +19,16 @@ The output plan should following the subsequent principles:
 1.The plan is a rough checklist of steps outlining the entire program's structure.Try to keep the number of steps fewer than 5.
 2.The steps should be written concisely and at a high level, avoiding overly detailed implementation specifics.
 3.The execution of the plan happens sequentially, but the plan can incorporate conditional (if) and looping(loop) keywords for more complex structures.
-4.Output carefully referenced "Format example" in format.
+
+Output the code steps in a JSON format, as shown in this example:
+```json
+{
+    "Step 1": "",
+    "Step 2": "",
+    "Step 3": "",
+    ...
+}
+```
 """
 
 STRUCTURAL_CONTEXT = """
@@ -51,10 +54,11 @@ class WriteCodeSteps(Action):
         """
 
         context = self.get_context(plan)
-        code_steps_prompt = CODE_STEPS_PROMPT_TEMPLATE.format(
-            context=context,
+        code_steps_prompt = CODE_STEPS_PROMPT_TEMPLATE.replace(
+            "{context}", context
         )
         code_steps = await self._aask(code_steps_prompt)
+        code_steps = CodeParser.parse_code(block=None, text=code_steps)
         return code_steps
 
     def get_context(self, plan: Plan):
@@ -74,4 +78,3 @@ class WriteCodeSteps(Action):
         )
         # print(context)
         return context
-
diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index 148851e9e..c2841be4c 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -294,7 +294,7 @@ if __name__ == "__main__":
     # requirement = "Run data analysis on sklearn Diabetes dataset, include a plot"
     # requirement = "Run data analysis on sklearn Wine recognition dataset, include a plot, and train a model to predict wine class (20% as validation), and show validation accuracy"
     # 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"
-    requirement = "Run EDA and visualization on this dataset, train a model to predict survival, report metrics on validation set (20%), dataset: workspace/titanic/train.csv"
+    # requirement = "Run EDA and visualization on this dataset, train a model to predict survival, report metrics on validation set (20%), dataset: workspace/titanic/train.csv"
 
     requirement = "Perform data analysis on the provided data. Train a model to predict the target variable Survived. Include data preprocessing, feature engineering, and modeling in your pipeline. The metric is accuracy."
     data_path = "/data/lidanyang/tabular_data/titanic"

From 757174366e49cb3f0a8c460b8ba8075baedc2ac7 Mon Sep 17 00:00:00 2001
From: stellahsr <stellahsr@126.com>
Date: Wed, 6 Dec 2023 20:45:37 +0800
Subject: [PATCH 05/49] update locally

---
 config/config.yaml                  | 15 ++++++++-------
 metagpt/roles/ml_engineer.py        | 18 +++++++++++-------
 metagpt/tools/functions/__init__.py |  2 +-
 metagpt/tools/web_browser_engine.py |  2 +-
 metagpt/utils/__init__.py           |  4 ++--
 requirements.txt                    |  2 --
 6 files changed, 23 insertions(+), 20 deletions(-)

diff --git a/config/config.yaml b/config/config.yaml
index bed67083c..694251f17 100644
--- a/config/config.yaml
+++ b/config/config.yaml
@@ -5,7 +5,7 @@
 ## The official OPENAI_API_BASE is https://api.openai.com/v1
 ## If the official OPENAI_API_BASE is not available, we recommend using the [openai-forward](https://github.com/beidongjiedeguang/openai-forward).
 ## Or, you can configure OPENAI_PROXY to access official OPENAI_API_BASE.
-OPENAI_API_BASE: "https://api.openai.com/v1"
+#OPENAI_API_BASE: "https://api.openai.com/v1"
 #OPENAI_PROXY: "http://127.0.0.1:8118"
 #OPENAI_API_KEY: "YOUR_API_KEY"   # set the value to sk-xxx if you host the openai interface for open llm model
 OPENAI_API_MODEL: "gpt-4"
@@ -24,12 +24,13 @@ RPM: 10
 
 #### if AZURE, check https://github.com/openai/openai-cookbook/blob/main/examples/azure/chat.ipynb
 #### You can use ENGINE or DEPLOYMENT mode
-#OPENAI_API_TYPE: "azure"
-#OPENAI_API_BASE: "YOUR_AZURE_ENDPOINT"
-#OPENAI_API_KEY: "YOUR_AZURE_API_KEY"
-#OPENAI_API_VERSION: "YOUR_AZURE_API_VERSION"
-#DEPLOYMENT_NAME: "YOUR_DEPLOYMENT_NAME"
-#DEPLOYMENT_ID: "YOUR_DEPLOYMENT_ID"
+OPENAI_API_TYPE: "azure"
+OPENAI_API_BASE: "https://deepwisdom.openai.azure.com/"
+OPENAI_API_KEY: "02ae6058d09849c691176befeae2107c"
+#OPENAI_API_VERSION: "2023-05-15"
+OPENAI_API_VERSION: "2023-07-01-preview"
+DEPLOYMENT_ID: "GPT-4"
+OPENAI_API_ENGINE: "gpt-4"
 
 #### if zhipuai from `https://open.bigmodel.cn`. You can set here or export API_KEY="YOUR_API_KEY"
 # ZHIPUAI_API_KEY: "YOUR_API_KEY"
diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index 15edb2b06..c088ff104 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -10,7 +10,7 @@ from metagpt.actions import Action
 from metagpt.actions.execute_code import ExecutePyCode
 from metagpt.actions.write_analysis_code import WriteCodeByGenerate, WriteCodeWithTools
 from metagpt.actions.write_plan import WritePlan
-from metagpt.actions.write_task_guide import WriteTaskGuide
+# from metagpt.actions.write_task_guide import WriteTaskGuide
 from metagpt.logs import logger
 from metagpt.prompts.ml_engineer import GEN_DATA_DESC_PROMPT
 from metagpt.roles import Role
@@ -39,7 +39,7 @@ catboost
 def truncate(result: str, keep_len: int = 1000) -> str:
     desc = "Truncated to show only the last 1000 characters\n"
     if result.startswith(desc):
-        result = result[-len(desc) :]
+        result = result[-len(desc):]
 
     if len(result) > keep_len:
         result = result[-keep_len:]
@@ -110,9 +110,9 @@ class AskReview(Action):
         logger.info("most recent context:")
         latest_action = context[-1].cause_by.__name__ if context[-1].cause_by else ""
         prompt = f"\nPlease review output from {latest_action}:\n" \
-            "If you want to change a task in the plan, say 'change task task_id, ... (things to change)'\n" \
-            "If you confirm the output and wish to continue with the current process, type CONFIRM\n" \
-            "If you want to terminate the process, type exit:\n"
+                 "If you want to change a task in the plan, say 'change task task_id, ... (things to change)'\n" \
+                 "If you confirm the output and wish to continue with the current process, type CONFIRM\n" \
+                 "If you want to terminate the process, type exit:\n"
         rsp = input(prompt)
 
         if rsp.lower() in ("exit"):
@@ -148,7 +148,7 @@ class GenerateDataDesc(Action):
 
 class MLEngineer(Role):
     def __init__(
-        self, name="ABC", profile="MLEngineer", goal="", auto_run: bool = False, data_path: str = None
+            self, name="ABC", profile="MLEngineer", goal="", auto_run: bool = False, data_path: str = None
     ):
         super().__init__(name=name, profile=profile, goal=goal)
         self._set_react_mode(react_mode="plan_and_act")
@@ -300,11 +300,15 @@ if __name__ == "__main__":
     # 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"
     # requirement = "Run EDA and visualization on this dataset, train a model to predict survival, report metrics on validation set (20%), dataset: workspace/titanic/train.csv"
 
+    from metagpt.const import DATA_PATH
+
     requirement = "Perform data analysis on the provided data. Train a model to predict the target variable Survived. Include data preprocessing, feature engineering, and modeling in your pipeline. The metric is accuracy."
-    data_path = "/data/lidanyang/tabular_data/titanic"
+    data_path = f"{DATA_PATH}/titanic"
+
 
     async def main(requirement: str = requirement, auto_run: bool = True, data_path: str = data_path):
         role = MLEngineer(goal=requirement, auto_run=auto_run, data_path=data_path)
         await role.run(requirement)
 
+
     fire.Fire(main)
diff --git a/metagpt/tools/functions/__init__.py b/metagpt/tools/functions/__init__.py
index 30ee10827..d4a1ff73b 100644
--- a/metagpt/tools/functions/__init__.py
+++ b/metagpt/tools/functions/__init__.py
@@ -6,4 +6,4 @@
 # @Desc    :
 from metagpt.tools.functions.register.register import registry
 import metagpt.tools.functions.libs.feature_engineering
-import metagpt.tools.functions.libs.data_preprocess
+# import metagpt.tools.functions.libs.data_preprocess
diff --git a/metagpt/tools/web_browser_engine.py b/metagpt/tools/web_browser_engine.py
index 453d87f31..7228ae9cf 100644
--- a/metagpt/tools/web_browser_engine.py
+++ b/metagpt/tools/web_browser_engine.py
@@ -7,7 +7,7 @@ from typing import Any, Callable, Coroutine, Literal, overload
 
 from metagpt.config import CONFIG
 from metagpt.tools import WebBrowserEngineType
-from metagpt.utils.parse_html import WebPage
+# from metagpt.utils.parse_html import WebPage
 
 
 class WebBrowserEngine:
diff --git a/metagpt/utils/__init__.py b/metagpt/utils/__init__.py
index f13175cf8..86cac50db 100644
--- a/metagpt/utils/__init__.py
+++ b/metagpt/utils/__init__.py
@@ -6,7 +6,7 @@
 @File    : __init__.py
 """
 
-from metagpt.utils.read_document import read_docx
+# from metagpt.utils.read_document import read_docx
 from metagpt.utils.singleton import Singleton
 from metagpt.utils.token_counter import (
     TOKEN_COSTS,
@@ -16,7 +16,7 @@ from metagpt.utils.token_counter import (
 
 
 __all__ = [
-    "read_docx",
+    # "read_docx",
     "Singleton",
     "TOKEN_COSTS",
     "count_message_tokens",
diff --git a/requirements.txt b/requirements.txt
index 1d1bc95a1..9b75fd200 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -35,7 +35,6 @@ tqdm==4.64.0
 # webdriver_manager<3.9
 anthropic==0.3.6
 typing-inspect==0.8.0
-typing_extensions==4.5.0
 libcst==1.0.1
 qdrant-client==1.4.0
 pytest-mock==3.11.1
@@ -46,7 +45,6 @@ wrapt==1.15.0
 websocket-client==0.58.0
 zhipuai==1.0.7
 rich==13.6.0
-nbclient==0.9.0
 nbformat==5.9.2
 ipython==8.17.2
 ipykernel==6.27.0

From f26b2c135922eeb539fc4c907b086bbefdddff19 Mon Sep 17 00:00:00 2001
From: stellahsr <stellahsr@126.com>
Date: Thu, 7 Dec 2023 19:21:27 +0800
Subject: [PATCH 06/49] =?UTF-8?q?=E5=8F=96=E6=B6=88=E6=B3=A8=E9=87=8A?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 metagpt/tools/functions/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/metagpt/tools/functions/__init__.py b/metagpt/tools/functions/__init__.py
index d4a1ff73b..30ee10827 100644
--- a/metagpt/tools/functions/__init__.py
+++ b/metagpt/tools/functions/__init__.py
@@ -6,4 +6,4 @@
 # @Desc    :
 from metagpt.tools.functions.register.register import registry
 import metagpt.tools.functions.libs.feature_engineering
-# import metagpt.tools.functions.libs.data_preprocess
+import metagpt.tools.functions.libs.data_preprocess

From 204cda844fba774910baaa21417a40c9ae8171d8 Mon Sep 17 00:00:00 2001
From: stellahsr <stellahsr@126.com>
Date: Thu, 7 Dec 2023 19:22:19 +0800
Subject: [PATCH 07/49] fix typo

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

diff --git a/metagpt/actions/write_analysis_code.py b/metagpt/actions/write_analysis_code.py
index c8a28edd1..957d35f7e 100644
--- a/metagpt/actions/write_analysis_code.py
+++ b/metagpt/actions/write_analysis_code.py
@@ -192,7 +192,7 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
             output_desc = TOOL_OUTPUT_DESC.get(task_type, "")
 
             hist_info = f"Previous finished code is \n\n ```Python {final_code} ``` \n\n " \
-                        f"Conde runtime result is {result} \n\n"
+                        f"Runtime result is {result} \n\n"
 
             prompt = TOOL_USAGE_PROMPT.format(
                 goal=plan.current_task.instruction,
@@ -213,7 +213,7 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
 
         else:
             hist_info = f"Previous finished code is \n\n ```Python {code_context} ``` \n\n " \
-                     f"Conde runtime result is {result} \n\n"
+                     f"runtime result is {result} \n\n"
 
             prompt = GENERATE_CODE_PROMPT.format(
                 goal=plan.current_task.instruction,

From ba6a62f55aa5546d9ac274db1416d90b91c17bfb Mon Sep 17 00:00:00 2001
From: stellahsr <stellahsr@126.com>
Date: Thu, 7 Dec 2023 19:24:21 +0800
Subject: [PATCH 08/49] update ignore

---
 .gitignore | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/.gitignore b/.gitignore
index e03eab3d3..d01469a36 100644
--- a/.gitignore
+++ b/.gitignore
@@ -148,6 +148,9 @@ allure-results
 .DS_Store
 .vscode
 
+# Config
+config/config.yaml
+
 log.txt
 docs/scripts/set_env.sh
 key.yaml

From 7e343a100b449a8441ab55063ad76661d0391f46 Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Thu, 7 Dec 2023 20:45:08 +0800
Subject: [PATCH 09/49] update ml functions

---
 .../tools/functions/libs/data_preprocess.py   | 29 ++++++-------
 .../functions/libs/feature_engineering.py     | 42 +++++++++++++------
 .../functions/schemas/data_preprocess.py      | 21 ++++++----
 .../functions/schemas/feature_engineering.py  | 36 ++++++++++------
 4 files changed, 80 insertions(+), 48 deletions(-)

diff --git a/metagpt/tools/functions/libs/data_preprocess.py b/metagpt/tools/functions/libs/data_preprocess.py
index 68c96bbc9..5579c5bd8 100644
--- a/metagpt/tools/functions/libs/data_preprocess.py
+++ b/metagpt/tools/functions/libs/data_preprocess.py
@@ -1,15 +1,12 @@
-
-import pandas as pd
 import numpy as np
-
 from sklearn.impute import SimpleImputer
-from sklearn.preprocessing import LabelEncoder
 from sklearn.preprocessing import KBinsDiscretizer
-from sklearn.preprocessing import MinMaxScaler
-from sklearn.preprocessing import StandardScaler
 from sklearn.preprocessing import MaxAbsScaler
-from sklearn.preprocessing import RobustScaler
+from sklearn.preprocessing import MinMaxScaler
+from sklearn.preprocessing import OneHotEncoder
 from sklearn.preprocessing import OrdinalEncoder
+from sklearn.preprocessing import RobustScaler
+from sklearn.preprocessing import StandardScaler
 
 from metagpt.tools.functions import registry
 from metagpt.tools.functions.schemas.data_preprocess import *
@@ -21,13 +18,6 @@ def fill_missing_value(df: pd.DataFrame, features: list, strategy: str = 'mean',
     return df
 
 
-# @registry.register("data_preprocess", FillMissingValue)
-# def label_encode(df: pd.DataFrame, features: list,):
-#     for col in features:
-#         df[col] = LabelEncoder().fit_transform(df[col])
-#     return df
-
-
 @registry.register("data_preprocess", SplitBins)
 def split_bins(df: pd.DataFrame, features: list, strategy: str = 'quantile',):
     df[features] = KBinsDiscretizer(strategy=strategy, encode='ordinal').fit_transform(df[features])
@@ -73,6 +63,17 @@ def ordinal_encode(df: pd.DataFrame, features: list,):
     return df
 
 
+@registry.register("data_preprocess", OneHotEncoding)
+def one_hot_encoding(df, cols):
+    enc = OneHotEncoder(handle_unknown="ignore", sparse=False)
+    ts_data = enc.fit_transform(df[cols])
+    new_columns = enc.get_feature_names_out(cols)
+    ts_data = pd.DataFrame(ts_data, columns=new_columns, index=df.index)
+    df.drop(cols, axis=1, inplace=True)
+    df = pd.concat([df, ts_data], axis=1)
+    return df
+
+
 if __name__ == '__main__':
     def run():
         V = {
diff --git a/metagpt/tools/functions/libs/feature_engineering.py b/metagpt/tools/functions/libs/feature_engineering.py
index 0573f362d..4780e4fa0 100644
--- a/metagpt/tools/functions/libs/feature_engineering.py
+++ b/metagpt/tools/functions/libs/feature_engineering.py
@@ -8,7 +8,8 @@ import itertools
 
 from dateutil.relativedelta import relativedelta
 from pandas.api.types import is_numeric_dtype
-from sklearn.preprocessing import PolynomialFeatures, OneHotEncoder
+from sklearn.model_selection import KFold
+from sklearn.preprocessing import PolynomialFeatures
 
 from metagpt.tools.functions import registry
 from metagpt.tools.functions.schemas.feature_engineering import *
@@ -29,17 +30,6 @@ def polynomial_expansion(df, cols, degree=2):
     return df
 
 
-@registry.register("feature_engineering", OneHotEncoding)
-def one_hot_encoding(df, cols):
-    enc = OneHotEncoder(handle_unknown="ignore", sparse=False)
-    ts_data = enc.fit_transform(df[cols])
-    new_columns = enc.get_feature_names_out(cols)
-    ts_data = pd.DataFrame(ts_data, columns=new_columns, index=df.index)
-    df.drop(cols, axis=1, inplace=True)
-    df = pd.concat([df, ts_data], axis=1)
-    return df
-
-
 @registry.register("feature_engineering", FrequencyEncoding)
 def frequency_encoding(df, cols):
     for col in cols:
@@ -48,6 +38,31 @@ def frequency_encoding(df, cols):
     return df
 
 
+@registry.register("feature_engineering", TargetMeanEncoder)
+def target_mean_encoder(df, col, label):
+    encoder_dict = df.groupby(col)[label].mean().to_dict()
+    df[f"{col}_target_mean"] = df[col].map(encoder_dict)
+    return df
+
+
+@registry.register("feature_engineering", KFoldTargetMeanEncoder)
+def k_fold_target_mean_encoder(df, col, label, n_splits=5, random_state=2021):
+    tmp = df.copy()
+    kf = KFold(n_splits=n_splits, shuffle=True, random_state=random_state)
+
+    global_mean = tmp[label].mean()
+    col_name = f"{col}_kf_target_mean"
+    for trn_idx, val_idx in kf.split(tmp, tmp[label]):
+        _trn, _val = tmp.iloc[trn_idx], tmp.iloc[val_idx]
+        tmp.loc[tmp.index[val_idx], col_name] = _val[col].map(
+            _trn.groupby(col)[label].mean()
+        )
+    tmp[col_name].fillna(global_mean, inplace=True)
+    encoder_dict = tmp.groupby(col)[col_name].mean().to_dict()
+    df[f"{col}_kf_target_mean"] = df[col].map(encoder_dict)
+    return df
+
+
 @registry.register("feature_engineering", CatCross)
 def cat_cross(df, cols, max_cat_num=100):
     for col in cols:
@@ -56,7 +71,8 @@ def cat_cross(df, cols, max_cat_num=100):
 
     for col1, col2 in itertools.combinations(cols, 2):
         cross_col = f"{col1}_cross_{col2}"
-        df[cross_col] = df[col1].astype(str) + "_" + df[col2].astype(str)
+        crossed = df[col1].astype(str) + "_" + df[col2].astype(str)
+        df[cross_col] = crossed.astype('category').cat.codes
     return df
 
 
diff --git a/metagpt/tools/functions/schemas/data_preprocess.py b/metagpt/tools/functions/schemas/data_preprocess.py
index 40e1d64e0..16b97aeac 100644
--- a/metagpt/tools/functions/schemas/data_preprocess.py
+++ b/metagpt/tools/functions/schemas/data_preprocess.py
@@ -8,14 +8,13 @@ class FillMissingValue(ToolSchema):
     """Completing missing values with simple strategies"""
     df: pd.DataFrame = tool_field(description="input dataframe")
     features: list = tool_field(description="columns to be processed")
-    strategy: str = tool_field(description="the imputation strategy", default='mean')
-    fill_value: int = tool_field(description="fill_value is used to replace all occurrences of missing_values", default=None)
-
-
-# class LabelEncode(ToolSchema):
-#     """Completing missing values with simple strategies"""
-#     df: pd.DataFrame = tool_field(description="input dataframe")
-#     features: list = tool_field(description="columns to be processed")
+    strategy: str = tool_field(
+        description="the imputation strategy",
+        default='mean',
+        enum=['mean', 'median', 'most_frequent', 'constant']
+    )
+    fill_value: int = tool_field(
+        description="fill_value is used to replace all occurrences of missing_values", default=None)
 
 
 class SplitBins(ToolSchema):
@@ -60,3 +59,9 @@ class OrdinalEncode(ToolSchema):
     df: pd.DataFrame = tool_field(description="input dataframe")
     features: list = tool_field(description="columns to be processed")
 
+
+class OneHotEncoding(ToolSchema):
+    """Apply one-hot encoding to specified categorical columns, the original columns will be dropped."""
+
+    df: pd.DataFrame = tool_field(description="DataFrame to process.")
+    cols: list = tool_field(description="Categorical columns to be one-hot encoded and dropped.")
diff --git a/metagpt/tools/functions/schemas/feature_engineering.py b/metagpt/tools/functions/schemas/feature_engineering.py
index df2eebff6..5c89d9b16 100644
--- a/metagpt/tools/functions/schemas/feature_engineering.py
+++ b/metagpt/tools/functions/schemas/feature_engineering.py
@@ -12,29 +12,39 @@ from metagpt.tools.functions.schemas.base import ToolSchema, tool_field
 
 
 class PolynomialExpansion(ToolSchema):
-    """Generate polynomial and interaction features from selected columns, excluding the bias column."""
+    """Add polynomial and interaction features from selected numeric columns, excluding the bias column."""
 
     df: pd.DataFrame = tool_field(description="DataFrame to process.")
     cols: list = tool_field(description="Columns for polynomial expansion.")
     degree: int = tool_field(description="Degree of polynomial features.", default=2)
 
 
-class OneHotEncoding(ToolSchema):
-    """Apply one-hot encoding to specified categorical columns, the original columns will be dropped."""
-
-    df: pd.DataFrame = tool_field(description="DataFrame to process.")
-    cols: list = tool_field(description="Categorical columns to be one-hot encoded.")
-
-
 class FrequencyEncoding(ToolSchema):
-    """Convert categorical columns to frequency encoding."""
+    """Add value counts of categorical columns as new features."""
 
     df: pd.DataFrame = tool_field(description="DataFrame to process.")
     cols: list = tool_field(description="Categorical columns to be frequency encoded.")
 
 
+class TargetMeanEncoder(ToolSchema):
+    """Encodes a categorical column by the mean of the label column, and adds the result as a new feature."""
+
+    df: pd.DataFrame = tool_field(description="DataFrame to process.")
+    col: str = tool_field(description="Column to be mean encoded.")
+    label: str = tool_field(description="Predicted label column.")
+
+
+class KFoldTargetMeanEncoder(ToolSchema):
+    """Adds a new feature to the DataFrame by k-fold mean encoding of a categorical column using the label column."""
+    df: pd.DataFrame = tool_field(description="DataFrame to process.")
+    col: str = tool_field(description="Column to be k-fold mean encoded.")
+    label: str = tool_field(description="Predicted label column.")
+    n_splits: int = tool_field(description="Number of splits for K-fold.", default=5)
+    random_state: int = tool_field(description="Random seed.", default=2021)
+
+
 class CatCross(ToolSchema):
-    """Create pairwise crossed features from categorical columns, joining values with '_'."""
+    """Add pairwise crossed features and convert them to numerical features."""
 
     df: pd.DataFrame = tool_field(description="DataFrame to process.")
     cols: list = tool_field(description="Columns to be pairwise crossed.")
@@ -44,7 +54,7 @@ class CatCross(ToolSchema):
 
 
 class GroupStat(ToolSchema):
-    """Perform aggregation operations on a specified column grouped by certain categories."""
+    """Aggregate specified column in a DataFrame grouped by another column, adding new features named '<agg_col>_<agg_func>_by_<group_col>'."""
 
     df: pd.DataFrame = tool_field(description="DataFrame to process.")
     group_col: str = tool_field(description="Column used for grouping.")
@@ -56,7 +66,7 @@ class GroupStat(ToolSchema):
 
 
 class ExtractTimeComps(ToolSchema):
-    """Extract specific time components from a designated time column in a DataFrame."""
+    """Extract and add specific time components as new features from a designated time column."""
 
     df: pd.DataFrame = tool_field(description="DataFrame to process.")
     time_col: str = tool_field(
@@ -69,7 +79,7 @@ class ExtractTimeComps(ToolSchema):
 
 
 class FeShiftByTime(ToolSchema):
-    """Shift column values in a DataFrame based on specified time intervals."""
+    """Shift column values based on specified time intervals and add the resulting new features to the DataFrame. New features are named in the format of '<group_col>_<shift_col>_lag_<period>_<freq>'."""
 
     df: pd.DataFrame = tool_field(description="DataFrame to process.")
     time_col: str = tool_field(description="Column for time-based shifting.")

From fe2b79fedc407afe72ad855ea6187afe11108beb Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Thu, 7 Dec 2023 20:48:00 +0800
Subject: [PATCH 10/49] refine ml prompt

---
 metagpt/actions/write_analysis_code.py | 114 +++++++++---------------
 metagpt/prompts/ml_engineer.py         | 118 ++++++++++++++++++++++---
 metagpt/roles/ml_engineer.py           |  75 ++++++----------
 metagpt/utils/common.py                |  14 +++
 4 files changed, 192 insertions(+), 129 deletions(-)

diff --git a/metagpt/actions/write_analysis_code.py b/metagpt/actions/write_analysis_code.py
index 957d35f7e..f96ade1b4 100644
--- a/metagpt/actions/write_analysis_code.py
+++ b/metagpt/actions/write_analysis_code.py
@@ -15,15 +15,11 @@ from metagpt.prompts.ml_engineer import (
     TOO_ORGANIZATION_PROMPT,
     ML_SPECIFIC_PROMPT,
     ML_MODULE_MAP,
-    TOOL_OUTPUT_DESC,
-    TOOL_USAGE_PROMPT,
+    TOOL_OUTPUT_DESC, DATA_PROCESS_PROMPT,
 )
 from metagpt.schema import Message, Plan
 from metagpt.tools.functions import registry
-from metagpt.utils.common import create_func_config
-from metagpt.prompts.ml_engineer import GEN_DATA_DESC_PROMPT, GENERATE_CODE_PROMPT
-from metagpt.utils.common import CodeParser
-from metagpt.actions.execute_code import ExecutePyCode
+from metagpt.utils.common import create_func_config, remove_comments
 
 
 class BaseWriteAnalysisCode(Action):
@@ -51,13 +47,13 @@ class BaseWriteAnalysisCode(Action):
 
         # 添加默认的提示词
         if (
-                default_system_msg not in messages[0]["content"]
-                and messages[0]["role"] != "system"
+            default_system_msg not in messages[0]["content"]
+            and messages[0]["role"] != "system"
         ):
             messages.insert(0, {"role": "system", "content": default_system_msg})
         elif (
-                default_system_msg not in messages[0]["content"]
-                and messages[0]["role"] == "system"
+            default_system_msg not in messages[0]["content"]
+            and messages[0]["role"] == "system"
         ):
             messages[0] = {
                 "role": "system",
@@ -66,7 +62,7 @@ class BaseWriteAnalysisCode(Action):
         return messages
 
     async def run(
-            self, context: List[Message], plan: Plan = None, code_steps: str = ""
+        self, context: List[Message], plan: Plan = None, code_steps: str = ""
     ) -> str:
         """Run of a code writing action, used in data analysis or modeling
 
@@ -87,12 +83,12 @@ class WriteCodeByGenerate(BaseWriteAnalysisCode):
         super().__init__(name, context, llm)
 
     async def run(
-            self,
-            context: [List[Message]],
-            plan: Plan = None,
-            code_steps: str = "",
-            system_msg: str = None,
-            **kwargs,
+        self,
+        context: [List[Message]],
+        plan: Plan = None,
+        code_steps: str = "",
+        system_msg: str = None,
+        **kwargs,
     ) -> str:
         context.append(Message(content=self.REUSE_CODE_INSTRUCTION, role="user"))
         prompt = self.process_msg(context, system_msg)
@@ -102,7 +98,6 @@ class WriteCodeByGenerate(BaseWriteAnalysisCode):
 
 class WriteCodeWithTools(BaseWriteAnalysisCode):
     """Write code with help of local available tools. Choose tools first, then generate code to use the tools"""
-    execute_code = ExecutePyCode()
 
     @staticmethod
     def _parse_recommend_tools(module: str, recommend_tools: list) -> List[Dict]:
@@ -126,10 +121,10 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
         return tool_catalog
 
     async def _tool_recommendation(
-            self,
-            context: [List[Message]],
-            code_steps: str,
-            available_tools: list
+        self,
+        task: str,
+        code_steps: str,
+        available_tools: list
     ) -> list:
         """
         Recommend tools for the specified task.
@@ -142,86 +137,63 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
         Returns:
             list: recommended tools for the specified task
         """
-        system_prompt = TOOL_RECOMMENDATION_PROMPT.format(
+        prompt = TOOL_RECOMMENDATION_PROMPT.format(
+            current_task=task,
             code_steps=code_steps,
             available_tools=available_tools,
         )
-        prompt = self.process_msg(context, system_prompt)
-
         tool_config = create_func_config(SELECT_FUNCTION_TOOLS)
         rsp = await self.llm.aask_code(prompt, **tool_config)
         recommend_tools = rsp["recommend_tools"]
         return recommend_tools
 
-
     async def run(
-            self,
-            context: List[Message],
-            plan: Plan = None,
-            code_steps: str = "",
-            **kwargs,
+        self,
+        context: List[Message],
+        plan: Plan = None,
+        code_steps: str = "",
+        column_info: str = "",
     ) -> str:
         task_type = plan.current_task.task_type
-        logger.info(f"task_type is: {task_type}")
         available_tools = registry.get_all_schema_by_module(task_type)
-
-        # special_prompt = ML_SPECIFIC_PROMPT.get(task_type, "")
+        special_prompt = ML_SPECIFIC_PROMPT.get(task_type, "")
 
         finished_tasks = plan.get_finished_tasks()
-        code_context = [task.code for task in finished_tasks]
-
+        code_context = [remove_comments(task.code) for task in finished_tasks]
         code_context = "\n\n".join(code_context)
 
-        ### add runtime info
-        result, success = await self.execute_code.run(code_context)
-        logger.info(result)
-
         if len(available_tools) > 0:
             available_tools = [
                 {k: tool[k] for k in ["name", "description"] if k in tool}
                 for tool in available_tools
             ]
 
-            final_code = code_context
-
-            recommend_tools = await self._tool_recommendation(context, code_steps, available_tools)
+            recommend_tools = await self._tool_recommendation(
+                plan.current_task.instruction,
+                code_steps,
+                available_tools
+            )
             tool_catalog = self._parse_recommend_tools(task_type, recommend_tools)
             logger.info(f"Recommended tools: \n{recommend_tools}")
 
             module_name = ML_MODULE_MAP[task_type]
             output_desc = TOOL_OUTPUT_DESC.get(task_type, "")
-
-            hist_info = f"Previous finished code is \n\n ```Python {final_code} ``` \n\n " \
-                        f"Runtime result is {result} \n\n"
-
-            prompt = TOOL_USAGE_PROMPT.format(
-                goal=plan.current_task.instruction,
-                context=hist_info,
+            prompt = DATA_PROCESS_PROMPT.format(
+                user_requirement=plan.goal,
+                history_code=code_context,
+                current_task=plan.current_task.instruction,
+                column_info=column_info,
+                special_prompt=special_prompt,
                 code_steps=code_steps,
                 module_name=module_name,
                 output_desc=output_desc,
                 function_catalog=tool_catalog,
             )
-
-            tool_config = create_func_config(CODE_GENERATOR_WITH_TOOLS)
-
-            rsp = await self.llm.aask_code(prompt, **tool_config)
-            logger.info(f"rsp is: {rsp}")
-            final_code = final_code + "\n\n" + rsp["code"]
-
-            return final_code
-
         else:
-            hist_info = f"Previous finished code is \n\n ```Python {code_context} ``` \n\n " \
-                     f"runtime result is {result} \n\n"
+            context.append(Message(content=self.REUSE_CODE_INSTRUCTION, role="user"))
+            context.append(Message(content=special_prompt, role="user"))
+            prompt = self.process_msg(context)
 
-            prompt = GENERATE_CODE_PROMPT.format(
-                goal=plan.current_task.instruction,
-                context=hist_info,
-            )
-
-            tool_config = create_func_config(CODE_GENERATOR_WITH_TOOLS)
-            logger.info(f"prompt is: {prompt}")
-            rsp = await self.llm.aask_code(prompt, **tool_config)
-            logger.info(f"rsp is: {rsp}")
-            return rsp["code"]
+        tool_config = create_func_config(CODE_GENERATOR_WITH_TOOLS)
+        rsp = await self.llm.aask_code(prompt, **tool_config)
+        return rsp['code']
diff --git a/metagpt/prompts/ml_engineer.py b/metagpt/prompts/ml_engineer.py
index b68dadc9a..88cebf68a 100644
--- a/metagpt/prompts/ml_engineer.py
+++ b/metagpt/prompts/ml_engineer.py
@@ -8,19 +8,22 @@ GEN_DATA_DESC_PROMPT = """
 Here is the head 5 rows of the dataset:
 {data_head}
 
-Please provide a brief one-sentence background of the dataset, and concise descriptions for each column. Keep descriptions short yet informative.
+Please provide a brief one-sentence background of the dataset, and concise meaning for each column. Keep descriptions short.
 
 Output the information in a JSON format, as shown in this example:
 ```json
 {
     "data_desc": "Brief dataset background.",
     "column_desc": {
-        "column_name1": "Description of the first column.",
-        "column_name2": "Description of the second column.",
+        "column_name1": "Abstract meaning of the first column.",
+        "column_name2": "Abstract meaning of the second column.",
         ...
     }
 }
 ```
+
+# Constraints:
+- Don't contain specific values or examples found in the data column.
 """
 
 ASSIGN_TASK_TYPE_PROMPT = """
@@ -53,19 +56,22 @@ ASSIGN_TASK_TYPE = {
 }
 
 TOOL_RECOMMENDATION_PROMPT = """
-Your are a tool recommender, the main goal is to recommend suitable tools for current task before coding. A tool means a function that can be used to help you solve the task.
+## User Requirement:
+{current_task}
 
-## List of Available Tools:
-{available_tools}
-
-This is a task guide for the current task, including detailed code steps. You can refer to it when recommending tools.
+## Task
+Recommend up to five tools from 'Available Tools' that can help solve the 'User Requirement'. 
+This is a detailed code steps for current task. You can refer to it when recommending tools.
 {code_steps}
 
+## Available Tools:
+{available_tools}
+
 ## Tool Selection and Instructions:
-- For the task, choose up to five tools that are most likely to be useful in solving the task.
+- Select tools most relevant to completing the 'User Requirement'.
 - If you believe that no tools are suitable, indicate with an empty list.
 - Only list the names of the tools, not the full schema of each tool.
-- The result should only contain tool names that are in the list of available tools.
+- Ensure selected tools are listed in 'Available Tools'.
 """
 
 SELECT_FUNCTION_TOOLS = {
@@ -149,6 +155,34 @@ Finish your coding tasks as a helpful programmer based on the tools.
 
 """
 
+TOOL_USAGE_PROMPT = """
+## Target
+{goal}
+
+## History Info
+{context}
+
+## Available Tools:
+Each function is described in JSON format, including the function name and parameters. {output_desc}
+{function_catalog}
+
+When you call a function above, you should import the function from `{module_name}` first, e.g.:
+```python
+from metagpt.tools.functions.libs.data_preprocess import fill_missing_value
+```end
+
+## Your Output Format:
+Generate the complete code for this task:
+```python
+# Tools used: [function names or 'none']
+<your code for the current task, without any comments>
+```end
+
+## Attention:
+Make sure use the columns from the dataset columns
+Finish your coding tasks as a helpful programmer based on the tools.
+"""
+
 TOO_ORGANIZATION_PROMPT = """
 The previous conversation has provided all tasks step-by-step for the use goal and their statuses. 
 Now, begin writing code for the current task. This code should writen strictly on the basis of all previous completed tasks code, not a standalone code. And avoid writing duplicate code that has already been written in previous tasks, such as repeated import of packages, reading data, etc.
@@ -197,6 +231,66 @@ The current task is about feature engineering. when performing it, please adhere
 - Before generating a new feature, ensure the used features are already processed and ready to use.
 """
 
+DATA_PROCESS_PROMPT = """
+# Background
+As a data scientist, you need to help user to achieve the goal [{user_requirement}] step-by-step in an continuous Jupyter notebook.
+
+## Done Tasks
+```python
+{history_code}
+```end
+
+## Current Task
+{current_task}
+
+# Latest Data Info
+Latest data info after previous tasks:
+{column_info}
+
+# Task
+Write a Python function for 'Current Task'. Start by copying the input DataFrame. Avoid duplicating code from 'Done Tasks'.
+Specifically, {special_prompt}
+
+# Code Steps:
+Follow steps below when you writing code if it's convenient.
+{code_steps}
+
+# Capabilities
+- You can utilize pre-defined tools in any code lines from 'Available Tools' in the form of python functions.
+- You can freely combine the use of any other public packages, like sklearn, numpy, pandas, etc..
+- You can do anything about data preprocessing, feature engineering, model training, etc..
+
+# Available Tools:
+Each function tool is described in JSON format. {output_desc}
+When you call a function below, import the function from `{module_name}` first.
+{function_catalog}
+
+# Output Example:
+when current task is "fill missing value and handle outliers", the output code be like:
+```python
+from metagpt.tools.functions.libs.data_preprocess import fill_missing_value
+
+def function_name(df):
+    df_processed = df.copy()
+    num_cols = df_processed.select_dtypes(include='number').columns.tolist()
+    df_processed = fill_missing_value(df_processed, num_cols, 'mean')
+    
+    for col in num_cols:
+        low, high = df_processed[col].quantile([0.01, 0.99])
+        df_processed[col] = df_processed[col].clip(low, high)
+    return df_processed
+
+df_processed = function_name(df)
+print(df_processed.info())
+```end
+
+# Constraints:
+- Ensure the output new code is executable in the same Jupyter notebook with previous tasks code have been executed.
+- Prioritize using pre-defined tools for the same functionality.
+- Return DataFrame should always be named `df_processed`, while the input DataFrame should based on the done tasks' output DataFrame.
+- Limit to one print statement for the output DataFrame's info.
+"""
+
 MODEL_TRAIN_PROMPT = """
 The current task is about training a model, please ensure high 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.
@@ -204,9 +298,9 @@ The current task is about training a model, please ensure high performance:
 - Use the data from previous task result directly, do not mock or reload data yourself.
 """
 
-DATA_PREPROCESS_OUTPUT_DESC = "Please note that all functions uniformly output a processed pandas.DataFrame, facilitating seamless integration into the broader workflow."
+DATA_PREPROCESS_OUTPUT_DESC = "Please note that all functions output a updated pandas.DataFrame after data preprocessing."
 
-FEATURE_ENGINEERING_OUTPUT_DESC = "Please note that all functions uniformly output updated pandas.DataFrame with feature engineering applied."
+FEATURE_ENGINEERING_OUTPUT_DESC = "Please note that all functions output a updated pandas.DataFrame with new features added or existing features modified."
 
 CLASSIFICATION_MODEL_OUTPUT_DESC = ""
 
diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index deb76f0a9..4ad24df52 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -1,21 +1,21 @@
-import glob
 import json
+import re
 from typing import List
 
 import fire
 import pandas as pd
-import re
 
 from metagpt.actions import Action
 from metagpt.actions.execute_code import ExecutePyCode
 from metagpt.actions.write_analysis_code import WriteCodeByGenerate, WriteCodeWithTools
+from metagpt.actions.write_code_steps import WriteCodeSteps
 from metagpt.actions.write_plan import WritePlan
+from metagpt.const import DATA_PATH
 from metagpt.logs import logger
 from metagpt.prompts.ml_engineer import GEN_DATA_DESC_PROMPT
 from metagpt.roles import Role
 from metagpt.schema import Message, Plan
 from metagpt.utils.common import CodeParser
-from metagpt.actions.write_code_steps import WriteCodeSteps
 
 STRUCTURAL_CONTEXT = """
 ## User Requirement
@@ -70,32 +70,16 @@ def read_data(file: str) -> pd.DataFrame:
     return df
 
 
-def get_samples(df: pd.DataFrame) -> str:
+def get_column_info(df: pd.DataFrame) -> str:
     data = []
-
-    if len(df) > 5:
-        df_ = df.sample(5, random_state=0)
-    else:
-        df_ = df
-
-    for i in list(df_):
+    for i in df.columns:
         nan_freq = float("%.2g" % (df[i].isna().mean() * 100))
         n_unique = df[i].nunique()
-        s = df_[i].tolist()
+        data.append([i, df[i].dtype, nan_freq, n_unique])
 
-        if str(df[i].dtype) == "float64":
-            s = [round(sample, 2) if not pd.isna(sample) else None for sample in s]
-
-        data.append([df_[i].name, df[i].dtype, nan_freq, n_unique, s])
     samples = pd.DataFrame(
         data,
-        columns=[
-            "Column_name",
-            "Data_type",
-            "NaN_Frequency(%)",
-            "N_unique",
-            "Samples",
-        ],
+        columns=["Column_name", "Data_type", "NaN_Frequency(%)", "N_unique"],
     )
     return samples.to_string(index=False)
 
@@ -124,20 +108,19 @@ class AskReview(Action):
 
 
 class GenerateDataDesc(Action):
-    async def run(self, files: list) -> dict:
+    async def run(self, file: str) -> dict:
         data_desc = {}
-        for file in files:
-            df = read_data(file)
-            file_name = file.split("/")[-1]
-            data_head = df.head().to_dict(orient="list")
-            data_head = json.dumps(data_head, indent=4, ensure_ascii=False)
-            prompt = GEN_DATA_DESC_PROMPT.replace("{data_head}", data_head)
-            rsp = await self._aask(prompt)
-            rsp = CodeParser.parse_code(block=None, text=rsp)
-            data_desc[file_name] = {}
-            data_desc[file_name]["path"] = file
-            data_desc[file_name]["description"] = rsp
-            data_desc[file_name]["column_info"] = get_samples(df)
+        df = read_data(file)
+        data_head = df.head().to_dict(orient="list")
+        data_head = json.dumps(data_head, indent=4, ensure_ascii=False)
+        prompt = GEN_DATA_DESC_PROMPT.replace("{data_head}", data_head)
+        rsp = await self._aask(prompt)
+        rsp = CodeParser.parse_code(block=None, text=rsp)
+        rsp = json.loads(rsp)
+        data_desc["path"] = file
+        data_desc["data_desc"] = rsp["data_desc"]
+        data_desc["column_desc"] = rsp["column_desc"]
+        data_desc["column_info"] = get_column_info(df)
         return data_desc
 
 
@@ -159,7 +142,6 @@ class MLEngineer(Role):
         if self.data_path:
             self.data_desc = await self._generate_data_desc()
 
-
         # create initial plan and update until confirmation
         await self._update_plan()
 
@@ -181,13 +163,14 @@ class MLEngineer(Role):
                 self.plan.finish_current_task()
                 self.working_memory.clear()
 
+                if "print(df_processed.info())" in code:
+                    self.data_desc["column_info"] = result
             else:
                 # update plan according to user's feedback and to take on changed tasks
                 await self._update_plan()
 
     async def _generate_data_desc(self):
-        files = glob.glob(self.data_path + "/*.csv")
-        data_desc = await GenerateDataDesc().run(files=files)
+        data_desc = await GenerateDataDesc().run(self.data_path)
         return data_desc
 
     async def _write_and_exec_code(self, max_retry: int = 3):
@@ -201,9 +184,11 @@ class MLEngineer(Role):
         success = False
         while not success and counter < max_retry:
             context = self.get_useful_memories()
-            # breakpoint()
 
-            column_names_dict = {key: value["column_info"] for key,value in self.data_desc.items()}
+            # print("*" * 10)
+            # print(context)
+            # print("*" * 10)
+            # breakpoint()
 
             if not self.use_tools or self.plan.current_task.task_type == "other":
                 logger.info("Write code with pure generation")
@@ -214,9 +199,9 @@ class MLEngineer(Role):
                 cause_by = WriteCodeByGenerate
             else:
                 logger.info("Write code with tools")
-
+                column_info = self.data_desc['column_info']
                 code = await WriteCodeWithTools().run(
-                    context=context, plan=self.plan, code_steps=code_steps, **{"column_names": column_names_dict}
+                    context=context, plan=self.plan, code_steps=code_steps, column_info=column_info
                 )
                 cause_by = WriteCodeWithTools
 
@@ -296,10 +281,8 @@ if __name__ == "__main__":
     # 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"
     # requirement = "Run EDA and visualization on this dataset, train a model to predict survival, report metrics on validation set (20%), dataset: workspace/titanic/train.csv"
 
-    from metagpt.const import DATA_PATH
-
     requirement = "Perform data analysis on the provided data. Train a model to predict the target variable Survived. Include data preprocessing, feature engineering, and modeling in your pipeline. The metric is accuracy."
-    data_path = f"{DATA_PATH}/titanic"
+    data_path = f"{DATA_PATH}/titanic.csv"
 
     async def main(requirement: str = requirement, auto_run: bool = True, data_path: str = data_path):
         role = MLEngineer(goal=requirement, auto_run=auto_run, data_path=data_path)
diff --git a/metagpt/utils/common.py b/metagpt/utils/common.py
index 8f8edbc6d..168966ef7 100644
--- a/metagpt/utils/common.py
+++ b/metagpt/utils/common.py
@@ -315,3 +315,17 @@ def create_func_config(func_schema: dict) -> dict:
         "tools": tools,
         "tool_choice": tool_choice,
     }
+
+
+def remove_comments(code_str):
+    """Remove comments from code."""
+    pattern = r"(\".*?\"|\'.*?\')|(\#.*?$)"
+    def replace_func(match):
+        if match.group(2) is not None:
+            return ""
+        else:
+            return match.group(1)
+
+    clean_code = re.sub(pattern, replace_func, code_str, flags=re.MULTILINE)
+    clean_code = os.linesep.join([s.rstrip() for s in clean_code.splitlines() if s.strip()])
+    return clean_code

From 13e2b058125f45f43ff998483a7e175ddaeb5883 Mon Sep 17 00:00:00 2001
From: stellahsr <stellahsr@126.com>
Date: Fri, 8 Dec 2023 11:01:13 +0800
Subject: [PATCH 11/49] add reflection change write code internal ppl

---
 metagpt/actions/debug_code.py          | 111 +++++++++++++++++++++++++
 metagpt/actions/write_analysis_code.py |  65 ++++++++-------
 metagpt/prompts/ml_engineer.py         |  25 ++++--
 metagpt/roles/ml_engineer.py           |  68 +++++++++++----
 4 files changed, 219 insertions(+), 50 deletions(-)
 create mode 100644 metagpt/actions/debug_code.py

diff --git a/metagpt/actions/debug_code.py b/metagpt/actions/debug_code.py
new file mode 100644
index 000000000..3d460fa40
--- /dev/null
+++ b/metagpt/actions/debug_code.py
@@ -0,0 +1,111 @@
+from typing import Dict, List, Union, Tuple, Optional, Any
+
+from metagpt.actions import Action
+from metagpt.logs import logger
+from metagpt.schema import Message, Plan
+from metagpt.utils.common import CodeParser
+from metagpt.actions.write_analysis_code import BaseWriteAnalysisCode
+
+DEBUG_REFLECTION_EXAMPLE = '''Example 1:
+                           [previous impl]:
+                           ```python
+                           def add(a: int, b: int) -> int:
+                               """
+                               Given integers a and b, return the total value of a and b.
+                               """
+                               return a - b
+                           ```
+
+                           [runtime Error]:
+                           Tested passed:
+
+                           Tests failed:
+                           assert add(1, 2) == 3 # output: -1
+                           assert add(1, 2) == 4 # output: -1
+
+                           [reflection on previous impl]:
+                           The implementation failed the test cases where the input integers are 1 and 2. The issue arises because the code does not add the two integers together, but instead subtracts the second integer from the first. To fix this issue, we should change the operator from `-` to `+` in the return statement. This will ensure that the function returns the correct output for the given input.
+
+                           [improved impl]:
+                           ```python
+                           def add(a: int, b: int) -> int:
+                               """
+                               Given integers a and b, return the total value of a and b.
+                               """
+                               return a + b
+                           ```
+                           '''
+
+REFLECTION_PROMPT = """
+                       Here is an example for you.
+                       {debug_example}
+                       [requirement]
+                       {goal}
+                       [previous impl]
+                       {code}
+                       [runtime Error]
+                       {runtime_result}
+
+                        Analysis the error step by step, provide me improve method. Do not repeat [previous impl]
+                        [reflection on previous impl]:
+                        xxx
+
+                       """
+
+
+def message_to_str(message: Message) -> str:
+    return f"{message.role}: {message.content}"
+
+
+def messages_to_str(messages: List[Message]) -> str:
+    return "\n".join([message_to_str(message) for message in messages])
+
+
+class DebugCode(BaseWriteAnalysisCode):
+    name: str = "debugcode"
+    context: Optional[str] = None
+    llm: None
+
+    def __init__(self, **kwargs: Any):
+        super().__init__(**kwargs)
+
+    async def run_reflection(self, plan, code, runtime_result) -> str:
+        info = []
+        reflection_prompt = REFLECTION_PROMPT.format(debug_example=DEBUG_REFLECTION_EXAMPLE,
+                                                     goal=plan.goal,
+                                                     code=code,
+                                                     runtime_result=runtime_result
+                                                     )
+        system_prompt = "You are an AI Python assistant. You will be given your previous implementation of a function, runtime error results, and a hint to change the implementation appropriately. Write your full implementation "
+        info.append(Message(role="system", content=system_prompt))
+        info.append(Message(role="assistant", content=reflection_prompt))
+
+        msg = messages_to_str(info)
+        resp = await self.llm.aask(msg=msg)
+        logger.info(f"reflection is {resp}")
+        return resp
+
+    async def rewrite_code(self, reflection: str = "") -> str:
+        """
+        根据reflection重写代码
+        """
+        info = []
+        info.append(Message(role="assistant", content=f"[reflection]: \n {reflection}"))
+        info.append(Message(role="user", content=f"[improved impl]:\n Return in Python block"))
+        msg = messages_to_str(info)
+        resp = await self.llm.aask(msg=msg)
+        logger.info(f"improve code is {resp}")
+        improv_code = CodeParser.parse_code(block=None, text=resp)
+        return improv_code
+
+    async def run(self,
+                  plan: Plan = None,
+                  code: str = "",
+                  runtime_result: str = "") -> str:
+        """
+        根据当前运行代码和报错信息进行reflection和纠错
+        """
+        reflection = await self.run_reflection(plan, code, runtime_result)
+        # 根据reflection结果重写代码
+        improv_code = await self.rewrite_code(reflection)
+        return improv_code
diff --git a/metagpt/actions/write_analysis_code.py b/metagpt/actions/write_analysis_code.py
index 957d35f7e..777064f93 100644
--- a/metagpt/actions/write_analysis_code.py
+++ b/metagpt/actions/write_analysis_code.py
@@ -4,7 +4,7 @@
 @Author  :   orange-crow
 @File    :   write_code_v2.py
 """
-from typing import Dict, List, Union, Tuple
+from typing import Dict, List, Union, Tuple, Optional, Any
 
 from metagpt.actions import Action
 from metagpt.logs import logger
@@ -12,7 +12,7 @@ from metagpt.prompts.ml_engineer import (
     TOOL_RECOMMENDATION_PROMPT,
     SELECT_FUNCTION_TOOLS,
     CODE_GENERATOR_WITH_TOOLS,
-    TOO_ORGANIZATION_PROMPT,
+    TOOL_ORGANIZATION_PROMPT,
     ML_SPECIFIC_PROMPT,
     ML_MODULE_MAP,
     TOOL_OUTPUT_DESC,
@@ -22,10 +22,13 @@ from metagpt.schema import Message, Plan
 from metagpt.tools.functions import registry
 from metagpt.utils.common import create_func_config
 from metagpt.prompts.ml_engineer import GEN_DATA_DESC_PROMPT, GENERATE_CODE_PROMPT
-from metagpt.utils.common import CodeParser
+
 from metagpt.actions.execute_code import ExecutePyCode
 
 
+
+
+
 class BaseWriteAnalysisCode(Action):
     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."""  # 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!"""
@@ -80,6 +83,8 @@ class BaseWriteAnalysisCode(Action):
         """
 
 
+
+
 class WriteCodeByGenerate(BaseWriteAnalysisCode):
     """Write code fully by generation"""
 
@@ -153,7 +158,6 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
         recommend_tools = rsp["recommend_tools"]
         return recommend_tools
 
-
     async def run(
             self,
             context: List[Message],
@@ -164,25 +168,23 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
         task_type = plan.current_task.task_type
         logger.info(f"task_type is: {task_type}")
         available_tools = registry.get_all_schema_by_module(task_type)
+        special_prompt = ML_SPECIFIC_PROMPT.get(task_type, "")
 
-        # special_prompt = ML_SPECIFIC_PROMPT.get(task_type, "")
-
+        column_names = kwargs.get("column_names", {})
         finished_tasks = plan.get_finished_tasks()
         code_context = [task.code for task in finished_tasks]
 
         code_context = "\n\n".join(code_context)
 
-        ### add runtime info
-        result, success = await self.execute_code.run(code_context)
-        logger.info(result)
-
         if len(available_tools) > 0:
             available_tools = [
                 {k: tool[k] for k in ["name", "description"] if k in tool}
                 for tool in available_tools
             ]
 
-            final_code = code_context
+            final_code = {}
+            new_code = ""
+            code_steps_dict = eval(code_steps)
 
             recommend_tools = await self._tool_recommendation(context, code_steps, available_tools)
             tool_catalog = self._parse_recommend_tools(task_type, recommend_tools)
@@ -191,33 +193,40 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
             module_name = ML_MODULE_MAP[task_type]
             output_desc = TOOL_OUTPUT_DESC.get(task_type, "")
 
-            hist_info = f"Previous finished code is \n\n ```Python {final_code} ``` \n\n " \
-                        f"Runtime result is {result} \n\n"
 
-            prompt = TOOL_USAGE_PROMPT.format(
-                goal=plan.current_task.instruction,
-                context=hist_info,
-                code_steps=code_steps,
-                module_name=module_name,
-                output_desc=output_desc,
-                function_catalog=tool_catalog,
-            )
+            for idx, tool in enumerate(recommend_tools):
+                hist_info = f"Previous finished code is \n\n ```Python {code_context} ``` \n\n "
 
-            tool_config = create_func_config(CODE_GENERATOR_WITH_TOOLS)
+                prompt = TOOL_USAGE_PROMPT.format(
+                    goal=plan.current_task.instruction,
+                    context=hist_info,
+                    code_steps=code_steps,
+                    column_names=column_names,
+                    special_prompt=special_prompt,
+                    module_name=module_name,
+                    output_desc=output_desc,
+                    function_catalog=tool_catalog[idx],
+                )
 
-            rsp = await self.llm.aask_code(prompt, **tool_config)
-            logger.info(f"rsp is: {rsp}")
-            final_code = final_code + "\n\n" + rsp["code"]
+                tool_config = create_func_config(CODE_GENERATOR_WITH_TOOLS)
 
-            return final_code
+                rsp = await self.llm.aask_code(prompt, **tool_config)
+                logger.info(f"rsp is: {rsp}")
+                # final_code = final_code + "\n\n" + rsp["code"]
+                # final_code[key] = rsp["code"]
+                new_code = new_code + "\n\n" + rsp["code"]
+                code_context = code_context + "\n\n" + rsp["code"]
+            return new_code
 
         else:
-            hist_info = f"Previous finished code is \n\n ```Python {code_context} ``` \n\n " \
-                     f"runtime result is {result} \n\n"
+            hist_info = f"Previous finished code is \n\n ```Python {code_context} ``` \n\n "
 
             prompt = GENERATE_CODE_PROMPT.format(
                 goal=plan.current_task.instruction,
                 context=hist_info,
+                code_steps=code_steps,
+                special_prompt=special_prompt,
+                # column_names=column_names
             )
 
             tool_config = create_func_config(CODE_GENERATOR_WITH_TOOLS)
diff --git a/metagpt/prompts/ml_engineer.py b/metagpt/prompts/ml_engineer.py
index b68dadc9a..9a234478c 100644
--- a/metagpt/prompts/ml_engineer.py
+++ b/metagpt/prompts/ml_engineer.py
@@ -105,9 +105,15 @@ TOOL_USAGE_PROMPT = """
 ## Target
 {goal}
 
+Specifically, {special_prompt}
+
 ## History Info
 {context}
 
+## Code Steps for Current Task:
+Follow steps below when you writing code if it's convenient.
+{code_steps}
+
 ## Available Tools:
 Each function is described in JSON format, including the function name and parameters. {output_desc}
 {function_catalog}
@@ -125,7 +131,7 @@ Generate the complete code for this task:
 ```end
 
 ## Attention:
-Make sure use the columns from the dataset columns
+Make sure use the columns from the dataset columns: {column_names}
 Finish your coding tasks as a helpful programmer based on the tools.
 
 """
@@ -133,23 +139,30 @@ GENERATE_CODE_PROMPT = """
 ## Target
 {goal}
 
+Specifically, {special_prompt}
+
+
 ## History Info
 {context}
 
+## Code Steps for Current Task:
+Follow steps below when you writing code if it's convenient.
+{code_steps}
+
 ## Your Output Format:
 Generate the complete code for this task:
 ```python
-# Tools used: [function names or 'none']
-<your code for the current task>
-```end
+import pandas as pd
+
+```
 
 ## Attention:
 Make sure use the columns from the dataset columns
-Finish your coding tasks as a helpful programmer based on the tools.
+Finish your coding tasks as a helpful programmer based on the code.
 
 """
 
-TOO_ORGANIZATION_PROMPT = """
+TOOL_ORGANIZATION_PROMPT = """
 The previous conversation has provided all tasks step-by-step for the use goal and their statuses. 
 Now, begin writing code for the current task. This code should writen strictly on the basis of all previous completed tasks code, not a standalone code. And avoid writing duplicate code that has already been written in previous tasks, such as repeated import of packages, reading data, etc.
 Specifically, {special_prompt}
diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index deb76f0a9..b5904213c 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -16,6 +16,7 @@ from metagpt.roles import Role
 from metagpt.schema import Message, Plan
 from metagpt.utils.common import CodeParser
 from metagpt.actions.write_code_steps import WriteCodeSteps
+from metagpt.actions.debug_code import DebugCode
 
 STRUCTURAL_CONTEXT = """
 ## User Requirement
@@ -36,10 +37,13 @@ catboost
 """
 
 
+
+
+
 def truncate(result: str, keep_len: int = 1000) -> str:
     desc = "Truncated to show only the last 1000 characters\n"
     if result.startswith(desc):
-        result = result[-len(desc) :]
+        result = result[-len(desc):]
 
     if len(result) > keep_len:
         result = result[-keep_len:]
@@ -110,9 +114,9 @@ class AskReview(Action):
         logger.info("most recent context:")
         latest_action = context[-1].cause_by.__name__ if context[-1].cause_by else ""
         prompt = f"\nPlease review output from {latest_action}:\n" \
-            "If you want to change a task in the plan, say 'change task task_id, ... (things to change)'\n" \
-            "If you confirm the output and wish to continue with the current process, type CONFIRM\n" \
-            "If you want to terminate the process, type exit:\n"
+                 "If you want to change a task in the plan, say 'change task task_id, ... (things to change)'\n" \
+                 "If you confirm the output and wish to continue with the current process, type CONFIRM\n" \
+                 "If you want to terminate the process, type exit:\n"
         rsp = input(prompt)
 
         if rsp.lower() in ("exit"):
@@ -143,7 +147,7 @@ class GenerateDataDesc(Action):
 
 class MLEngineer(Role):
     def __init__(
-        self, name="ABC", profile="MLEngineer", goal="", auto_run: bool = False, data_path: str = None
+            self, name="ABC", profile="MLEngineer", goal="", auto_run: bool = False, data_path: str = None
     ):
         super().__init__(name=name, profile=profile, goal=goal)
         self._set_react_mode(react_mode="plan_and_act")
@@ -159,7 +163,6 @@ class MLEngineer(Role):
         if self.data_path:
             self.data_desc = await self._generate_data_desc()
 
-
         # create initial plan and update until confirmation
         await self._update_plan()
 
@@ -185,6 +188,15 @@ class MLEngineer(Role):
                 # update plan according to user's feedback and to take on changed tasks
                 await self._update_plan()
 
+
+        finished_tasks = self.plan.get_finished_tasks()
+        if len(finished_tasks) == len(self.plan.tasks):
+            code_context = [task.code for task in finished_tasks]
+            code_context = "\n\n".join(code_context)
+            result, success = await self.execute_code.run(code_context)
+            # truncated the result
+            print(truncate(result))
+
     async def _generate_data_desc(self):
         files = glob.glob(self.data_path + "/*.csv")
         data_desc = await GenerateDataDesc().run(files=files)
@@ -198,16 +210,29 @@ class MLEngineer(Role):
         )
 
         counter = 0
+        improve_code = ""
         success = False
+
+        finished_tasks = self.plan.get_finished_tasks()
+        code_context = [task.code for task in finished_tasks]
+        code_context = "\n\n".join(code_context)
+
         while not success and counter < max_retry:
-            context = self.get_useful_memories()
+            if counter == 0:
+                context = self.get_useful_memories()
+            else:
+                # improve_code = await DebugCode().run(plan=self.plan,
+                #                                      code= code_context + "\n\n" + code,
+                #                                      runtime_result=self.working_memory.get())
+                improve_code = ""
+
             # breakpoint()
 
-            column_names_dict = {key: value["column_info"] for key,value in self.data_desc.items()}
+            column_names_dict = {key: value["column_info"] for key, value in self.data_desc.items()}
 
             if not self.use_tools or self.plan.current_task.task_type == "other":
                 logger.info("Write code with pure generation")
-                # code = "print('abc')"
+
                 code = await WriteCodeByGenerate().run(
                     context=context, plan=self.plan, code_steps=code_steps, temperature=0.0
                 )
@@ -215,16 +240,24 @@ class MLEngineer(Role):
             else:
                 logger.info("Write code with tools")
 
-                code = await WriteCodeWithTools().run(
-                    context=context, plan=self.plan, code_steps=code_steps, **{"column_names": column_names_dict}
-                )
-                cause_by = WriteCodeWithTools
+                if improve_code!="":
+                    code = improve_code
+                    logger.info(f"new code {code}")
+                    cause_by = DebugCode
+                else:
+                    code = await WriteCodeWithTools().run(
+                        context=context, plan=self.plan, code_steps=code_steps, **{"column_names": column_names_dict}
+                    )
+
+                    cause_by = WriteCodeWithTools
 
             self.working_memory.add(
                 Message(content=code, role="assistant", cause_by=cause_by)
             )
 
-            result, success = await self.execute_code.run(code)
+            # debug on code, run on runcode with finished code and new_df
+            runcode = code_context + "\n\n" + code
+            result, success = await self.execute_code.run(runcode)
             # truncated the result
             print(truncate(result))
             # print(result)
@@ -266,6 +299,7 @@ class MLEngineer(Role):
         self.plan.add_tasks(tasks)
         self.working_memory.clear()
 
+
     def get_useful_memories(self) -> List[Message]:
         """find useful memories only to reduce context length and improve performance"""
         # TODO dataset description , code steps
@@ -298,11 +332,13 @@ if __name__ == "__main__":
 
     from metagpt.const import DATA_PATH
 
-    requirement = "Perform data analysis on the provided data. Train a model to predict the target variable Survived. Include data preprocessing, feature engineering, and modeling in your pipeline. The metric is accuracy."
+    # requirement = "Perform data analysis on the provided data. Train a model to predict the target variable Survived. Include data preprocessing, feature engineering, and modeling in your pipeline. The metric is accuracy."
     data_path = f"{DATA_PATH}/titanic"
+    requirement = f"This is a titanic passenger survival dataset, your goal is to predict passenger survival outcome. The target column is Survived. Perform data analysis, data preprocessing, feature engineering, and modeling to predict the target. Report accuracy on the eval data. Train data path: '{data_path}/split_train.csv', eval data path: '{data_path}/split_eval.csv'."
 
-    async def main(requirement: str = requirement, auto_run: bool = True, data_path: str = data_path):
+    async def main(requirement: str = requirement, auto_run: bool = True, data_path: str = ""):
         role = MLEngineer(goal=requirement, auto_run=auto_run, data_path=data_path)
         await role.run(requirement)
 
+
     fire.Fire(main)

From c3a06ad20365a89ce3209f33fa33c9ae7e98af67 Mon Sep 17 00:00:00 2001
From: stellahsr <stellahsr@126.com>
Date: Tue, 12 Dec 2023 10:10:07 +0800
Subject: [PATCH 12/49] =?UTF-8?q?=E6=9B=B4=E6=96=B0reflection=EF=BC=8C?=
 =?UTF-8?q?=E5=88=86=E5=BC=80=E5=8E=86=E5=8F=B2code=E5=92=8C=E5=B7=B2?=
 =?UTF-8?q?=E6=9C=89=E8=BF=90=E8=A1=8C=E7=BB=93=E6=9E=9C?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 metagpt/actions/debug_code.py | 39 ++++++++++++++++++++++++-----------
 1 file changed, 27 insertions(+), 12 deletions(-)

diff --git a/metagpt/actions/debug_code.py b/metagpt/actions/debug_code.py
index 3d460fa40..9efe93efc 100644
--- a/metagpt/actions/debug_code.py
+++ b/metagpt/actions/debug_code.py
@@ -41,6 +41,12 @@ REFLECTION_PROMPT = """
                        {debug_example}
                        [requirement]
                        {goal}
+                       [finished code]
+                       finished code are executable, and you should based on the code to continue your current code debug
+                       {finished_code}
+                       
+                       try to reuse the code here to understand the coding task.
+                       
                        [previous impl]
                        {code}
                        [runtime Error]
@@ -65,47 +71,56 @@ class DebugCode(BaseWriteAnalysisCode):
     name: str = "debugcode"
     context: Optional[str] = None
     llm: None
-
+    
     def __init__(self, **kwargs: Any):
         super().__init__(**kwargs)
-
-    async def run_reflection(self, plan, code, runtime_result) -> str:
+    
+    async def run_reflection(self, goal, finished_code, finished_code_result, code, runtime_result) -> str:
         info = []
+        finished_code_and_result = finished_code + "\n [finished results]\n\n" + finished_code_result
         reflection_prompt = REFLECTION_PROMPT.format(debug_example=DEBUG_REFLECTION_EXAMPLE,
-                                                     goal=plan.goal,
+                                                     goal=goal,
+                                                     finished_code=finished_code_and_result,
                                                      code=code,
                                                      runtime_result=runtime_result
                                                      )
         system_prompt = "You are an AI Python assistant. You will be given your previous implementation of a function, runtime error results, and a hint to change the implementation appropriately. Write your full implementation "
         info.append(Message(role="system", content=system_prompt))
         info.append(Message(role="assistant", content=reflection_prompt))
-
+        
         msg = messages_to_str(info)
         resp = await self.llm.aask(msg=msg)
         logger.info(f"reflection is {resp}")
         return resp
-
-    async def rewrite_code(self, reflection: str = "") -> str:
+    
+    async def rewrite_code(self, reflection: str = "", code_context: str = "") -> str:
         """
         根据reflection重写代码
         """
         info = []
-        info.append(Message(role="assistant", content=f"[reflection]: \n {reflection}"))
+        info.append(Message(role="assistant", content=f"[code context]:{code_context}"
+                                                      f"finished code are executable, and you should based on the code to continue your current code debug and improvement"
+                                                      f"[reflection]: \n {reflection}"))
         info.append(Message(role="user", content=f"[improved impl]:\n Return in Python block"))
         msg = messages_to_str(info)
         resp = await self.llm.aask(msg=msg)
         logger.info(f"improve code is {resp}")
         improv_code = CodeParser.parse_code(block=None, text=resp)
         return improv_code
-
+    
     async def run(self,
-                  plan: Plan = None,
+                  plan: str = "",
+                  finished_code: str = "",
+                  finished_code_result: str = "",
                   code: str = "",
                   runtime_result: str = "") -> str:
         """
         根据当前运行代码和报错信息进行reflection和纠错
         """
-        reflection = await self.run_reflection(plan, code, runtime_result)
+        reflection = await self.run_reflection(plan, finished_code=finished_code,
+                                               finished_code_result=finished_code_result,
+                                               code=code,
+                                               runtime_result=runtime_result)
         # 根据reflection结果重写代码
-        improv_code = await self.rewrite_code(reflection)
+        improv_code = await self.rewrite_code(reflection, code_context=finished_code)
         return improv_code

From 4f1aa0333ec9cae6bf69c711735794c8c6677693 Mon Sep 17 00:00:00 2001
From: stellahsr <stellahsr@126.com>
Date: Tue, 12 Dec 2023 10:10:19 +0800
Subject: [PATCH 13/49] =?UTF-8?q?=E5=A2=9E=E5=8A=A0retry?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 metagpt/provider/openai_api.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/metagpt/provider/openai_api.py b/metagpt/provider/openai_api.py
index 34e5693f8..d8d2e9a4f 100644
--- a/metagpt/provider/openai_api.py
+++ b/metagpt/provider/openai_api.py
@@ -15,6 +15,7 @@ from tenacity import (
     retry,
     retry_if_exception_type,
     stop_after_attempt,
+    wait_random_exponential,
     wait_fixed,
 )
 
@@ -259,7 +260,8 @@ class OpenAIGPTAPI(BaseGPTAPI, RateLimiter):
         rsp = self.llm.ChatCompletion.create(**self._func_configs(messages, **kwargs))
         self._update_costs(rsp.get("usage"))
         return rsp
-
+    
+    @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
     async def _achat_completion_function(self, messages: list[dict], **chat_configs) -> dict:
         rsp = await self.llm.ChatCompletion.acreate(**self._func_configs(messages, **chat_configs))
         self._update_costs(rsp.get("usage"))

From 4634415e378dc3b02659721ff97cd9b852d53cd4 Mon Sep 17 00:00:00 2001
From: stellahsr <stellahsr@126.com>
Date: Tue, 12 Dec 2023 10:15:21 +0800
Subject: [PATCH 14/49] =?UTF-8?q?=E5=8F=AA=E4=BD=BF=E7=94=A8=E5=BD=93?=
 =?UTF-8?q?=E5=89=8Dcode=E8=BF=90=E8=A1=8C=EF=BC=8C=E4=B8=8D=E8=BF=AD?=
 =?UTF-8?q?=E4=BB=A3=E5=8E=86=E5=8F=B2code=20=E5=88=86=E5=BC=80=E5=BD=93?=
 =?UTF-8?q?=E5=89=8Dcode=E5=92=8C=E5=8E=86=E5=8F=B2=E4=BB=A3=E7=A0=81?=
 =?UTF-8?q?=E7=BB=99reflection?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 metagpt/roles/ml_engineer.py | 114 +++++++++++++++++++----------------
 1 file changed, 62 insertions(+), 52 deletions(-)

diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index 1b191c8ba..45fe728dd 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -37,17 +37,14 @@ catboost
 """
 
 
-
-
-
 def truncate(result: str, keep_len: int = 1000) -> str:
     desc = "Truncated to show only the last 1000 characters\n"
     if result.startswith(desc):
         result = result[-len(desc):]
-
+    
     if len(result) > keep_len:
         result = result[-keep_len:]
-
+    
     if not result.startswith(desc):
         return desc + result
     return desc
@@ -80,7 +77,7 @@ def get_column_info(df: pd.DataFrame) -> str:
         nan_freq = float("%.2g" % (df[i].isna().mean() * 100))
         n_unique = df[i].nunique()
         data.append([i, df[i].dtype, nan_freq, n_unique])
-
+    
     samples = pd.DataFrame(
         data,
         columns=["Column_name", "Data_type", "NaN_Frequency(%)", "N_unique"],
@@ -94,7 +91,7 @@ class AskReview(Action):
         logger.info(
             "\n".join([f"{task.task_id}: {task.instruction}, is_finished: {task.is_finished}" for task in plan.tasks])
         )
-
+        
         logger.info("most recent context:")
         latest_action = context[-1].cause_by.__name__ if context[-1].cause_by else ""
         prompt = f"\nPlease review output from {latest_action}:\n" \
@@ -102,12 +99,12 @@ class AskReview(Action):
                  "If you confirm the output and wish to continue with the current process, type CONFIRM\n" \
                  "If you want to terminate the process, type exit:\n"
         rsp = input(prompt)
-
+        
         if rsp.lower() in ("exit"):
             exit()
-
+        
         confirmed = rsp.lower() in ("confirm", "yes", "y")
-
+        
         return rsp, confirmed
 
 
@@ -141,24 +138,24 @@ class MLEngineer(Role):
         self.auto_run = auto_run
         self.data_path = data_path
         self.data_desc = {}
-
+    
     async def _plan_and_act(self):
         if self.data_path:
             self.data_desc = await self._generate_data_desc()
-
+        
         # create initial plan and update until confirmation
         await self._update_plan()
-
+        
         while self.plan.current_task:
             task = self.plan.current_task
             logger.info(f"ready to take on task {task}")
-
+            
             # take on current task
             code, result, success, code_steps = await self._write_and_exec_code()
-
+            
             # ask for acceptance, users can other refuse and change tasks in the plan
             task_result_confirmed = await self._ask_review()
-
+            
             if success and task_result_confirmed:
                 # tick off this task and record progress
                 task.code = code
@@ -166,14 +163,13 @@ class MLEngineer(Role):
                 task.code_steps = code_steps
                 self.plan.finish_current_task()
                 self.working_memory.clear()
-
+                
                 if "print(df_processed.info())" in code:
                     self.data_desc["column_info"] = result
             else:
                 # update plan according to user's feedback and to take on changed tasks
                 await self._update_plan()
-
-
+        
         finished_tasks = self.plan.get_finished_tasks()
         if len(finished_tasks) == len(self.plan.tasks):
             code_context = [task.code for task in finished_tasks]
@@ -181,46 +177,51 @@ class MLEngineer(Role):
             result, success = await self.execute_code.run(code_context)
             # truncated the result
             print(truncate(result))
-
+    
     async def _generate_data_desc(self):
         data_desc = await GenerateDataDesc().run(self.data_path)
         return data_desc
-
+    
     async def _write_and_exec_code(self, max_retry: int = 3):
         code_steps = (
             await WriteCodeSteps().run(self.plan)
             if self.use_code_steps
             else ""
         )
-
+        
         counter = 0
         improve_code = ""
         success = False
-
+        
         finished_tasks = self.plan.get_finished_tasks()
         code_context = [task.code for task in finished_tasks]
+        code_result = [task.result for task in finished_tasks]
         code_context = "\n\n".join(code_context)
-
+        code_result = "\n\n".join(code_result)
+        
         while not success and counter < max_retry:
             if counter == 0:
                 context = self.get_useful_memories()
             else:
-                improve_code = await DebugCode().run(plan=self.plan,
-                                                     code= code_context + "\n\n" + code,
+                # context = self.get_useful_memories()
+                # logger.info(f"context {context}")
+                improve_code = await DebugCode().run(plan=self.plan.current_task.instruction,
+                                                     finished_code=code_context,
+                                                     finished_code_result=code_result,
+                                                     code=code,
                                                      runtime_result=self.working_memory.get())
-
-
+            
             if not self.use_tools or self.plan.current_task.task_type == "other":
                 logger.info("Write code with pure generation")
-
+                
                 code = await WriteCodeByGenerate().run(
                     context=context, plan=self.plan, code_steps=code_steps, temperature=0.0
                 )
                 cause_by = WriteCodeByGenerate
             else:
                 logger.info("Write code with tools")
-
-                if improve_code!="":
+                
+                if improve_code != "":
                     code = improve_code
                     logger.info(f"new code {code}")
                     cause_by = DebugCode
@@ -228,15 +229,17 @@ class MLEngineer(Role):
                     code = await WriteCodeWithTools().run(
                         context=context, plan=self.plan, code_steps=code_steps, **{"column_names": {}}
                     )
-
+                    
                     cause_by = WriteCodeWithTools
-
+            
             self.working_memory.add(
                 Message(content=code, role="assistant", cause_by=cause_by)
             )
-
+            
             # debug on code, run on runcode with finished code and new_df
-            runcode = code_context + "\n\n" + code
+            # runcode = code_context + "\n\n" + code
+            runcode = code
+            
             result, success = await self.execute_code.run(runcode)
             # truncated the result
             print(truncate(result))
@@ -244,16 +247,16 @@ class MLEngineer(Role):
             self.working_memory.add(
                 Message(content=truncate(remove_escape_and_color_codes(result)), role="user", cause_by=ExecutePyCode)
             )
-
+            
             if "!pip" in code:
-                success = False
+                 success = False
             # if not success:
             #     await self._ask_review()
-
+            
             counter += 1
-
+        
         return code, result, success, code_steps
-
+    
     async def _ask_review(self):
         if not self.auto_run:
             context = self.get_useful_memories()
@@ -262,9 +265,10 @@ class MLEngineer(Role):
                 self.working_memory.add(Message(content=review, role="user", cause_by=AskReview))
             return confirmed
         return True
-
+    
     async def _update_plan(self, max_tasks: int = 3):
         plan_confirmed = False
+        
         while not plan_confirmed:
             context = self.get_useful_memories()
             rsp = await WritePlan().run(
@@ -274,12 +278,17 @@ class MLEngineer(Role):
                 Message(content=rsp, role="assistant", cause_by=WritePlan)
             )
             plan_confirmed = await self._ask_review()
-
-        tasks = WritePlan.rsp_to_tasks(rsp)
+        
+        new_tasks = WritePlan.rsp_to_tasks(rsp)
+        logger.debug(len(self.plan.tasks))
+        logger.debug(len(new_tasks))
+        ## fixme: 能重复执行多轮重新plan，但应该有更优处理逻辑
+        ## fixme: do not overwrite original tasks
+        tasks = self.plan.tasks + new_tasks
+        
         self.plan.add_tasks(tasks)
         self.working_memory.clear()
-
-
+    
     def get_useful_memories(self) -> List[Message]:
         """find useful memories only to reduce context length and improve performance"""
         # TODO dataset description , code steps
@@ -295,9 +304,9 @@ class MLEngineer(Role):
             current_task=current_task
         )
         context_msg = [Message(content=context, role="user")]
-
+        
         return context_msg + self.working_memory.get()
-
+    
     @property
     def working_memory(self):
         return self._rc.memory
@@ -309,15 +318,16 @@ if __name__ == "__main__":
     # requirement = "Run data analysis on sklearn Wine recognition dataset, include a plot, and train a model to predict wine class (20% as validation), and show validation accuracy"
     # 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"
     # requirement = "Run EDA and visualization on this dataset, train a model to predict survival, report metrics on validation set (20%), dataset: workspace/titanic/train.csv"
-
+    
     # requirement = "Perform data analysis on the provided data. Train a model to predict the target variable Survived. Include data preprocessing, feature engineering, and modeling in your pipeline. The metric is accuracy."
-
+    
     data_path = f"{DATA_PATH}/titanic"
     requirement = f"This is a titanic passenger survival dataset, your goal is to predict passenger survival outcome. The target column is Survived. Perform data analysis, data preprocessing, feature engineering, and modeling to predict the target. Report accuracy on the eval data. Train data path: '{data_path}/split_train.csv', eval data path: '{data_path}/split_eval.csv'."
-
+    
+    
     async def main(requirement: str = requirement, auto_run: bool = True, data_path: str = ""):
         role = MLEngineer(goal=requirement, auto_run=auto_run, data_path=data_path)
         await role.run(requirement)
-
-
+    
+    
     fire.Fire(main)

From fd31cc065a74ce8b17765ab7b44ff51ce0adc833 Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Tue, 12 Dec 2023 10:30:05 +0800
Subject: [PATCH 15/49] save jupyter file

---
 metagpt/actions/execute_code.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/metagpt/actions/execute_code.py b/metagpt/actions/execute_code.py
index 981aa894c..6fd980494 100644
--- a/metagpt/actions/execute_code.py
+++ b/metagpt/actions/execute_code.py
@@ -156,6 +156,11 @@ class ExecutePyCode(ExecuteCode, Action):
 
         return code, language
 
+    def save_notebook(self, path: str):
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        nbformat.write(self.nb, path)
+
     async def run(self, code: Union[str, Dict, Message], language: str = "python") -> Tuple[str, bool]:
         code, language = self._process_code(code, language)
 

From 4f0d55656e17c2247b84d748f8cb0cc0ebba5176 Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Tue, 12 Dec 2023 10:56:05 +0800
Subject: [PATCH 16/49] update ml tool from Function to Class

---
 metagpt/tools/functions/libs/base.py          |  16 +
 .../tools/functions/libs/data_preprocess.py   | 248 ++++++----
 .../functions/libs/feature_engineering.py     | 427 +++++++++++-------
 3 files changed, 445 insertions(+), 246 deletions(-)
 create mode 100644 metagpt/tools/functions/libs/base.py

diff --git a/metagpt/tools/functions/libs/base.py b/metagpt/tools/functions/libs/base.py
new file mode 100644
index 000000000..c39adc66b
--- /dev/null
+++ b/metagpt/tools/functions/libs/base.py
@@ -0,0 +1,16 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# @Time    : 2023/12/10 20:12
+# @Author  : lidanyang
+# @File    : base
+# @Desc    :
+class MLProcess(object):
+    def fit(self, df):
+        raise NotImplementedError
+
+    def transform(self, df):
+        raise NotImplementedError
+
+    def fit_transform(self, df):
+        self.fit(df)
+        return self.transform(df)
diff --git a/metagpt/tools/functions/libs/data_preprocess.py b/metagpt/tools/functions/libs/data_preprocess.py
index 5579c5bd8..39474b0fd 100644
--- a/metagpt/tools/functions/libs/data_preprocess.py
+++ b/metagpt/tools/functions/libs/data_preprocess.py
@@ -1,6 +1,6 @@
 import numpy as np
 from sklearn.impute import SimpleImputer
-from sklearn.preprocessing import KBinsDiscretizer
+from sklearn.preprocessing import KBinsDiscretizer, LabelEncoder
 from sklearn.preprocessing import MaxAbsScaler
 from sklearn.preprocessing import MinMaxScaler
 from sklearn.preprocessing import OneHotEncoder
@@ -9,31 +9,52 @@ from sklearn.preprocessing import RobustScaler
 from sklearn.preprocessing import StandardScaler
 
 from metagpt.tools.functions import registry
+from metagpt.tools.functions.libs.base import MLProcess
 from metagpt.tools.functions.schemas.data_preprocess import *
 
 
-@registry.register("data_preprocess", FillMissingValue)
-def fill_missing_value(df: pd.DataFrame, features: list, strategy: str = 'mean', fill_value=None,):
-    df[features] = SimpleImputer(strategy=strategy, fill_value=fill_value).fit_transform(df[features])
-    return df
+class FillMissingValue(MLProcess):
+    def __init__(self, features: list, strategy: str = 'mean', fill_value=None,):
+        self.features = features
+        self.strategy = strategy
+        self.fill_value = fill_value
+        self.si = None
+
+    def fit(self, df: pd.DataFrame):
+        self.si = SimpleImputer(strategy=self.strategy, fill_value=self.fill_value)
+        self.si.fit(df[self.features])
+
+    def transform(self, df: pd.DataFrame):
+        df[self.features] = self.si.transform(df[self.features])
+        return df
 
 
-@registry.register("data_preprocess", SplitBins)
-def split_bins(df: pd.DataFrame, features: list, strategy: str = 'quantile',):
-    df[features] = KBinsDiscretizer(strategy=strategy, encode='ordinal').fit_transform(df[features])
-    return df
+class MinMaxScale(MLProcess):
+    def __init__(self, features: list,):
+        self.features = features
+        self.mms = None
+
+    def fit(self, df: pd.DataFrame):
+        self.mms = MinMaxScaler()
+        self.mms.fit(df[self.features])
+
+    def transform(self, df: pd.DataFrame):
+        df[self.features] = self.mms.transform(df[self.features])
+        return df
 
 
-@registry.register("data_preprocess", MinMaxScale)
-def min_max_scale(df: pd.DataFrame, features: list, ):
-    df[features] = MinMaxScaler().fit_transform(df[features])
-    return df
+class StandardScale(MLProcess):
+    def __init__(self, features: list,):
+        self.features = features
+        self.ss = None
 
+    def fit(self, df: pd.DataFrame):
+        self.ss = StandardScaler()
+        self.ss.fit(df[self.features])
 
-@registry.register("data_preprocess", StandardScale)
-def standard_scale(df: pd.DataFrame, features: list, ):
-    df[features] = StandardScaler().fit_transform(df[features])
-    return df
+    def transform(self, df: pd.DataFrame):
+        df[self.features] = self.ss.transform(df[self.features])
+        return df
 
 
 @registry.register("data_preprocess", LogTransform)
@@ -45,80 +66,145 @@ def log_transform(df: pd.DataFrame, features: list, ):
     return df
 
 
-@registry.register("data_preprocess", MaxAbsScale)
-def max_abs_scale(df: pd.DataFrame, features: list, ):
-    df[features] = MaxAbsScaler().fit_transform(df[features])
-    return df
+class MaxAbsScale(MLProcess):
+    def __init__(self, features: list,):
+        self.features = features
+        self.mas = None
+
+    def fit(self, df: pd.DataFrame):
+        self.mas = MaxAbsScaler()
+        self.mas.fit(df[self.features])
+
+    def transform(self, df: pd.DataFrame):
+        df[self.features] = self.mas.transform(df[self.features])
+        return df
 
 
-@registry.register("data_preprocess", RobustScale)
-def robust_scale(df: pd.DataFrame, features: list, ):
-    df[features] = RobustScaler().fit_transform(df[features])
-    return df
+class RobustScale(MLProcess):
+    def __init__(self, features: list,):
+        self.features = features
+        self.rs = None
+
+    def fit(self, df: pd.DataFrame):
+        self.rs = RobustScaler()
+        self.rs.fit(df[self.features])
+
+    def transform(self, df: pd.DataFrame):
+        df[self.features] = self.rs.transform(df[self.features])
+        return df
 
 
-@registry.register("data_preprocess", OrdinalEncode)
-def ordinal_encode(df: pd.DataFrame, features: list,):
-    df[features] = OrdinalEncoder().fit_transform(df[features])
-    return df
+class OrdinalEncode(MLProcess):
+    def __init__(self, features: list,):
+        self.features = features
+        self.oe = None
+
+    def fit(self, df: pd.DataFrame):
+        self.oe = OrdinalEncoder()
+        self.oe.fit(df[self.features])
+
+    def transform(self, df: pd.DataFrame):
+        df[self.features] = self.oe.transform(df[self.features])
+        return df
 
 
-@registry.register("data_preprocess", OneHotEncoding)
-def one_hot_encoding(df, cols):
-    enc = OneHotEncoder(handle_unknown="ignore", sparse=False)
-    ts_data = enc.fit_transform(df[cols])
-    new_columns = enc.get_feature_names_out(cols)
-    ts_data = pd.DataFrame(ts_data, columns=new_columns, index=df.index)
-    df.drop(cols, axis=1, inplace=True)
-    df = pd.concat([df, ts_data], axis=1)
-    return df
+class OneHotEncode(MLProcess):
+    def __init__(self, features: list,):
+        self.features = features
+        self.ohe = None
+
+    def fit(self, df: pd.DataFrame):
+        self.ohe = OneHotEncoder(handle_unknown="ignore", sparse=False)
+        self.ohe.fit(df[self.features])
+
+    def transform(self, df: pd.DataFrame):
+        ts_data = self.ohe.transform(df[self.features])
+        new_columns = self.ohe.get_feature_names_out(self.features)
+        ts_data = pd.DataFrame(ts_data, columns=new_columns, index=df.index)
+        df.drop(self.features, axis=1, inplace=True)
+        df = pd.concat([df, ts_data], axis=1)
+        return df
 
 
-if __name__ == '__main__':
-    def run():
-        V = {
-            'a': [-1, 2, 3, 6, 5, 4],
-            'b': [1.1, 2.2, 3.3, 6.6, 5.5, 4.4],
-            'c': ['aa', 'bb', 'cc', 'dd', 'ee', 'ff'],
-            'd': [1, None, 3, None, 5, 4],
-            'e': [1.1, np.NAN, 3.3, None, 5.5, 4.4],
-            'f': ['aa', np.NAN, 'cc', None, '', 'ff'],
+class LabelEncode(MLProcess):
+    def __init__(self, features: list,):
+        self.features = features
+        self.le_encoders = []
 
-        }
+    def fit(self, df: pd.DataFrame):
+        for col in self.features:
+            le = LabelEncoder().fit(df[col].astype(str).unique().tolist() + ['unknown'])
+            self.le_encoders.append(le)
 
-        df = pd.DataFrame(V)
-        print(df.dtypes)
+    def transform(self, df: pd.DataFrame):
+        for i in range(len(self.features)):
+            data_list = df[self.features[i]].astype(str).tolist()
+            for unique_item in np.unique(df[self.features[i]].astype(str)):
+                if unique_item not in self.le_encoders[i].classes_:
+                    data_list = ['unknown' if x == unique_item else x for x in data_list]
+            df[self.features[i]] = self.le_encoders[i].transform(data_list)
+        return df
 
-        numeric_features = ['a', 'b', 'd', 'e']
-        numeric_features_wo_miss = ['a', 'b', ]
-        categorial_features = ['c', 'f']
 
-        df_ = fill_missing_value(df.copy(), numeric_features)
-        print(df_)
-        df_ = fill_missing_value(df.copy(), categorial_features, strategy='constant', fill_value='hehe')
-        print(df_)
+def get_column_info(df: pd.DataFrame) -> str:
+    data = []
+    for i in df.columns:
+        nan_freq = float("%.2g" % (df[i].isna().mean() * 100))
+        n_unique = df[i].nunique()
+        data.append([i, df[i].dtype, nan_freq, n_unique])
 
-        df_ = fill_missing_value(df.copy(), numeric_features, strategy='constant', fill_value=999)
-        print(df_)
-
-        # df_ = label_encode(df.copy(), numeric_features + categorial_features, )
-        # print(df_)
-
-        df_ = split_bins(df.copy(), numeric_features_wo_miss, strategy='quantile')
-        print(df_)
-
-        df_ = min_max_scale(df.copy(), numeric_features, )
-        print(df_)
-
-        df_ = standard_scale(df.copy(), numeric_features, )
-        print(df_)
-
-        df_ = log_transform(df.copy(), numeric_features, )
-        print(df_)
-
-        df_ = max_abs_scale(df.copy(), numeric_features, )
-        print(df_)
-
-        df_ = robust_scale(df.copy(), numeric_features, )
-        print(df_)
-    run()
\ No newline at end of file
+    samples = pd.DataFrame(
+        data,
+        columns=["Column_name", "Data_type", "NaN_Frequency(%)", "N_unique"],
+    )
+    return samples.to_string(index=False)
+#
+#
+# if __name__ == '__main__':
+#     def run():
+#         V = {
+#             'a': [-1, 2, 3, 6, 5, 4],
+#             'b': [1.1, 2.2, 3.3, 6.6, 5.5, 4.4],
+#             'c': ['aa', 'bb', 'cc', 'dd', 'ee', 'ff'],
+#             'd': [1, None, 3, None, 5, 4],
+#             'e': [1.1, np.NAN, 3.3, None, 5.5, 4.4],
+#             'f': ['aa', np.NAN, 'cc', None, '', 'ff'],
+#
+#         }
+#
+#         df = pd.DataFrame(V)
+#         print(df.dtypes)
+#
+#         numeric_features = ['a', 'b', 'd', 'e']
+#         numeric_features_wo_miss = ['a', 'b', ]
+#         categorial_features = ['c', 'f']
+#
+#         df_ = fill_missing_value(df.copy(), numeric_features)
+#         print(df_)
+#         df_ = fill_missing_value(df.copy(), categorial_features, strategy='constant', fill_value='hehe')
+#         print(df_)
+#
+#         df_ = fill_missing_value(df.copy(), numeric_features, strategy='constant', fill_value=999)
+#         print(df_)
+#
+#         # df_ = label_encode(df.copy(), numeric_features + categorial_features, )
+#         # print(df_)
+#
+#         df_ = split_bins(df.copy(), numeric_features_wo_miss, strategy='quantile')
+#         print(df_)
+#
+#         df_ = min_max_scale(df.copy(), numeric_features, )
+#         print(df_)
+#
+#         df_ = standard_scale(df.copy(), numeric_features, )
+#         print(df_)
+#
+#         df_ = log_transform(df.copy(), numeric_features, )
+#         print(df_)
+#
+#         df_ = max_abs_scale(df.copy(), numeric_features, )
+#         print(df_)
+#
+#         df_ = robust_scale(df.copy(), numeric_features, )
+#         print(df_)
+#     run()
\ No newline at end of file
diff --git a/metagpt/tools/functions/libs/feature_engineering.py b/metagpt/tools/functions/libs/feature_engineering.py
index 4780e4fa0..06a988d9a 100644
--- a/metagpt/tools/functions/libs/feature_engineering.py
+++ b/metagpt/tools/functions/libs/feature_engineering.py
@@ -3,188 +3,285 @@
 # @Time    : 2023/11/17 10:33
 # @Author  : lidanyang
 # @File    : feature_engineering.py
-# @Desc    : Feature Engineering Functions
+# @Desc    : Feature Engineering Tools
 import itertools
 
+import numpy as np
 from dateutil.relativedelta import relativedelta
+from joblib import Parallel, delayed
 from pandas.api.types import is_numeric_dtype
 from sklearn.model_selection import KFold
-from sklearn.preprocessing import PolynomialFeatures
+from sklearn.preprocessing import PolynomialFeatures, KBinsDiscretizer
 
-from metagpt.tools.functions import registry
+from metagpt.tools.functions.libs.base import MLProcess
 from metagpt.tools.functions.schemas.feature_engineering import *
 
 
-@registry.register("feature_engineering", PolynomialExpansion)
-def polynomial_expansion(df, cols, degree=2):
-    for col in cols:
-        if not is_numeric_dtype(df[col]):
-            raise ValueError(f"Column '{col}' must be numeric.")
+class PolynomialExpansion(MLProcess):
+    def __init__(self, cols: list, degree: int = 2):
+        self.cols = cols
+        self.degree = degree
+        self.poly = PolynomialFeatures(degree=degree, include_bias=False)
 
-    poly = PolynomialFeatures(degree=degree, include_bias=False)
-    ts_data = poly.fit_transform(df[cols].fillna(0))
-    new_columns = poly.get_feature_names_out(cols)
-    ts_data = pd.DataFrame(ts_data, columns=new_columns, index=df.index)
-    ts_data = ts_data.drop(cols, axis=1)
-    df = pd.concat([df, ts_data], axis=1)
-    return df
+    def fit(self, df: pd.DataFrame):
+        self.poly.fit(df[self.cols].fillna(0))
+
+    def transform(self, df: pd.DataFrame) -> pd.DataFrame:
+        ts_data = self.poly.transform(df[self.cols].fillna(0))
+        column_name = self.poly.get_feature_names_out(self.cols)
+        ts_data = pd.DataFrame(ts_data, index=df.index, columns=column_name)
+        df.drop(self.cols, axis=1, inplace=True)
+        df = pd.concat([df, ts_data], axis=1)
+        return df
 
 
-@registry.register("feature_engineering", FrequencyEncoding)
-def frequency_encoding(df, cols):
-    for col in cols:
-        encoder_dict = df[col].value_counts().to_dict()
-        df[f"{col}_cnt"] = df[col].map(encoder_dict)
-    return df
+class CatCount(MLProcess):
+    def __init__(self, col: str):
+        self.col = col
+        self.encoder_dict = None
+
+    def fit(self, df: pd.DataFrame):
+        self.encoder_dict = df[self.col].value_counts().to_dict()
+
+    def transform(self, df: pd.DataFrame) -> pd.DataFrame:
+        df[f"{self.col}_cnt"] = df[self.col].map(self.encoder_dict)
+        return df
 
 
-@registry.register("feature_engineering", TargetMeanEncoder)
-def target_mean_encoder(df, col, label):
-    encoder_dict = df.groupby(col)[label].mean().to_dict()
-    df[f"{col}_target_mean"] = df[col].map(encoder_dict)
-    return df
+class TargetMeanEncoder(MLProcess):
+    def __init__(self, col: str, label: str):
+        self.col = col
+        self.label = label
+        self.encoder_dict = None
+
+    def fit(self, df: pd.DataFrame):
+        self.encoder_dict = df.groupby(self.col)[self.label].mean().to_dict()
+
+    def transform(self, df: pd.DataFrame) -> pd.DataFrame:
+        df[f"{self.col}_target_mean"] = df[self.col].map(self.encoder_dict)
+        return df
 
 
-@registry.register("feature_engineering", KFoldTargetMeanEncoder)
-def k_fold_target_mean_encoder(df, col, label, n_splits=5, random_state=2021):
-    tmp = df.copy()
-    kf = KFold(n_splits=n_splits, shuffle=True, random_state=random_state)
+class KFoldTargetMeanEncoder(MLProcess):
+    def __init__(self, col: str, label: str, n_splits: int = 5, random_state: int = 2021):
+        self.col = col
+        self.label = label
+        self.n_splits = n_splits
+        self.random_state = random_state
+        self.encoder_dict = None
 
-    global_mean = tmp[label].mean()
-    col_name = f"{col}_kf_target_mean"
-    for trn_idx, val_idx in kf.split(tmp, tmp[label]):
-        _trn, _val = tmp.iloc[trn_idx], tmp.iloc[val_idx]
-        tmp.loc[tmp.index[val_idx], col_name] = _val[col].map(
-            _trn.groupby(col)[label].mean()
-        )
-    tmp[col_name].fillna(global_mean, inplace=True)
-    encoder_dict = tmp.groupby(col)[col_name].mean().to_dict()
-    df[f"{col}_kf_target_mean"] = df[col].map(encoder_dict)
-    return df
+    def fit(self, df: pd.DataFrame):
+        tmp = df.copy()
+        kf = KFold(n_splits=self.n_splits, shuffle=True, random_state=self.random_state)
 
-
-@registry.register("feature_engineering", CatCross)
-def cat_cross(df, cols, max_cat_num=100):
-    for col in cols:
-        if df[col].nunique() > max_cat_num:
-            cols.remove(col)
-
-    for col1, col2 in itertools.combinations(cols, 2):
-        cross_col = f"{col1}_cross_{col2}"
-        crossed = df[col1].astype(str) + "_" + df[col2].astype(str)
-        df[cross_col] = crossed.astype('category').cat.codes
-    return df
-
-
-@registry.register("feature_engineering", GroupStat)
-def group_stat(df, group_col, agg_col, agg_funcs):
-    group_df = df.groupby(group_col)[agg_col].agg(agg_funcs).reset_index()
-    group_df.columns = group_col + [
-        f"{agg_col}_{agg_func}_by_{group_col}" for agg_func in agg_funcs
-    ]
-    df = df.merge(group_df, on=group_col, how="left")
-    return df
-
-
-@registry.register("feature_engineering", ExtractTimeComps)
-def extract_time_comps(df, time_col, time_comps):
-    time_s = pd.to_datetime(df[time_col], errors="coerce")
-    time_comps_df = pd.DataFrame()
-
-    if "year" in time_comps:
-        time_comps_df["year"] = time_s.dt.year
-    if "month" in time_comps:
-        time_comps_df["month"] = time_s.dt.month
-    if "day" in time_comps:
-        time_comps_df["day"] = time_s.dt.day
-    if "hour" in time_comps:
-        time_comps_df["hour"] = time_s.dt.hour
-    if "dayofweek" in time_comps:
-        time_comps_df["dayofweek"] = time_s.dt.dayofweek + 1
-    if "is_weekend" in time_comps:
-        time_comps_df["is_weekend"] = time_s.dt.dayofweek.isin([5, 6]).astype(int)
-    df = pd.concat([df, time_comps_df], axis=1)
-    return df
-
-
-@registry.register("feature_engineering", FeShiftByTime)
-def fe_shift_by_time(df, time_col, group_col, shift_col, periods, freq):
-    df[time_col] = pd.to_datetime(df[time_col])
-
-    def shift_datetime(date, offset, unit):
-        if unit in ["year", "y", "Y"]:
-            return date + relativedelta(years=offset)
-        elif unit in ["month", "m", "M"]:
-            return date + relativedelta(months=offset)
-        elif unit in ["day", "d", "D"]:
-            return date + relativedelta(days=offset)
-        elif unit in ["week", "w", "W"]:
-            return date + relativedelta(weeks=offset)
-        elif unit in ["hour", "h", "H"]:
-            return date + relativedelta(hours=offset)
-        else:
-            return date
-
-    def shift_by_time_on_key(
-        inner_df, time_col, group_col, shift_col, offset, unit, col_name
-    ):
-        inner_df = inner_df.drop_duplicates()
-        inner_df[time_col] = inner_df[time_col].map(
-            lambda x: shift_datetime(x, offset, unit)
-        )
-        inner_df = inner_df.groupby([time_col, group_col], as_index=False)[
-            shift_col
-        ].mean()
-        inner_df.rename(columns={shift_col: col_name}, inplace=True)
-        return inner_df
-
-    shift_df = df[[time_col, group_col, shift_col]].copy()
-    for period in periods:
-        new_col_name = f"{group_col}_{shift_col}_lag_{period}_{freq}"
-        tmp = shift_by_time_on_key(
-            shift_df, time_col, group_col, shift_col, period, freq, new_col_name
-        )
-        df = df.merge(tmp, on=[time_col, group_col], how="left")
-
-    return df
-
-
-@registry.register("feature_engineering", FeRollingByTime)
-def fe_rolling_by_time(df, time_col, group_col, rolling_col, periods, freq, agg_funcs):
-    df[time_col] = pd.to_datetime(df[time_col])
-
-    def rolling_by_time_on_key(inner_df, offset, unit, agg_func, col_name):
-        time_freq = {
-            "Y": [365 * offset, "D"],
-            "M": [30 * offset, "D"],
-            "D": [offset, "D"],
-            "W": [7 * offset, "D"],
-            "H": [offset, "h"],
-        }
-
-        if agg_func not in ["mean", "std", "max", "min", "median", "sum", "count"]:
-            raise ValueError(f"Invalid agg function: {agg_func}")
-
-        rolling_feat = inner_df.rolling(
-            f"{time_freq[unit][0]}{time_freq[unit][1]}", closed="left"
-        )
-        rolling_feat = getattr(rolling_feat, agg_func)()
-        depth = df.columns.nlevels
-        rolling_feat = rolling_feat.stack(list(range(depth)))
-        rolling_feat.name = col_name
-        return rolling_feat
-
-    rolling_df = df[[time_col, group_col, rolling_col]].copy()
-    for period in periods:
-        for func in agg_funcs:
-            new_col_name = f"{group_col}_{rolling_col}_rolling_{period}_{freq}_{func}"
-            tmp = pd.pivot_table(
-                rolling_df,
-                index=time_col,
-                values=rolling_col,
-                columns=group_col,
+        global_mean = tmp[self.label].mean()
+        col_name = f"{self.col}_kf_target_mean"
+        for trn_idx, val_idx in kf.split(tmp, tmp[self.label]):
+            _trn, _val = tmp.iloc[trn_idx], tmp.iloc[val_idx]
+            tmp.loc[tmp.index[val_idx], col_name] = _val[self.col].map(
+                _trn.groupby(self.col)[self.label].mean()
             )
-            tmp = rolling_by_time_on_key(tmp, period, freq, func, new_col_name)
-            df = df.merge(tmp, on=[time_col, group_col], how="left")
+        tmp[col_name].fillna(global_mean, inplace=True)
+        self.encoder_dict = tmp.groupby(self.col)[col_name].mean().to_dict()
 
-    return df
+    def transform(self, df: pd.DataFrame) -> pd.DataFrame:
+        df[f"{self.col}_kf_target_mean"] = df[self.col].map(self.encoder_dict)
+        return df
+
+
+class CatCross(MLProcess):
+    def __init__(self, cols: list, max_cat_num: int = 100):
+        self.cols = cols
+        self.max_cat_num = max_cat_num
+        self.combs = []
+        self.combs_map = {}
+
+    @staticmethod
+    def cross_two(comb, df):
+        new_col = f'{comb[0]}_{comb[1]}'
+        new_col_combs = list(itertools.product(df[comb[0]].unique(), df[comb[1]].unique()))
+        ll = list(range(len(new_col_combs)))
+        comb_map = dict(zip(new_col_combs, ll))
+        return new_col, comb_map
+
+    def fit(self, df: pd.DataFrame):
+        for col in self.cols:
+            if df[col].nunique() > self.max_cat_num:
+                self.cols.remove(col)
+        self.combs = list(itertools.combinations(self.cols, 2))
+        res = Parallel(n_jobs=4, require='sharedmem')(
+            delayed(self.cross_two)(comb, df) for comb in self.combs)
+        self.combs_map = dict(res)
+
+    def transform(self, df: pd.DataFrame) -> pd.DataFrame:
+        for comb in self.combs:
+            new_col = f'{comb[0]}_{comb[1]}'
+            _map = self.combs_map[new_col]
+            df[new_col] = pd.Series(zip(df[comb[0]], df[comb[1]])).map(_map)
+            # set the unknown value to a new number
+            df[new_col].fillna(max(_map.values()) + 1, inplace=True)
+            df[new_col] = df[new_col].astype(int)
+        return df
+
+
+class GroupStat(MLProcess):
+    def __init__(self, group_col: str, agg_col: str, agg_funcs: list):
+        self.group_col = group_col
+        self.agg_col = agg_col
+        self.agg_funcs = agg_funcs
+        self.group_df = None
+
+    def fit(self, df: pd.DataFrame):
+        group_df = df.groupby(self.group_col)[self.agg_col].agg(self.agg_funcs).reset_index()
+        group_df.columns = [self.group_col] + [
+            f"{self.agg_col}_{agg_func}_by_{self.group_col}" for agg_func in self.agg_funcs
+        ]
+        self.group_df = group_df
+
+    def transform(self, df: pd.DataFrame) -> pd.DataFrame:
+        df = df.merge(self.group_df, on=self.group_col, how="left")
+        return df
+
+
+class SplitBins(MLProcess):
+    def __init__(self, cols: str, strategy: str = 'quantile'):
+        self.cols = cols
+        self.strategy = strategy
+        self.encoder = None
+
+    def fit(self, df: pd.DataFrame):
+        self.encoder = KBinsDiscretizer(strategy=self.strategy, encode='ordinal')
+        self.encoder.fit(df[self.cols].fillna(0))
+
+    def transform(self, df: pd.DataFrame) -> pd.DataFrame:
+        df[self.cols] = self.encoder.transform(df[self.cols].fillna(0))
+        return df
+
+# @registry.register("feature_engineering", ExtractTimeComps)
+# def extract_time_comps(df, time_col, time_comps):
+#     time_s = pd.to_datetime(df[time_col], errors="coerce")
+#     time_comps_df = pd.DataFrame()
+#
+#     if "year" in time_comps:
+#         time_comps_df["year"] = time_s.dt.year
+#     if "month" in time_comps:
+#         time_comps_df["month"] = time_s.dt.month
+#     if "day" in time_comps:
+#         time_comps_df["day"] = time_s.dt.day
+#     if "hour" in time_comps:
+#         time_comps_df["hour"] = time_s.dt.hour
+#     if "dayofweek" in time_comps:
+#         time_comps_df["dayofweek"] = time_s.dt.dayofweek + 1
+#     if "is_weekend" in time_comps:
+#         time_comps_df["is_weekend"] = time_s.dt.dayofweek.isin([5, 6]).astype(int)
+#     df = pd.concat([df, time_comps_df], axis=1)
+#     return df
+#
+#
+# @registry.register("feature_engineering", FeShiftByTime)
+# def fe_shift_by_time(df, time_col, group_col, shift_col, periods, freq):
+#     df[time_col] = pd.to_datetime(df[time_col])
+#
+#     def shift_datetime(date, offset, unit):
+#         if unit in ["year", "y", "Y"]:
+#             return date + relativedelta(years=offset)
+#         elif unit in ["month", "m", "M"]:
+#             return date + relativedelta(months=offset)
+#         elif unit in ["day", "d", "D"]:
+#             return date + relativedelta(days=offset)
+#         elif unit in ["week", "w", "W"]:
+#             return date + relativedelta(weeks=offset)
+#         elif unit in ["hour", "h", "H"]:
+#             return date + relativedelta(hours=offset)
+#         else:
+#             return date
+#
+#     def shift_by_time_on_key(
+#         inner_df, time_col, group_col, shift_col, offset, unit, col_name
+#     ):
+#         inner_df = inner_df.drop_duplicates()
+#         inner_df[time_col] = inner_df[time_col].map(
+#             lambda x: shift_datetime(x, offset, unit)
+#         )
+#         inner_df = inner_df.groupby([time_col, group_col], as_index=False)[
+#             shift_col
+#         ].mean()
+#         inner_df.rename(columns={shift_col: col_name}, inplace=True)
+#         return inner_df
+#
+#     shift_df = df[[time_col, group_col, shift_col]].copy()
+#     for period in periods:
+#         new_col_name = f"{group_col}_{shift_col}_lag_{period}_{freq}"
+#         tmp = shift_by_time_on_key(
+#             shift_df, time_col, group_col, shift_col, period, freq, new_col_name
+#         )
+#         df = df.merge(tmp, on=[time_col, group_col], how="left")
+#
+#     return df
+#
+#
+# @registry.register("feature_engineering", FeRollingByTime)
+# def fe_rolling_by_time(df, time_col, group_col, rolling_col, periods, freq, agg_funcs):
+#     df[time_col] = pd.to_datetime(df[time_col])
+#
+#     def rolling_by_time_on_key(inner_df, offset, unit, agg_func, col_name):
+#         time_freq = {
+#             "Y": [365 * offset, "D"],
+#             "M": [30 * offset, "D"],
+#             "D": [offset, "D"],
+#             "W": [7 * offset, "D"],
+#             "H": [offset, "h"],
+#         }
+#
+#         if agg_func not in ["mean", "std", "max", "min", "median", "sum", "count"]:
+#             raise ValueError(f"Invalid agg function: {agg_func}")
+#
+#         rolling_feat = inner_df.rolling(
+#             f"{time_freq[unit][0]}{time_freq[unit][1]}", closed="left"
+#         )
+#         rolling_feat = getattr(rolling_feat, agg_func)()
+#         depth = df.columns.nlevels
+#         rolling_feat = rolling_feat.stack(list(range(depth)))
+#         rolling_feat.name = col_name
+#         return rolling_feat
+#
+#     rolling_df = df[[time_col, group_col, rolling_col]].copy()
+#     for period in periods:
+#         for func in agg_funcs:
+#             new_col_name = f"{group_col}_{rolling_col}_rolling_{period}_{freq}_{func}"
+#             tmp = pd.pivot_table(
+#                 rolling_df,
+#                 index=time_col,
+#                 values=rolling_col,
+#                 columns=group_col,
+#             )
+#             tmp = rolling_by_time_on_key(tmp, period, freq, func, new_col_name)
+#             df = df.merge(tmp, on=[time_col, group_col], how="left")
+#
+#     return df
+
+
+class GeneralSelection(MLProcess):
+    def __init__(self, label_col: str):
+        self.label_col = label_col
+        self.feats = []
+
+    def fit(self, df: pd.DataFrame):
+        feats = [f for f in df.columns if f != self.label_col]
+        for col in df.columns:
+            if df[col].isnull().sum() / df.shape[0] == 1:
+                feats.remove(col)
+
+            if df[col].nunique() == 1:
+                feats.remove(col)
+
+            if (
+                df.loc[df[col] == np.inf].shape[0] != 0
+                or df.loc[df[col] == np.inf].shape[0] != 0
+            ):
+                feats.remove(col)
+        self.feats = feats
+
+    def transform(self, df: pd.DataFrame) -> pd.DataFrame:
+        df = df[self.feats]
+        return df

From 07771a769955f900b305334920d2ef5c70eae5bc Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Tue, 12 Dec 2023 10:56:51 +0800
Subject: [PATCH 17/49] add ml Class tool schema

---
 .../functions/schemas/data_preprocess.yml     | 306 +++++++++++++
 .../functions/schemas/feature_engineering.yml | 429 ++++++++++++++++++
 2 files changed, 735 insertions(+)
 create mode 100644 metagpt/tools/functions/schemas/data_preprocess.yml
 create mode 100644 metagpt/tools/functions/schemas/feature_engineering.yml

diff --git a/metagpt/tools/functions/schemas/data_preprocess.yml b/metagpt/tools/functions/schemas/data_preprocess.yml
new file mode 100644
index 000000000..95b0124cc
--- /dev/null
+++ b/metagpt/tools/functions/schemas/data_preprocess.yml
@@ -0,0 +1,306 @@
+FillMissingValue:
+  type: class
+  description: "Completing missing values with simple strategies"
+  methods:
+    __init__:
+      description: "Initialize self."
+      parameters:
+        properties:
+          features:
+            type: list
+            description: "columns to be processed"
+          strategy:
+            type: str
+            description: "the imputation strategy"
+            default: mean
+            enum:
+              - mean
+              - median
+              - most_frequent
+              - constant
+          fill_value:
+            type: int
+            description: "fill_value is used to replace all occurrences of missing_values"
+            default: null
+        required:
+          - features
+    fit:
+      description: "Fit the FillMissingValue model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+    transform:
+      description: "Transform the input DataFrame with the fitted model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+    fit_transform:
+      description: "Fit and transform the input DataFrame."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+
+MinMaxScale:
+  type: class
+  description: "Transform features by scaling each feature to a range, witch is (0, 1)"
+  methods:
+    __init__:
+      description: "Initialize self."
+      parameters:
+        properties:
+          features:
+            type: list
+            description: "columns to be processed"
+        required:
+          - features
+    fit:
+      description: "Fit the MinMaxScale model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+    transform:
+      description: "Transform the input DataFrame with the fitted model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+    fit_transform:
+      description: "Fit and transform the input DataFrame."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+
+StandardScale:
+  type: class
+  description: "Standardize features by removing the mean and scaling to unit variance"
+  methods:
+    __init__:
+      description: "Initialize self."
+      parameters:
+        properties:
+          features:
+            type: list
+            description: "columns to be processed"
+        required:
+          - features
+    fit:
+      description: "Fit the StandardScale model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+    transform:
+      description: "Transform the input DataFrame with the fitted model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+    fit_transform:
+      description: "Fit and transform the input DataFrame."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+
+MaxAbsScale:
+  type: class
+  description: "cale each feature by its maximum absolute value"
+  methods:
+    __init__:
+      description: "Initialize self."
+      parameters:
+        properties:
+          features:
+            type: list
+            description: "columns to be processed"
+        required:
+          - features
+    fit:
+      description: "Fit the MaxAbsScale model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+    transform:
+      description: "Transform the input DataFrame with the fitted model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+    fit_transform:
+      description: "Fit and transform the input DataFrame."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+
+LabelEncode:
+  type: class
+  description: "Apply label encoding to specified categorical columns in-place."
+  methods:
+    __init__:
+      description: "Initialize self."
+      parameters:
+        properties:
+          features:
+            type: list
+            description: "Categorical columns to be label encoded"
+        required:
+          - features
+    fit:
+      description: "Fit the LabelEncode model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+    transform:
+      description: "Transform the input DataFrame with the fitted model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+    fit_transform:
+      description: "Fit and transform the input DataFrame."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+
+OneHotEncode:
+  type: class
+  description: "Apply one-hot encoding to specified categorical columns, the original columns will be dropped."
+  methods:
+    __init__:
+      description: "Initialize self."
+      parameters:
+        properties:
+          features:
+            type: list
+            description: "Categorical columns to be one-hot encoded and dropped"
+        required:
+          - features
+    fit:
+      description: "Fit the OneHotEncoding model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+    transform:
+      description: "Transform the input DataFrame with the fitted model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+    fit_transform:
+      description: "Fit and transform the input DataFrame."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
\ No newline at end of file
diff --git a/metagpt/tools/functions/schemas/feature_engineering.yml b/metagpt/tools/functions/schemas/feature_engineering.yml
new file mode 100644
index 000000000..2cc4ec2fa
--- /dev/null
+++ b/metagpt/tools/functions/schemas/feature_engineering.yml
@@ -0,0 +1,429 @@
+PolynomialExpansion:
+  type: class
+  description: "Add polynomial and interaction features from selected numeric columns, excluding the bias column."
+  methods:
+    __init__:
+      description: "Initialize self."
+      parameters:
+        properties:
+          cols:
+            type: list
+            description: "Columns for polynomial expansion."
+          degree:
+            type: int
+            description: "The degree of the polynomial features."
+            default: 2
+        required:
+          - cols
+    fit:
+      description: "Fit the PolynomialExpansion model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+    transform:
+      description: "Transform the input DataFrame with the fitted model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+    fit_transform:
+      description: "Fit and transform the input DataFrame."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+
+CatCount:
+  type: class
+  description: "Add value counts of categorical columns as new features."
+  methods:
+    __init__:
+      description: "Initialize self."
+      parameters:
+        properties:
+          cols:
+            type: list
+            description: "Columns for value counts."
+        required:
+          - cols
+    fit:
+      description: "Fit the CatCount model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+    transform:
+      description: "Transform the input DataFrame with the fitted model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+    fit_transform:
+      description: "Fit and transform the input DataFrame."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+
+TargetMeanEncoder:
+  type: class
+  description: "Encodes a categorical column by the mean of the label column, and adds the result as a new feature."
+  methods:
+    __init__:
+      description: "Initialize self."
+      parameters:
+        properties:
+          col:
+            type: str
+            description: "Column to be mean encoded."
+          label:
+            type: str
+            description: "Predicted label column."
+        required:
+          - col
+          - label
+    fit:
+      description: "Fit the TargetMeanEncoder model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+    transform:
+      description: "Transform the input DataFrame with the fitted model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+    fit_transform:
+      description: "Fit and transform the input DataFrame."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+
+KFoldTargetMeanEncoder:
+  type: class
+  description: "Adds a new feature to the DataFrame by k-fold mean encoding of a categorical column using the label column."
+  methods:
+    __init__:
+      description: "Initialize self."
+      parameters:
+        properties:
+          col:
+            type: str
+            description: "Column to be k-fold mean encoded."
+          label:
+            type: str
+            description: "Predicted label column."
+          n_splits:
+            type: int
+            description: "Number of splits for K-fold."
+            default: 5
+          random_state:
+            type: int
+            description: "Random seed."
+            default: 2021
+        required:
+          - col
+          - label
+    fit:
+      description: "Fit the KFoldTargetMeanEncoder model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+    transform:
+      description: "Transform the input DataFrame with the fitted model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+    fit_transform:
+      description: "Fit and transform the input DataFrame."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+
+CatCross:
+  type: class
+  description: "Add pairwise crossed features and convert them to numerical features."
+  methods:
+    __init__:
+      description: "Initialize self."
+      parameters:
+        properties:
+          cols:
+            type: list
+            description: "Columns to be pairwise crossed."
+          max_cat_num:
+            type: int
+            description: "Maximum unique categories per crossed feature."
+            default: 100
+      required:
+        - cols
+    fit:
+      description: "Fit the CatCross model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+    transform:
+      description: "Transform the input DataFrame with the fitted model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+    fit_transform:
+      description: "Fit and transform the input DataFrame."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+
+GroupStat:
+  type: class
+  description: "Aggregate specified column in a DataFrame grouped by another column, adding new features named '<agg_col>_<agg_func>_by_<group_col>'."
+  methods:
+    __init__:
+      description: "Initialize self."
+      parameters:
+        properties:
+          group_col:
+            type: str
+            description: "Column used for grouping."
+          agg_col:
+            type: str
+            description: "Column on which aggregation is performed."
+          agg_funcs:
+            type: list
+            description: >-
+              List of aggregation functions to apply, such as ['mean', 'std'].
+              Each function must be supported by pandas.
+        required:
+          - group_col
+          - agg_col
+          - agg_funcs
+    fit:
+      description: "Fit the GroupStat model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+    transform:
+      description: "Transform the input DataFrame with the fitted model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+    fit_transform:
+      description: "Fit and transform the input DataFrame."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+
+SplitBins:
+  type: class
+  description: "Bin continuous data into intervals and return the bin identifier encoded as an integer value"
+  methods:
+    __init__:
+      description: "Initialize self."
+      parameters:
+        properties:
+          cols:
+            type: list
+            description: "Columns to be binned."
+          strategy:
+            type: str
+            description: "Strategy used to define the widths of the bins."
+            default: quantile
+        required:
+          - cols
+    fit:
+      description: "Fit the SplitBins model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+    transform:
+      description: "Transform the input DataFrame with the fitted model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+    fit_transform:
+      description: "Fit and transform the input DataFrame."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+
+GeneralSelection:
+  type: class
+  description: "Drop all nan feats and feats with only one unique value."
+  methods:
+    __init__:
+      description: "Initialize self."
+      parameters:
+        properties:
+          label_col:
+            type: str
+            description: "Label column name."
+        required:
+          - label_col
+    fit:
+      description: "Fit the GeneralSelection model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+    transform:
+      description: "Transform the input DataFrame with the fitted model."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
+    fit_transform:
+      description: "Fit and transform the input DataFrame."
+      parameters:
+        properties:
+          df:
+            type: DataFrame
+            description: "The input DataFrame."
+        required:
+          - df
+      returns:
+        df:
+          type: DataFrame
+          description: "The transformed DataFrame."
\ No newline at end of file

From 988c7072ef6084b0a4cf46d55cc6023f36c0b8b8 Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Tue, 12 Dec 2023 13:36:54 +0800
Subject: [PATCH 18/49] give history code for current code steps

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

diff --git a/metagpt/actions/write_code_steps.py b/metagpt/actions/write_code_steps.py
index 0bfb9c225..a19549b71 100644
--- a/metagpt/actions/write_code_steps.py
+++ b/metagpt/actions/write_code_steps.py
@@ -63,7 +63,7 @@ class WriteCodeSteps(Action):
 
     def get_context(self, plan: Plan):
         user_requirement = plan.goal
-        select_task_keys = ['task_id', 'instruction', 'is_finished', 'code_steps']
+        select_task_keys = ['task_id', 'instruction', 'is_finished', 'code']
 
         def process_task(task):
             task_dict = task.dict()

From 9c426f73dc7ba876fbe076a4d5f71996424fcfcf Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Tue, 12 Dec 2023 13:39:41 +0800
Subject: [PATCH 19/49] fix bug

---
 metagpt/tools/functions/libs/feature_engineering.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/metagpt/tools/functions/libs/feature_engineering.py b/metagpt/tools/functions/libs/feature_engineering.py
index 06a988d9a..67247d0d1 100644
--- a/metagpt/tools/functions/libs/feature_engineering.py
+++ b/metagpt/tools/functions/libs/feature_engineering.py
@@ -283,5 +283,5 @@ class GeneralSelection(MLProcess):
         self.feats = feats
 
     def transform(self, df: pd.DataFrame) -> pd.DataFrame:
-        df = df[self.feats]
+        df = df[self.feats + [self.label_col]]
         return df

From 1da4409475579705e5c7a44e1a873e337d02eb83 Mon Sep 17 00:00:00 2001
From: stellahsr <stellahsr@126.com>
Date: Tue, 12 Dec 2023 16:10:05 +0800
Subject: [PATCH 20/49] add step plan

---
 metagpt/roles/ml_engineer.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index 45fe728dd..8ad75b399 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -156,6 +156,11 @@ class MLEngineer(Role):
             # ask for acceptance, users can other refuse and change tasks in the plan
             task_result_confirmed = await self._ask_review()
             
+            # 针对当前task进行单独plan
+            if not success or not task_result_confirmed:
+                # fixme: 增加对应plan
+                self.state.plan()
+                
             if success and task_result_confirmed:
                 # tick off this task and record progress
                 task.code = code
@@ -203,8 +208,6 @@ class MLEngineer(Role):
             if counter == 0:
                 context = self.get_useful_memories()
             else:
-                # context = self.get_useful_memories()
-                # logger.info(f"context {context}")
                 improve_code = await DebugCode().run(plan=self.plan.current_task.instruction,
                                                      finished_code=code_context,
                                                      finished_code_result=code_result,
@@ -255,6 +258,8 @@ class MLEngineer(Role):
             
             counter += 1
         
+        success = False
+        
         return code, result, success, code_steps
     
     async def _ask_review(self):

From 49779d8615e4b05b759b549b6d7ceb9b5258ec0a Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Wed, 13 Dec 2023 13:35:22 +0800
Subject: [PATCH 21/49] refine schema desc

---
 metagpt/tools/functions/schemas/feature_engineering.yml | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/metagpt/tools/functions/schemas/feature_engineering.yml b/metagpt/tools/functions/schemas/feature_engineering.yml
index 2cc4ec2fa..4f2a7100d 100644
--- a/metagpt/tools/functions/schemas/feature_engineering.yml
+++ b/metagpt/tools/functions/schemas/feature_engineering.yml
@@ -328,7 +328,7 @@ GroupStat:
 
 SplitBins:
   type: class
-  description: "Bin continuous data into intervals and return the bin identifier encoded as an integer value"
+  description: "Inplace binning of continuous data into intervals, returning integer-encoded bin identifiers directly."
   methods:
     __init__:
       description: "Initialize self."
@@ -336,11 +336,15 @@ SplitBins:
         properties:
           cols:
             type: list
-            description: "Columns to be binned."
+            description: "Columns to be binned inplace."
           strategy:
             type: str
             description: "Strategy used to define the widths of the bins."
             default: quantile
+            enum:
+              - quantile
+              - uniform
+              - kmeans
         required:
           - cols
     fit:

From f614fbfa7c3e22e968bc4229271df092c3be9575 Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Wed, 13 Dec 2023 13:37:40 +0800
Subject: [PATCH 22/49] update ml tools

---
 metagpt/tools/functions/libs/data_preprocess.py  | 16 +++-------------
 .../tools/functions/libs/feature_engineering.py  |  5 +++++
 2 files changed, 8 insertions(+), 13 deletions(-)

diff --git a/metagpt/tools/functions/libs/data_preprocess.py b/metagpt/tools/functions/libs/data_preprocess.py
index 39474b0fd..fa70bf8fc 100644
--- a/metagpt/tools/functions/libs/data_preprocess.py
+++ b/metagpt/tools/functions/libs/data_preprocess.py
@@ -1,6 +1,6 @@
 import numpy as np
 from sklearn.impute import SimpleImputer
-from sklearn.preprocessing import KBinsDiscretizer, LabelEncoder
+from sklearn.preprocessing import LabelEncoder
 from sklearn.preprocessing import MaxAbsScaler
 from sklearn.preprocessing import MinMaxScaler
 from sklearn.preprocessing import OneHotEncoder
@@ -8,7 +8,6 @@ from sklearn.preprocessing import OrdinalEncoder
 from sklearn.preprocessing import RobustScaler
 from sklearn.preprocessing import StandardScaler
 
-from metagpt.tools.functions import registry
 from metagpt.tools.functions.libs.base import MLProcess
 from metagpt.tools.functions.schemas.data_preprocess import *
 
@@ -57,15 +56,6 @@ class StandardScale(MLProcess):
         return df
 
 
-@registry.register("data_preprocess", LogTransform)
-def log_transform(df: pd.DataFrame, features: list, ):
-    for col in features:
-        if df[col].min() <= 0:
-            df[col] = df[col] - df[col].min() + 2
-        df[col] = np.log(df[col])
-    return df
-
-
 class MaxAbsScale(MLProcess):
     def __init__(self, features: list,):
         self.features = features
@@ -146,7 +136,7 @@ class LabelEncode(MLProcess):
         return df
 
 
-def get_column_info(df: pd.DataFrame) -> str:
+def get_column_info(df: pd.DataFrame) -> dict:
     data = []
     for i in df.columns:
         nan_freq = float("%.2g" % (df[i].isna().mean() * 100))
@@ -157,7 +147,7 @@ def get_column_info(df: pd.DataFrame) -> str:
         data,
         columns=["Column_name", "Data_type", "NaN_Frequency(%)", "N_unique"],
     )
-    return samples.to_string(index=False)
+    return samples.to_dict(orient='list')
 #
 #
 # if __name__ == '__main__':
diff --git a/metagpt/tools/functions/libs/feature_engineering.py b/metagpt/tools/functions/libs/feature_engineering.py
index 67247d0d1..de54e4db0 100644
--- a/metagpt/tools/functions/libs/feature_engineering.py
+++ b/metagpt/tools/functions/libs/feature_engineering.py
@@ -10,6 +10,7 @@ import numpy as np
 from dateutil.relativedelta import relativedelta
 from joblib import Parallel, delayed
 from pandas.api.types import is_numeric_dtype
+from pandas.core.dtypes.common import is_object_dtype
 from sklearn.model_selection import KFold
 from sklearn.preprocessing import PolynomialFeatures, KBinsDiscretizer
 
@@ -280,6 +281,10 @@ class GeneralSelection(MLProcess):
                 or df.loc[df[col] == np.inf].shape[0] != 0
             ):
                 feats.remove(col)
+
+            if is_object_dtype(df[col]) and df[col].nunique() == df.shape[0]:
+                feats.remove(col)
+
         self.feats = feats
 
     def transform(self, df: pd.DataFrame) -> pd.DataFrame:

From 92d59ea31bb7bcb563d2fdd94cd6b6af64963aa7 Mon Sep 17 00:00:00 2001
From: yzlin <yzlin@fuzhi.ai>
Date: Wed, 13 Dec 2023 13:48:18 +0800
Subject: [PATCH 23/49] save code steps early

---
 metagpt/actions/write_analysis_code.py |  7 ++-----
 metagpt/roles/ml_engineer.py           | 11 +++++------
 metagpt/schema.py                      |  2 +-
 3 files changed, 8 insertions(+), 12 deletions(-)

diff --git a/metagpt/actions/write_analysis_code.py b/metagpt/actions/write_analysis_code.py
index 1127dc78b..7e6483371 100644
--- a/metagpt/actions/write_analysis_code.py
+++ b/metagpt/actions/write_analysis_code.py
@@ -23,9 +23,7 @@ from metagpt.utils.common import create_func_config
 
 
 class BaseWriteAnalysisCode(Action):
-    async def run(
-        self, context: List[Message], plan: Plan = None, task_guide: str = ""
-    ) -> str:
+    async def run(self, context: List[Message], plan: Plan = None) -> str:
         """Run of a code writing action, used in data analysis or modeling
 
         Args:
@@ -85,7 +83,6 @@ class WriteCodeByGenerate(BaseWriteAnalysisCode):
         self,
         context: [List[Message]],
         plan: Plan = None,
-        code_steps: str = "",
         system_msg: str = None,
         **kwargs,
     ) -> str:
@@ -155,11 +152,11 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
         self,
         context: List[Message],
         plan: Plan = None,
-        code_steps: str = "",
         data_desc: str = "",
     ) -> str:
         task_type = plan.current_task.task_type
         task = plan.current_task.instruction
+        code_steps = plan.current_task.code_steps
         available_tools = registry.get_all_schema_by_module(task_type)
         available_tools = [
             {k: tool[k] for k in ["name", "description"] if k in tool}
diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index de649e857..3260dd43f 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -59,7 +59,7 @@ class MLEngineer(Role):
             logger.info(f"ready to take on task {task}")
 
             # take on current task
-            code, result, success, code_steps = await self._write_and_exec_code()
+            code, result, success = await self._write_and_exec_code()
 
             # ask for acceptance, users can other refuse and change tasks in the plan
             review, task_result_confirmed = await self._ask_review(trigger=ReviewConst.TASK_REVIEW_TRIGGER)
@@ -73,7 +73,6 @@ class MLEngineer(Role):
                 # tick off this task and record progress
                 task.code = code
                 task.result = result
-                task.code_steps = code_steps
                 self.plan.finish_current_task()
                 self.working_memory.clear()
 
@@ -102,7 +101,7 @@ class MLEngineer(Role):
         return rsp
 
     async def _write_and_exec_code(self, max_retry: int = 3):
-        code_steps = (
+        self.plan.current_task.code_steps = (
             await WriteCodeSteps().run(self.plan)
             if self.use_code_steps
             else ""
@@ -121,12 +120,12 @@ class MLEngineer(Role):
             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, code_steps=code_steps, temperature=0.0
+                    context=context, plan=self.plan, temperature=0.0
                 )
                 cause_by = WriteCodeByGenerate
             else:
                 code = await WriteCodeWithTools().run(
-                    context=context, plan=self.plan, code_steps=code_steps, data_desc=""
+                    context=context, plan=self.plan, data_desc=""
                 )
                 cause_by = WriteCodeWithTools
 
@@ -151,7 +150,7 @@ class MLEngineer(Role):
                 if ReviewConst.CHANGE_WORD[0] in review:
                     counter = 0  # redo the task again with help of human suggestions
 
-        return code, result, success, code_steps
+        return code, result, success
 
     async def _ask_review(self, auto_run: bool = None, trigger: str = ReviewConst.TASK_REVIEW_TRIGGER):
         auto_run = auto_run or self.auto_run
diff --git a/metagpt/schema.py b/metagpt/schema.py
index f91922535..8eb7e31ca 100644
--- a/metagpt/schema.py
+++ b/metagpt/schema.py
@@ -78,10 +78,10 @@ class Task(BaseModel):
     dependent_task_ids: list[str] = [] # Tasks prerequisite to this Task
     instruction: str = ""
     task_type: str = ""
+    code_steps: str = ""
     code: str = ""
     result: str = ""
     is_finished: bool = False
-    code_steps: str = ""
 
 
 class Plan(BaseModel):

From 33810829a072467a8f61f2f7dc14ffd1792e793a Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Wed, 13 Dec 2023 14:31:32 +0800
Subject: [PATCH 24/49] support tool in debug

---
 metagpt/actions/debug_code.py | 106 ++++++++++++++++++++++------------
 1 file changed, 68 insertions(+), 38 deletions(-)

diff --git a/metagpt/actions/debug_code.py b/metagpt/actions/debug_code.py
index 9efe93efc..53ca2f54d 100644
--- a/metagpt/actions/debug_code.py
+++ b/metagpt/actions/debug_code.py
@@ -3,7 +3,7 @@ from typing import Dict, List, Union, Tuple, Optional, Any
 from metagpt.actions import Action
 from metagpt.logs import logger
 from metagpt.schema import Message, Plan
-from metagpt.utils.common import CodeParser
+from metagpt.utils.common import CodeParser, create_func_config
 from metagpt.actions.write_analysis_code import BaseWriteAnalysisCode
 
 DEBUG_REFLECTION_EXAMPLE = '''Example 1:
@@ -39,25 +39,39 @@ DEBUG_REFLECTION_EXAMPLE = '''Example 1:
 REFLECTION_PROMPT = """
                        Here is an example for you.
                        {debug_example}
-                       [requirement]
-                       {goal}
-                       [finished code]
-                       finished code are executable, and you should based on the code to continue your current code debug
-                       {finished_code}
-                       
-                       try to reuse the code here to understand the coding task.
+                       [context]
+                       {context}
                        
                        [previous impl]
                        {code}
                        [runtime Error]
                        {runtime_result}
 
-                        Analysis the error step by step, provide me improve method. Do not repeat [previous impl]
+                        Analysis the error step by step, provide me improve method and code. Remember to follow [context] requirement.
                         [reflection on previous impl]:
                         xxx
 
                        """
 
+CODE_REFLECTION = {
+    "name": "execute_reflection_code",
+    "description": "Execute reflection code.",
+    "parameters": {
+        "type": "object",
+        "properties": {
+            "reflection": {
+                "type": "string",
+                "description": "Reflection on previous impl.",
+            },
+            "improved_impl": {
+                "type": "string",
+                "description": "Refined code after reflection.",
+            },
+        },
+        "required": ["reflection", "improved_impl"],
+    },
+}
+
 
 def message_to_str(message: Message) -> str:
     return f"{message.role}: {message.content}"
@@ -75,52 +89,68 @@ class DebugCode(BaseWriteAnalysisCode):
     def __init__(self, **kwargs: Any):
         super().__init__(**kwargs)
     
-    async def run_reflection(self, goal, finished_code, finished_code_result, code, runtime_result) -> str:
+    async def run_reflection(
+        self,
+        # goal,
+        # finished_code,
+        # finished_code_result,
+        context: List[Message],
+        code,
+        runtime_result,
+    ) -> dict:
         info = []
-        finished_code_and_result = finished_code + "\n [finished results]\n\n" + finished_code_result
+        # finished_code_and_result = finished_code + "\n [finished results]\n\n" + finished_code_result
         reflection_prompt = REFLECTION_PROMPT.format(debug_example=DEBUG_REFLECTION_EXAMPLE,
-                                                     goal=goal,
-                                                     finished_code=finished_code_and_result,
+                                                     context=context,
+                                                     # goal=goal,
+                                                     # finished_code=finished_code_and_result,
                                                      code=code,
                                                      runtime_result=runtime_result
                                                      )
-        system_prompt = "You are an AI Python assistant. You will be given your previous implementation of a function, runtime error results, and a hint to change the implementation appropriately. Write your full implementation "
+        system_prompt = "You are an AI Python assistant. You will be given your previous implementation code of a task, runtime error results, and a hint to change the implementation appropriately. Write your full implementation "
         info.append(Message(role="system", content=system_prompt))
-        info.append(Message(role="assistant", content=reflection_prompt))
+        info.append(Message(role="user", content=reflection_prompt))
         
-        msg = messages_to_str(info)
-        resp = await self.llm.aask(msg=msg)
+        # msg = messages_to_str(info)
+        # resp = await self.llm.aask(msg=msg)
+        resp = await self.llm.aask_code(messages=info, **create_func_config(CODE_REFLECTION))
         logger.info(f"reflection is {resp}")
         return resp
     
-    async def rewrite_code(self, reflection: str = "", code_context: str = "") -> str:
-        """
-        根据reflection重写代码
-        """
-        info = []
-        info.append(Message(role="assistant", content=f"[code context]:{code_context}"
-                                                      f"finished code are executable, and you should based on the code to continue your current code debug and improvement"
-                                                      f"[reflection]: \n {reflection}"))
-        info.append(Message(role="user", content=f"[improved impl]:\n Return in Python block"))
-        msg = messages_to_str(info)
-        resp = await self.llm.aask(msg=msg)
-        logger.info(f"improve code is {resp}")
-        improv_code = CodeParser.parse_code(block=None, text=resp)
-        return improv_code
+    # async def rewrite_code(self, reflection: str = "", context: List[Message] = None) -> str:
+    #     """
+    #     根据reflection重写代码
+    #     """
+    #     info = context
+    #     # info.append(Message(role="assistant", content=f"[code context]:{code_context}"
+    #     #                                               f"finished code are executable, and you should based on the code to continue your current code debug and improvement"
+    #     #                                               f"[reflection]: \n {reflection}"))
+    #     info.append(Message(role="assistant", content=f"[reflection]: \n {reflection}"))
+    #     info.append(Message(role="user", content=f"[improved impl]:\n Return in Python block"))
+    #     msg = messages_to_str(info)
+    #     resp = await self.llm.aask(msg=msg)
+    #     improv_code = CodeParser.parse_code(block=None, text=resp)
+    #     return improv_code
     
     async def run(self,
+                  context: List[Message] = None,
                   plan: str = "",
-                  finished_code: str = "",
-                  finished_code_result: str = "",
+                  # finished_code: str = "",
+                  # finished_code_result: str = "",
                   code: str = "",
                   runtime_result: str = "") -> str:
         """
         根据当前运行代码和报错信息进行reflection和纠错
         """
-        reflection = await self.run_reflection(plan, finished_code=finished_code,
-                                               finished_code_result=finished_code_result,
-                                               code=code,
-                                               runtime_result=runtime_result)
+        reflection = await self.run_reflection(
+            # plan,
+            # finished_code=finished_code,
+            # finished_code_result=finished_code_result,
+            code=code,
+            context=context,
+            runtime_result=runtime_result,
+        )
         # 根据reflection结果重写代码
-        improv_code = await self.rewrite_code(reflection, code_context=finished_code)
+        # improv_code = await self.rewrite_code(reflection, context=context)
+        improv_code = reflection['improved_impl']
         return improv_code

From ab7af7768c00acc0c3f900430b402c64637f7b0f Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Wed, 13 Dec 2023 14:31:50 +0800
Subject: [PATCH 25/49] refine prompt

---
 metagpt/prompts/ml_engineer.py | 298 +++++++++++++--------------------
 1 file changed, 117 insertions(+), 181 deletions(-)

diff --git a/metagpt/prompts/ml_engineer.py b/metagpt/prompts/ml_engineer.py
index 5c7b9f82e..d11cbf453 100644
--- a/metagpt/prompts/ml_engineer.py
+++ b/metagpt/prompts/ml_engineer.py
@@ -4,6 +4,31 @@
 # @Author  : lidanyang
 # @File    : ml_engineer
 # @Desc    :
+UPDATE_DATA_COLUMNS = """
+# Background
+Keep dataset column information updated to reflect changes in training or testing datasets, aiding in informed decision-making during data analysis.
+## Done Tasks
+```python
+{history_code}
+```end
+
+# Task
+Update and print the dataset's column information only if the train or test data has changed. Use the following code:
+```python
+from metagpt.tools.functions.libs.data_preprocess import get_column_info
+
+column_info = get_column_info(df)
+print("df_column_info")
+print(column_info)
+```end
+
+# Constraints:
+- Use the DataFrame variable from 'Done Tasks' in place of df.
+- Import `get_column_info` only if it's not already imported.
+- Skip update if no changes in training/testing data, except for initial data load.
+- No need to update info if only model evaluation is performed.
+"""
+
 GEN_DATA_DESC_PROMPT = """
 Here is the head 5 rows of the dataset:
 {data_head}
@@ -34,7 +59,8 @@ Please assign a task type to each task in the list below from the given categori
 - **feature_engineering**: Only for creating new columns for input data.
 - **data_preprocess**: Only for changing value inplace.
 - **model_train**: Only for training model.
-- **other**: Any tasks that do not fit into the previous categories, such as visualization, summarizing findings, build model, etc.
+- **model_evaluate**: Only for evaluating model.
+- **other**: Any tasks that do not fit into the previous categories, such as visualization, summarizing findings, etc.
 """
 
 ASSIGN_TASK_TYPE = {
@@ -107,206 +133,122 @@ CODE_GENERATOR_WITH_TOOLS = {
     },
 }
 
-TOOL_USAGE_PROMPT = """
-## Target
-{goal}
 
-Specifically, {special_prompt}
-
-## History Info
-{context}
-
-## Code Steps for Current Task:
-Follow steps below when you writing code if it's convenient.
-{code_steps}
-
-## Available Tools:
-Each function is described in JSON format, including the function name and parameters. {output_desc}
-{function_catalog}
-
-When you call a function above, you should import the function from `{module_name}` first, e.g.:
-```python
-from metagpt.tools.functions.libs.data_preprocess import fill_missing_value
-```end
-
-## Your Output Format:
-Generate the complete code for this task:
-```python
-# Tools used: [function names or 'none']
-<your code for the current task, without any comments>
-```end
-
-## Attention:
-Make sure use the columns from the dataset columns: {column_names}
-Finish your coding tasks as a helpful programmer based on the tools.
-
-"""
+PRINT_DATA_COLUMNS = {
+    "name": "print_column_info",
+    "description": "Print the latest column information after 'Done Tasks' code if first read or data changed.",
+    "parameters": {
+        "type": "object",
+        "properties": {
+            "is_update": {
+                "type": "boolean",
+                "description": "Whether need to update the column info.",
+            },
+            "code": {
+                "type": "string",
+                "description": "The code to be added to a new cell in jupyter.",
+            },
+        },
+        "required": ["is_update", "code"],
+    },
+}
 
 GENERATE_CODE_PROMPT = """
-## Target
-{goal}
-
-Specifically, {special_prompt}
-
-
-## Finished Task and Code 
-{context}
-
-## Code Steps for Current Task:
-Follow steps below when you writing code if it's convenient.
-{code_steps}
-
-## Instruction
-Finished task and code are executable, and you should based on the code to continue your current task
-Do not repeat functions and code, try to reuse the code in [Finished Task and Code]
- 
-## Your Output Format:
-Generate the complete code for this task:
-```python
-import pandas as pd
-
-```
-
-## Attention:
-Make sure use the columns from the dataset columns
-Finish your coding tasks as a helpful programmer based on the code.
-
-"""
-
-TOOL_USAGE_PROMPT = """
-## Target
-{goal}
-
-## History Info
-{context}
-
-## Available Tools:
-Each function is described in JSON format, including the function name and parameters. {output_desc}
-{function_catalog}
-
-When you call a function above, you should import the function from `{module_name}` first, e.g.:
-```python
-from metagpt.tools.functions.libs.data_preprocess import fill_missing_value
-```end
-
-## Your Output Format:
-Generate the complete code for this task:
-```python
-# Tools used: [function names or 'none']
-<your code for the current task, without any comments>
-```end
-
-## Attention:
-Make sure use the columns from the dataset columns
-Finish your coding tasks as a helpful programmer based on the tools.
-"""
-
-TOOL_ORGANIZATION_PROMPT = """
-The previous conversation has provided all tasks step-by-step for the use goal and their statuses. 
-Now, begin writing code for the current task. This code should writen strictly on the basis of all previous completed tasks code, not a standalone code. And avoid writing duplicate code that has already been written in previous tasks, such as repeated import of packages, reading data, etc.
-Specifically, {special_prompt}
-You can utilize pre-defined tools in 'Available Tools' if the tools are sufficient. And you should combine the use of other public packages if necessary, like sklearn, numpy, pandas, etc..
-
-## Code Steps for Current Task:
-Follow steps below when you writing code if it's convenient.
-{code_steps}
-
-## Available Tools:
-Each function is described in JSON format, including the function name and parameters. {output_desc}
-{function_catalog}
-
-When you call a function above, you should import the function from `{module_name}` first, e.g.:
-```python
-from metagpt.tools.functions.libs.data_preprocess import fill_missing_value
-```end
-
-## Your Output Format:
-Generate the complete code for this task:
-```python
-# Tools used: [function names or 'none']
-<your code for the current task, without any comments>
-```end
-
-*** Important Rules ***
-- If you use tool not in the list, you should implement it by yourself.
-- Ensure the output new code is executable in the same Jupyter notebook environment with previous tasks code have been executed.
-- When write code for current task, remember the code should be coherent with previous tasks code.
-- Remember that don't process the columns have been processed in previous tasks and don't mock data yourself.
-- Prioritize using tools for the same functionality.
-"""
-
-DATA_PREPROCESS_PROMPT = """
-The current task is about data preprocessing, closely monitor each column's data type. Apply suitable methods for various types (numerical, categorical, datetime, textual, etc.) to ensure the pandas.DataFrame is correctly formatted.
-Additionally, ensure that the columns being processed must be the ones that actually exist in the dataset.
-Don't write processed data to files.
-"""
-
-FEATURE_ENGINEERING_PROMPT = """
-The current task is about feature engineering. when performing it, please adhere to the following principles:
-- Ensure that the feature you're working with is indeed present in the dataset and consider the data type (numerical, categorical, etc.) and application scenario (classification, regression tasks, etc.).
-- When generate new features, you should combine real world knowledge and decide what features are useful for the task.
-- Generate as diverse features as possible to improve the model's performance.
-- Before generating a new feature, ensure the used features are already processed and ready to use.
-"""
-
-DATA_PROCESS_PROMPT = """
 # Background
-As a data scientist, you need to help user to achieve the goal [{user_requirement}] step-by-step in an continuous Jupyter notebook.
+Assist in completing [{user_requirement}] in a Jupyter notebook.
 
-## Done Tasks
+## Task Progress
+### Done Tasks
 ```python
 {history_code}
 ```end
 
-## Current Task
+### Current Task
 {current_task}
 
-# Latest Data Info
-Latest data info after previous tasks:
+## Latest Data Info
 {column_info}
 
 # Task
-Write a Python function for 'Current Task'. Start by copying the input DataFrame. Avoid duplicating code from 'Done Tasks'.
-Specifically, {special_prompt}
+Fully implement 'Current Task', ensuring all necessary steps are covered without repeating code from 'Done Tasks'. Specifically, {special_prompt}
+
+# Code Steps:
+Follow steps below when you writing code if it's convenient.
+{code_steps}
+"""
+
+TOOL_USAGE_PROMPT = """
+# Background
+Assist in completing [{user_requirement}] in a Jupyter notebook.
+
+## Task Progress
+### Done Tasks
+```python
+{history_code}
+```end
+
+### Current Task
+{current_task}
+
+## Latest Data Info
+{column_info}
+
+# Task
+Fully implement 'Current Task', ensuring all necessary steps are covered without repeating code from 'Done Tasks'. Specifically, {special_prompt}
 
 # Code Steps:
 Follow steps below when you writing code if it's convenient.
 {code_steps}
 
 # Capabilities
-- You can utilize pre-defined tools in any code lines from 'Available Tools' in the form of python functions.
+- You can utilize pre-defined tools in any code lines from 'Available Tools' in the form of Python Class.
 - You can freely combine the use of any other public packages, like sklearn, numpy, pandas, etc..
-- You can do anything about data preprocessing, feature engineering, model training, etc..
 
 # Available Tools:
-Each function tool is described in JSON format. {output_desc}
-When you call a function below, import the function from `{module_name}` first.
-{function_catalog}
+Each Class tool is described in JSON format. When you call it, import the tool from `{module_name}` first.
+{tool_catalog}
 
 # Output Example:
-when current task is "fill missing value and handle outliers", the output code be like:
+For "fill missing value and handle outliers", the output code be like when there are training data and test data:
 ```python
-from metagpt.tools.functions.libs.data_preprocess import fill_missing_value
+# Tools used: ['FillMissingValue']
+from metagpt.tools.functions.libs.data_preprocess import FillMissingValue
 
-def function_name(df):
-    df_processed = df.copy()
-    num_cols = df_processed.select_dtypes(include='number').columns.tolist()
-    df_processed = fill_missing_value(df_processed, num_cols, 'mean')
-    
-    for col in num_cols:
-        low, high = df_processed[col].quantile([0.01, 0.99])
-        df_processed[col] = df_processed[col].clip(low, high)
-    return df_processed
+train_processed = train.copy()
+test_processed = test.copy()
+num_cols = train_processed.select_dtypes(include='number').columns.tolist()
+fill_missing_value = FillMissingValue(features=num_cols, strategy='mean')
+fill_missing_value.fit(train_processed)
+train_processed = fill_missing_value.transform(train_processed)
+test_processed = fill_missing_value.transform(test_processed)
 
-df_processed = function_name(df)
-print(df_processed.info())
+for col in num_cols:
+    low, high = train_processed[col].quantile([0.01, 0.99])
+    train_processed[col] = train_processed[col].clip(low, high)
+    test_processed[col] = test_processed[col].clip(low, high)
 ```end
 
 # Constraints:
-- Ensure the output new code is executable in the same Jupyter notebook with previous tasks code have been executed.
 - Prioritize using pre-defined tools for the same functionality.
-- Return DataFrame should always be named `df_processed`, while the input DataFrame should based on the done tasks' output DataFrame.
-- Limit to one print statement for the output DataFrame's info.
+- Copy DataFrame before processing if needed.
+- If 'Code Steps' contains step done in 'Done Tasks', such as reading data, don't repeat it.
+"""
+
+DATA_PREPROCESS_PROMPT = """
+The current task is about data preprocessing, please note the following:
+- Monitor data types per column, applying appropriate methods.
+- Ensure operations are on existing dataset columns.
+- Avoid writing processed data to files.
+- Prefer alternatives to one-hot encoding for categorical data.
+- Only encode necessary categorical columns to allow for potential feature-specific engineering tasks later.
+"""
+
+FEATURE_ENGINEERING_PROMPT = """
+The current task is about feature engineering. when performing it, please adhere to the following principles:
+- Ensure operations are on existing dataset columns and consider the data type (numerical, categorical, etc.) and application scenario (classification, regression tasks, etc.).
+- Create impactful features based on real-world knowledge and column info.
+- Generate as diverse features as possible to improve the model's performance.
+- If potential impactful features are not included in 'Code Steps', add new steps to generate them.
 """
 
 MODEL_TRAIN_PROMPT = """
@@ -316,23 +258,17 @@ The current task is about training a model, please ensure high performance:
 - Use the data from previous task result directly, do not mock or reload data yourself.
 """
 
-DATA_PREPROCESS_OUTPUT_DESC = "Please note that all functions output a updated pandas.DataFrame after data preprocessing."
-
-FEATURE_ENGINEERING_OUTPUT_DESC = "Please note that all functions output a updated pandas.DataFrame with new features added or existing features modified."
-
-CLASSIFICATION_MODEL_OUTPUT_DESC = ""
-
-REGRESSION_MODEL_OUTPUT_DESC = ""
+MODEL_EVALUATE_PROMPT = """
+The current task is about evaluating a model, please note the following:
+- Ensure that the evaluated data is same processed as the training data.
+- Use trained model from previous task result directly, do not mock or reload model yourself.
+"""
 
 ML_SPECIFIC_PROMPT = {
     "data_preprocess": DATA_PREPROCESS_PROMPT,
     "feature_engineering": FEATURE_ENGINEERING_PROMPT,
     "model_train": MODEL_TRAIN_PROMPT,
-}
-
-TOOL_OUTPUT_DESC = {
-    "data_preprocess": DATA_PREPROCESS_OUTPUT_DESC,
-    "feature_engineering": FEATURE_ENGINEERING_OUTPUT_DESC,
+    "model_evaluate": MODEL_EVALUATE_PROMPT,
 }
 
 ML_MODULE_MAP = {

From 537d51c26e29a1774269825e3611667b2436e80d Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Wed, 13 Dec 2023 14:32:25 +0800
Subject: [PATCH 26/49] write code with class tool

---
 metagpt/actions/write_analysis_code.py | 130 +++++++++----------
 metagpt/roles/ml_engineer.py           | 170 ++++++++++---------------
 2 files changed, 131 insertions(+), 169 deletions(-)

diff --git a/metagpt/actions/write_analysis_code.py b/metagpt/actions/write_analysis_code.py
index 58cab9c6a..aceebbfeb 100644
--- a/metagpt/actions/write_analysis_code.py
+++ b/metagpt/actions/write_analysis_code.py
@@ -4,7 +4,9 @@
 @Author  :   orange-crow
 @File    :   write_code_v2.py
 """
-from typing import Dict, List, Union, Tuple, Optional, Any
+from typing import Dict, List, Union, Tuple
+
+import yaml
 
 from metagpt.actions import Action
 from metagpt.logs import logger
@@ -15,11 +17,9 @@ from metagpt.prompts.ml_engineer import (
     TOOL_USAGE_PROMPT,
     ML_SPECIFIC_PROMPT,
     ML_MODULE_MAP,
-    TOOL_OUTPUT_DESC, DATA_PROCESS_PROMPT,
-    GENERATE_CODE_PROMPT
+    GENERATE_CODE_PROMPT,
 )
 from metagpt.schema import Message, Plan
-from metagpt.tools.functions import registry
 from metagpt.utils.common import create_func_config, remove_comments
 
 
@@ -100,40 +100,55 @@ class WriteCodeByGenerate(BaseWriteAnalysisCode):
 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) -> List[Dict]:
+    def __init__(self, name: str = "", context=None, llm=None, schema_path=None):
+        super().__init__(name, context, llm)
+        self.schema_path = schema_path
+        self.available_tools = {}
+
+        if self.schema_path is not None:
+            self._load_tools(schema_path)
+
+    def _load_tools(self, schema_path):
+        """Load tools from yaml file"""
+        yml_files = schema_path.glob("*.yml")
+        for yml_file in yml_files:
+            module = yml_file.stem
+            with open(yml_file, "r", encoding="utf-8") as f:
+                self.available_tools[module] = yaml.safe_load(f)
+
+    def _parse_recommend_tools(self, module: str, recommend_tools: list) -> dict:
         """
         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.
+            recommend_tools (list): A list of recommended tools.
 
         Returns:
-            List[Dict]: A list of dicts of valid tool schemas.
+            dict: A dict of valid tool schemas.
         """
         valid_tools = []
-        available_tools = registry.get_all_by_module(module).keys()
+        available_tools = self.available_tools[module].keys()
         for tool in recommend_tools:
             if tool in available_tools:
                 valid_tools.append(tool)
 
-        tool_catalog = registry.get_schemas(module, valid_tools)
+        tool_catalog = {tool: self.available_tools[module][tool] for tool in valid_tools}
         return tool_catalog
 
     async def _tool_recommendation(
-            self,
-            task: str,
-            code_steps: str,
-            available_tools: list
+        self,
+        task: str,
+        code_steps: str,
+        available_tools: dict,
     ) -> list:
         """
         Recommend tools for the specified task.
 
         Args:
-            context (List[Message]): Action output history, source action denoted by Message.cause_by
+            task (str): the task to recommend tools for
             code_steps (str): the code steps to generate the full code for the task
-            available_tools (list): the available tools for the task
+            available_tools (dict): the available tools description
 
         Returns:
             list: recommended tools for the specified task
@@ -149,27 +164,23 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
         return recommend_tools
 
     async def run(
-            self,
-            context: List[Message],
-            plan: Plan = None,
-            code_steps: str = "",
-            column_info: str = "",
-            **kwargs,
-    ) -> str:
+        self,
+        context: List[Message],
+        plan: Plan = None,
+        code_steps: str = "",
+        column_info: str = "",
+        **kwargs,
+    ) -> Tuple[List[Message], str]:
         task_type = plan.current_task.task_type
-        available_tools = registry.get_all_schema_by_module(task_type)
+        available_tools = self.available_tools.get(task_type, {})
         special_prompt = ML_SPECIFIC_PROMPT.get(task_type, "")
 
-        column_names = kwargs.get("column_names", {})
         finished_tasks = plan.get_finished_tasks()
         code_context = [remove_comments(task.code) for task in finished_tasks]
         code_context = "\n\n".join(code_context)
 
         if len(available_tools) > 0:
-            available_tools = [
-                {k: tool[k] for k in ["name", "description"] if k in tool}
-                for tool in available_tools
-            ]
+            available_tools = {k: v["description"] for k, v in available_tools.items()}
 
             recommend_tools = await self._tool_recommendation(
                 plan.current_task.instruction,
@@ -180,46 +191,27 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
             logger.info(f"Recommended tools: \n{recommend_tools}")
 
             module_name = ML_MODULE_MAP[task_type]
-            output_desc = TOOL_OUTPUT_DESC.get(task_type, "")
-            new_code = ""
-
-            for idx, tool in enumerate(recommend_tools):
-                hist_info = f"Previous finished code is \n\n ```Python {code_context} ``` \n\n "
-
-                prompt = TOOL_USAGE_PROMPT.format(
-                    goal=plan.current_task.instruction,
-                    context=hist_info,
-                    code_steps=code_steps,
-                    column_names=column_names,
-                    special_prompt=special_prompt,
-                    module_name=module_name,
-                    output_desc=output_desc,
-                    function_catalog=tool_catalog[idx],
-                )
-
-                tool_config = create_func_config(CODE_GENERATOR_WITH_TOOLS)
-
-                rsp = await self.llm.aask_code(prompt, **tool_config)
-                logger.info(f"rsp is: {rsp}")
-                # final_code = final_code + "\n\n" + rsp["code"]
-                # final_code[key] = rsp["code"]
-                new_code = new_code + "\n\n" + rsp["code"]
-                code_context = code_context + "\n\n" + rsp["code"]
-            return new_code
-
-        else:
-            hist_info = f"Previous finished code is \n\n ```Python {code_context} ``` \n\n "
-
-            prompt = GENERATE_CODE_PROMPT.format(
-                goal=plan.current_task.instruction,
-                context=hist_info,
-                code_steps=code_steps,
+            prompt = TOOL_USAGE_PROMPT.format(
+                user_requirement=plan.goal,
+                history_code=code_context,
+                current_task=plan.current_task.instruction,
+                column_info=column_info,
                 special_prompt=special_prompt,
-                # column_names=column_names
+                code_steps=code_steps,
+                module_name=module_name,
+                tool_catalog=tool_catalog,
+            )
+        else:
+            prompt = GENERATE_CODE_PROMPT.format(
+                user_requirement=plan.goal,
+                history_code=code_context,
+                current_task=plan.current_task.instruction,
+                column_info=column_info,
+                special_prompt=special_prompt,
+                code_steps=code_steps,
             )
 
-            tool_config = create_func_config(CODE_GENERATOR_WITH_TOOLS)
-            logger.info(f"prompt is: {prompt}")
-            rsp = await self.llm.aask_code(prompt, **tool_config)
-            logger.info(f"rsp is: {rsp}")
-            return rsp["code"]
+        tool_config = create_func_config(CODE_GENERATOR_WITH_TOOLS)
+        rsp = await self.llm.aask_code(prompt, **tool_config)
+        context = [Message(content=prompt, role="user")]
+        return context, rsp["code"]
diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index 45fe728dd..20589079d 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -1,5 +1,6 @@
 import json
 import re
+from datetime import datetime
 from typing import List
 
 import fire
@@ -10,12 +11,16 @@ from metagpt.actions.execute_code import ExecutePyCode
 from metagpt.actions.write_analysis_code import WriteCodeByGenerate, WriteCodeWithTools
 from metagpt.actions.write_code_steps import WriteCodeSteps
 from metagpt.actions.write_plan import WritePlan
-from metagpt.const import DATA_PATH
+from metagpt.const import DATA_PATH, PROJECT_ROOT
 from metagpt.logs import logger
-from metagpt.prompts.ml_engineer import GEN_DATA_DESC_PROMPT
+from metagpt.prompts.ml_engineer import (
+    GEN_DATA_DESC_PROMPT,
+    UPDATE_DATA_COLUMNS,
+    PRINT_DATA_COLUMNS
+)
 from metagpt.roles import Role
 from metagpt.schema import Message, Plan
-from metagpt.utils.common import CodeParser
+from metagpt.utils.common import CodeParser, remove_comments, create_func_config
 from metagpt.actions.debug_code import DebugCode
 
 STRUCTURAL_CONTEXT = """
@@ -57,34 +62,6 @@ def remove_escape_and_color_codes(input_str):
     return result
 
 
-def read_data(file: str) -> pd.DataFrame:
-    if file.endswith(".csv"):
-        df = pd.read_csv(file, sep=",")
-        sep_list = [";", "\t", ":", " ", "|"]
-        for sep in sep_list:
-            if df.shape[1] == 1:
-                df = pd.read_csv(file, sep=sep)
-            else:
-                break
-    else:
-        raise ValueError(f"Unsupported file type: {file}")
-    return df
-
-
-def get_column_info(df: pd.DataFrame) -> str:
-    data = []
-    for i in df.columns:
-        nan_freq = float("%.2g" % (df[i].isna().mean() * 100))
-        n_unique = df[i].nunique()
-        data.append([i, df[i].dtype, nan_freq, n_unique])
-    
-    samples = pd.DataFrame(
-        data,
-        columns=["Column_name", "Data_type", "NaN_Frequency(%)", "N_unique"],
-    )
-    return samples.to_string(index=False)
-
-
 class AskReview(Action):
     async def run(self, context: List[Message], plan: Plan = None):
         logger.info("Current overall plan:")
@@ -108,26 +85,20 @@ class AskReview(Action):
         return rsp, confirmed
 
 
-class GenerateDataDesc(Action):
-    async def run(self, file: str) -> dict:
-        data_desc = {}
-        df = read_data(file)
-        data_head = df.head().to_dict(orient="list")
-        data_head = json.dumps(data_head, indent=4, ensure_ascii=False)
-        prompt = GEN_DATA_DESC_PROMPT.replace("{data_head}", data_head)
-        rsp = await self._aask(prompt)
-        rsp = CodeParser.parse_code(block=None, text=rsp)
-        rsp = json.loads(rsp)
-        data_desc["path"] = file
-        data_desc["data_desc"] = rsp["data_desc"]
-        data_desc["column_desc"] = rsp["column_desc"]
-        data_desc["column_info"] = get_column_info(df)
-        return data_desc
+class UpdateDataColumns(Action):
+    async def run(self, plan: Plan = None) -> dict:
+        finished_tasks = plan.get_finished_tasks()
+        code_context = [remove_comments(task.code) for task in finished_tasks]
+        code_context = "\n\n".join(code_context)
+        prompt = UPDATE_DATA_COLUMNS.format(history_code=code_context)
+        tool_config = create_func_config(PRINT_DATA_COLUMNS)
+        rsp = await self.llm.aask_code(prompt, **tool_config)
+        return rsp
 
 
 class MLEngineer(Role):
     def __init__(
-            self, name="ABC", profile="MLEngineer", goal="", auto_run: bool = False, data_path: str = None
+            self, name="ABC", profile="MLEngineer", goal="", auto_run: bool = False,
     ):
         super().__init__(name=name, profile=profile, goal=goal)
         self._set_react_mode(react_mode="plan_and_act")
@@ -136,13 +107,9 @@ class MLEngineer(Role):
         self.use_code_steps = True
         self.execute_code = ExecutePyCode()
         self.auto_run = auto_run
-        self.data_path = data_path
         self.data_desc = {}
     
     async def _plan_and_act(self):
-        if self.data_path:
-            self.data_desc = await self._generate_data_desc()
-        
         # create initial plan and update until confirmation
         await self._update_plan()
         
@@ -163,25 +130,27 @@ class MLEngineer(Role):
                 task.code_steps = code_steps
                 self.plan.finish_current_task()
                 self.working_memory.clear()
-                
-                if "print(df_processed.info())" in code:
-                    self.data_desc["column_info"] = result
+
+                success, new_code = await self._update_data_columns()
+                if success:
+                    task.code = task.code + "\n\n" + new_code
             else:
                 # update plan according to user's feedback and to take on changed tasks
                 await self._update_plan()
-        
-        finished_tasks = self.plan.get_finished_tasks()
-        if len(finished_tasks) == len(self.plan.tasks):
-            code_context = [task.code for task in finished_tasks]
-            code_context = "\n\n".join(code_context)
-            result, success = await self.execute_code.run(code_context)
-            # truncated the result
-            print(truncate(result))
-    
-    async def _generate_data_desc(self):
-        data_desc = await GenerateDataDesc().run(self.data_path)
-        return data_desc
-    
+
+        time = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
+        self.execute_code.save_notebook(f"{DATA_PATH}/notebooks/ml_{time}.ipynb")
+
+    async def _update_data_columns(self):
+        rsp = await UpdateDataColumns().run(self.plan)
+        is_update, code = rsp["is_update"], rsp["code"]
+        success = False
+        if is_update:
+            result, success = await self.execute_code.run(code)
+            if success:
+                self.data_desc["column_info"] = result
+        return success, code
+
     async def _write_and_exec_code(self, max_retry: int = 3):
         code_steps = (
             await WriteCodeSteps().run(self.plan)
@@ -192,6 +161,7 @@ class MLEngineer(Role):
         counter = 0
         improve_code = ""
         success = False
+        debug_context = []
         
         finished_tasks = self.plan.get_finished_tasks()
         code_context = [task.code for task in finished_tasks]
@@ -200,37 +170,38 @@ class MLEngineer(Role):
         code_result = "\n\n".join(code_result)
         
         while not success and counter < max_retry:
-            if counter == 0:
-                context = self.get_useful_memories()
-            else:
-                # context = self.get_useful_memories()
-                # logger.info(f"context {context}")
+            context = self.get_useful_memories()
+
+            if counter > 0:
                 improve_code = await DebugCode().run(plan=self.plan.current_task.instruction,
-                                                     finished_code=code_context,
-                                                     finished_code_result=code_result,
+                                                     # finished_code=code_context,
+                                                     # finished_code_result=code_result,
                                                      code=code,
-                                                     runtime_result=self.working_memory.get())
-            
-            if not self.use_tools or self.plan.current_task.task_type == "other":
+                                                     runtime_result=self.working_memory.get(),
+                                                     context=debug_context)
+
+            if improve_code != "":
+                code = improve_code
+                logger.info(f"new code \n{improve_code}")
+                cause_by = DebugCode
+            elif not self.use_tools or self.plan.current_task.task_type == "other":
                 logger.info("Write code with pure generation")
-                
                 code = await WriteCodeByGenerate().run(
                     context=context, plan=self.plan, code_steps=code_steps, temperature=0.0
                 )
+                debug_context = [self.get_useful_memories(task_exclude_field={'result', 'code_steps'})[0]]
                 cause_by = WriteCodeByGenerate
             else:
                 logger.info("Write code with tools")
-                
-                if improve_code != "":
-                    code = improve_code
-                    logger.info(f"new code {code}")
-                    cause_by = DebugCode
-                else:
-                    code = await WriteCodeWithTools().run(
-                        context=context, plan=self.plan, code_steps=code_steps, **{"column_names": {}}
-                    )
-                    
-                    cause_by = WriteCodeWithTools
+                schema_path = PROJECT_ROOT / "metagpt/tools/functions/schemas"
+                tool_context, code = await WriteCodeWithTools(schema_path=schema_path).run(
+                    context=context,
+                    plan=self.plan,
+                    code_steps=code_steps,
+                    column_info=self.data_desc.get("column_info", ""),
+                )
+                debug_context = tool_context
+                cause_by = WriteCodeWithTools
             
             self.working_memory.add(
                 Message(content=code, role="assistant", cause_by=cause_by)
@@ -238,9 +209,7 @@ class MLEngineer(Role):
             
             # debug on code, run on runcode with finished code and new_df
             # runcode = code_context + "\n\n" + code
-            runcode = code
-            
-            result, success = await self.execute_code.run(runcode)
+            result, success = await self.execute_code.run(code)
             # truncated the result
             print(truncate(result))
             
@@ -289,12 +258,12 @@ class MLEngineer(Role):
         self.plan.add_tasks(tasks)
         self.working_memory.clear()
     
-    def get_useful_memories(self) -> List[Message]:
+    def get_useful_memories(self, task_exclude_field: set = None) -> List[Message]:
         """find useful memories only to reduce context length and improve performance"""
         # TODO dataset description , code steps
         user_requirement = self.plan.goal
         tasks = json.dumps(
-            [task.dict() for task in self.plan.tasks], indent=4, ensure_ascii=False
+            [task.dict(exclude=task_exclude_field) for task in self.plan.tasks], indent=4, ensure_ascii=False
         )
         current_task = self.plan.current_task.json() if self.plan.current_task else {}
         context = STRUCTURAL_CONTEXT.format(
@@ -321,12 +290,13 @@ if __name__ == "__main__":
     
     # requirement = "Perform data analysis on the provided data. Train a model to predict the target variable Survived. Include data preprocessing, feature engineering, and modeling in your pipeline. The metric is accuracy."
     
-    data_path = f"{DATA_PATH}/titanic"
-    requirement = f"This is a titanic passenger survival dataset, your goal is to predict passenger survival outcome. The target column is Survived. Perform data analysis, data preprocessing, feature engineering, and modeling to predict the target. Report accuracy on the eval data. Train data path: '{data_path}/split_train.csv', eval data path: '{data_path}/split_eval.csv'."
-    
-    
-    async def main(requirement: str = requirement, auto_run: bool = True, data_path: str = ""):
-        role = MLEngineer(goal=requirement, auto_run=auto_run, data_path=data_path)
+    # data_path = f"{DATA_PATH}/titanic"
+    # requirement = f"This is a titanic passenger survival dataset, your goal is to predict passenger survival outcome. The target column is Survived. Perform data analysis, data preprocessing, feature engineering, and modeling to predict the target. Report accuracy on the eval data. Train data path: '{data_path}/split_train.csv', eval data path: '{data_path}/split_eval.csv'."
+    # requirement = f"Run data analysis on sklearn Wine recognition dataset, include a plot, and train a model to predict wine class (20% as validation), and show validation accuracy"
+    data_path = f"{DATA_PATH}/icr-identify-age-related-conditions"
+    requirement = f"This is a medical dataset with over fifty anonymized health characteristics linked to three age-related conditions. Your goal is to predict whether a subject has or has not been diagnosed with one of these conditions.The target column is Class. Perform data analysis, data preprocessing, feature engineering, and modeling to predict the target. Report f1 score on the eval data. Train data path: {data_path}/split_train.csv, eval data path: {data_path}/split_eval.csv."
+    async def main(requirement: str = requirement, auto_run: bool = True):
+        role = MLEngineer(goal=requirement, auto_run=auto_run)
         await role.run(requirement)
     
     

From ea0b93d2b94997db94f470dbd3141f3f9dd435a6 Mon Sep 17 00:00:00 2001
From: stellahsr <stellahsr@126.com>
Date: Wed, 13 Dec 2023 14:33:50 +0800
Subject: [PATCH 27/49] update code locally

---
 metagpt/actions/write_code_steps.py | 12 ++++++++----
 metagpt/roles/ml_engineer.py        | 26 +++++++++++++-------------
 2 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/metagpt/actions/write_code_steps.py b/metagpt/actions/write_code_steps.py
index a19549b71..6bf223701 100644
--- a/metagpt/actions/write_code_steps.py
+++ b/metagpt/actions/write_code_steps.py
@@ -63,18 +63,22 @@ class WriteCodeSteps(Action):
 
     def get_context(self, plan: Plan):
         user_requirement = plan.goal
-        select_task_keys = ['task_id', 'instruction', 'is_finished', 'code']
-
+        # select_task_keys = ['task_id', 'instruction', 'is_finished', 'code']
+        select_task_keys = ['task_id','code']
+        
         def process_task(task):
             task_dict = task.dict()
-            ptask = {k: task_dict[k] for k in task_dict if k in select_task_keys}
+            ptask = {k: task_dict[k] for k in task_dict if k in select_task_keys }
             return ptask
+        
         tasks = json.dumps(
-            [process_task(task) for task in plan.tasks], indent=4, ensure_ascii=False
+            [process_task(task) for task in plan.tasks if task.is_finished==True], indent=4, ensure_ascii=False
         )
+        
         current_task = json.dumps(process_task(plan.current_task)) if plan.current_task else {}
         context = STRUCTURAL_CONTEXT.format(
             user_requirement=user_requirement, tasks=tasks, current_task=current_task
         )
+        print(context)
         # print(context)
         return context
diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index 8ad75b399..f50b6d494 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -148,7 +148,7 @@ class MLEngineer(Role):
         
         while self.plan.current_task:
             task = self.plan.current_task
-            logger.info(f"ready to take on task {task}")
+            logger.info(f"ready to take on task: {task}")
             
             # take on current task
             code, result, success, code_steps = await self._write_and_exec_code()
@@ -157,9 +157,11 @@ class MLEngineer(Role):
             task_result_confirmed = await self._ask_review()
             
             # 针对当前task进行单独plan
-            if not success or not task_result_confirmed:
-                # fixme: 增加对应plan
-                self.state.plan()
+            # if not success or not task_result_confirmed:
+            #     # fixme: 增加对应plan
+            #     logger.info(task.result)
+            #     # import pdb;pdb.set_trace()
+            # #     self.state.plan()
                 
             if success and task_result_confirmed:
                 # tick off this task and record progress
@@ -175,13 +177,13 @@ class MLEngineer(Role):
                 # update plan according to user's feedback and to take on changed tasks
                 await self._update_plan()
         
-        finished_tasks = self.plan.get_finished_tasks()
-        if len(finished_tasks) == len(self.plan.tasks):
-            code_context = [task.code for task in finished_tasks]
-            code_context = "\n\n".join(code_context)
-            result, success = await self.execute_code.run(code_context)
-            # truncated the result
-            print(truncate(result))
+        # finished_tasks = self.plan.get_finished_tasks()
+        # if len(finished_tasks) == len(self.plan.tasks):
+        #     code_context = [task.code for task in finished_tasks]
+        #     code_context = "\n\n".join(code_context)
+        #     result, success = await self.execute_code.run(code_context)
+        #     # truncated the result
+        #     print(truncate(result))
     
     async def _generate_data_desc(self):
         data_desc = await GenerateDataDesc().run(self.data_path)
@@ -258,8 +260,6 @@ class MLEngineer(Role):
             
             counter += 1
         
-        success = False
-        
         return code, result, success, code_steps
     
     async def _ask_review(self):

From abad52da85773c3f762eea7f7c956140e0c0cd3f Mon Sep 17 00:00:00 2001
From: stellahsr <stellahsr@126.com>
Date: Wed, 13 Dec 2023 15:38:19 +0800
Subject: [PATCH 28/49] update locally

---
 metagpt/actions/write_code_steps.py | 48 ++++++++++++++++++++++++++---
 metagpt/roles/ml_engineer.py        |  8 ++---
 2 files changed, 48 insertions(+), 8 deletions(-)

diff --git a/metagpt/actions/write_code_steps.py b/metagpt/actions/write_code_steps.py
index 889c06679..efee96749 100644
--- a/metagpt/actions/write_code_steps.py
+++ b/metagpt/actions/write_code_steps.py
@@ -6,6 +6,31 @@ from metagpt.actions import Action
 from metagpt.schema import Message, Task, Plan
 from metagpt.utils.common import CodeParser
 
+# CODE_STEPS_PROMPT_TEMPLATE = """
+# # Context
+# {context}
+#
+# -----
+# Tasks are all code development tasks.
+# You are a professional engineer, the main goal is to plan out concise solution steps for Current Task before coding.
+# A planning process can reduce the difficulty and improve the quality of coding.
+# You may be given some code plans for the tasks ahead, but you don't have to follow the existing plan when planning the current task.
+# The output plan should following the subsequent principles:
+# 1.The plan is a rough checklist of steps outlining the entire program's structure.Try to keep the number of steps fewer than 5.
+# 2.The steps should be written concisely and at a high level, avoiding overly detailed implementation specifics.
+# 3.The execution of the plan happens sequentially, but the plan can incorporate conditional (if) and looping(loop) keywords for more complex structures.
+#
+# Output the code steps in a JSON format, as shown in this example:
+# ```json
+# {
+#     "Step 1": "",
+#     "Step 2": "",
+#     "Step 3": "",
+#     ...
+# }
+# ```
+# """
+
 CODE_STEPS_PROMPT_TEMPLATE = """
 # Context
 {context}
@@ -19,6 +44,7 @@ The output plan should following the subsequent principles:
 1.The plan is a rough checklist of steps outlining the entire program's structure.Try to keep the number of steps fewer than 5.
 2.The steps should be written concisely and at a high level, avoiding overly detailed implementation specifics.
 3.The execution of the plan happens sequentially, but the plan can incorporate conditional (if) and looping(loop) keywords for more complex structures.
+4.Follow the code logic to design and provide the code steps. You can analysis it step by step
 
 Output the code steps in a JSON format, as shown in this example:
 ```json
@@ -31,11 +57,22 @@ Output the code steps in a JSON format, as shown in this example:
 ```
 """
 
+# STRUCTURAL_CONTEXT = """
+# ## User Requirement
+# {user_requirement}
+# ## Current Plan
+# {tasks}
+# ## Current Task
+# {current_task}
+# """
+
 STRUCTURAL_CONTEXT = """
 ## User Requirement
 {user_requirement}
-## Current Plan
+## Plan
 {tasks}
+## Codes
+{codes}
 ## Current Task
 {current_task}
 """
@@ -63,21 +100,24 @@ class WriteCodeSteps(Action):
 
     def get_context(self, plan: Plan):
         user_requirement = plan.goal
-        select_task_keys = ['task_id', 'instruction', 'is_finished', 'code']
-        # select_task_keys = ['task_id','code']
+        # select_task_keys = ['task_id', 'instruction', 'is_finished', 'code']
+        select_task_keys = ['task_id','instruction']
         
         def process_task(task):
             task_dict = task.dict()
             ptask = {k: task_dict[k] for k in task_dict if k in select_task_keys }
             return ptask
         
+        
         tasks = json.dumps(
             [process_task(task) for task in plan.tasks], indent=4, ensure_ascii=False
         )
         
+        code_lists = [task.code for task in plan.tasks if task.is_finished==True]
+        codes = "\n\n".join(code_lists)
         current_task = json.dumps(process_task(plan.current_task)) if plan.current_task else {}
         context = STRUCTURAL_CONTEXT.format(
-            user_requirement=user_requirement, tasks=tasks, current_task=current_task
+            user_requirement=user_requirement, tasks=tasks, codes=codes, current_task=current_task
         )
         print(context)
         # print(context)
diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index 20589079d..c735eb983 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -290,11 +290,11 @@ if __name__ == "__main__":
     
     # requirement = "Perform data analysis on the provided data. Train a model to predict the target variable Survived. Include data preprocessing, feature engineering, and modeling in your pipeline. The metric is accuracy."
     
-    # data_path = f"{DATA_PATH}/titanic"
-    # requirement = f"This is a titanic passenger survival dataset, your goal is to predict passenger survival outcome. The target column is Survived. Perform data analysis, data preprocessing, feature engineering, and modeling to predict the target. Report accuracy on the eval data. Train data path: '{data_path}/split_train.csv', eval data path: '{data_path}/split_eval.csv'."
+    data_path = f"{DATA_PATH}/titanic"
+    requirement = f"This is a titanic passenger survival dataset, your goal is to predict passenger survival outcome. The target column is Survived. Perform data analysis, data preprocessing, feature engineering, and modeling to predict the target. Report accuracy on the eval data. Train data path: '{data_path}/split_train.csv', eval data path: '{data_path}/split_eval.csv'."
     # requirement = f"Run data analysis on sklearn Wine recognition dataset, include a plot, and train a model to predict wine class (20% as validation), and show validation accuracy"
-    data_path = f"{DATA_PATH}/icr-identify-age-related-conditions"
-    requirement = f"This is a medical dataset with over fifty anonymized health characteristics linked to three age-related conditions. Your goal is to predict whether a subject has or has not been diagnosed with one of these conditions.The target column is Class. Perform data analysis, data preprocessing, feature engineering, and modeling to predict the target. Report f1 score on the eval data. Train data path: {data_path}/split_train.csv, eval data path: {data_path}/split_eval.csv."
+    # data_path = f"{DATA_PATH}/icr-identify-age-related-conditions"
+    # requirement = f"This is a medical dataset with over fifty anonymized health characteristics linked to three age-related conditions. Your goal is to predict whether a subject has or has not been diagnosed with one of these conditions.The target column is Class. Perform data analysis, data preprocessing, feature engineering, and modeling to predict the target. Report f1 score on the eval data. Train data path: {data_path}/split_train.csv, eval data path: {data_path}/split_eval.csv."
     async def main(requirement: str = requirement, auto_run: bool = True):
         role = MLEngineer(goal=requirement, auto_run=auto_run)
         await role.run(requirement)

From 8d694d47d9f2372011d39f759042b48cc54c8c27 Mon Sep 17 00:00:00 2001
From: stellahsr <stellahsr@126.com>
Date: Wed, 13 Dec 2023 16:26:24 +0800
Subject: [PATCH 29/49] update code step prompts

---
 metagpt/actions/write_analysis_code.py | 57 ++++++++++++++------------
 metagpt/actions/write_code_steps.py    |  7 ++--
 metagpt/prompts/ml_engineer.py         |  2 +-
 3 files changed, 35 insertions(+), 31 deletions(-)

diff --git a/metagpt/actions/write_analysis_code.py b/metagpt/actions/write_analysis_code.py
index aceebbfeb..3e91f4b14 100644
--- a/metagpt/actions/write_analysis_code.py
+++ b/metagpt/actions/write_analysis_code.py
@@ -26,7 +26,7 @@ from metagpt.utils.common import create_func_config, remove_comments
 class BaseWriteAnalysisCode(Action):
     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."""  # 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 process_msg(self, prompt: Union[str, List[Dict], Message, List[Message]], system_msg: str = None):
         default_system_msg = system_msg or self.DEFAULT_SYSTEM_MSG
         # 全部转成list
@@ -45,7 +45,7 @@ class BaseWriteAnalysisCode(Action):
                     messages.append(p.to_dict())
                 elif isinstance(p.content, dict) and "code" in p.content:
                     messages.append(p.content["code"])
-
+        
         # 添加默认的提示词
         if (
                 default_system_msg not in messages[0]["content"]
@@ -61,7 +61,7 @@ class BaseWriteAnalysisCode(Action):
                 "content": messages[0]["content"] + default_system_msg,
             }
         return messages
-
+    
     async def run(
             self, context: List[Message], plan: Plan = None, code_steps: str = ""
     ) -> str:
@@ -79,10 +79,10 @@ class BaseWriteAnalysisCode(Action):
 
 class WriteCodeByGenerate(BaseWriteAnalysisCode):
     """Write code fully by generation"""
-
+    
     def __init__(self, name: str = "", context=None, llm=None) -> str:
         super().__init__(name, context, llm)
-
+    
     async def run(
             self,
             context: [List[Message]],
@@ -99,15 +99,15 @@ class WriteCodeByGenerate(BaseWriteAnalysisCode):
 
 class WriteCodeWithTools(BaseWriteAnalysisCode):
     """Write code with help of local available tools. Choose tools first, then generate code to use the tools"""
-
+    
     def __init__(self, name: str = "", context=None, llm=None, schema_path=None):
         super().__init__(name, context, llm)
         self.schema_path = schema_path
         self.available_tools = {}
-
+        
         if self.schema_path is not None:
             self._load_tools(schema_path)
-
+    
     def _load_tools(self, schema_path):
         """Load tools from yaml file"""
         yml_files = schema_path.glob("*.yml")
@@ -115,7 +115,7 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
             module = yml_file.stem
             with open(yml_file, "r", encoding="utf-8") as f:
                 self.available_tools[module] = yaml.safe_load(f)
-
+    
     def _parse_recommend_tools(self, module: str, recommend_tools: list) -> dict:
         """
         Parses and validates a list of recommended tools, and retrieves their schema from registry.
@@ -132,15 +132,15 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
         for tool in recommend_tools:
             if tool in available_tools:
                 valid_tools.append(tool)
-
+        
         tool_catalog = {tool: self.available_tools[module][tool] for tool in valid_tools}
         return tool_catalog
-
+    
     async def _tool_recommendation(
-        self,
-        task: str,
-        code_steps: str,
-        available_tools: dict,
+            self,
+            task: str,
+            code_steps: str,
+            available_tools: dict,
     ) -> list:
         """
         Recommend tools for the specified task.
@@ -162,26 +162,26 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
         rsp = await self.llm.aask_code(prompt, **tool_config)
         recommend_tools = rsp["recommend_tools"]
         return recommend_tools
-
+    
     async def run(
-        self,
-        context: List[Message],
-        plan: Plan = None,
-        code_steps: str = "",
-        column_info: str = "",
-        **kwargs,
+            self,
+            context: List[Message],
+            plan: Plan = None,
+            code_steps: str = "",
+            column_info: str = "",
+            **kwargs,
     ) -> Tuple[List[Message], str]:
         task_type = plan.current_task.task_type
         available_tools = self.available_tools.get(task_type, {})
         special_prompt = ML_SPECIFIC_PROMPT.get(task_type, "")
-
+        
         finished_tasks = plan.get_finished_tasks()
         code_context = [remove_comments(task.code) for task in finished_tasks]
         code_context = "\n\n".join(code_context)
-
+        
         if len(available_tools) > 0:
             available_tools = {k: v["description"] for k, v in available_tools.items()}
-
+            
             recommend_tools = await self._tool_recommendation(
                 plan.current_task.instruction,
                 code_steps,
@@ -189,8 +189,9 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
             )
             tool_catalog = self._parse_recommend_tools(task_type, recommend_tools)
             logger.info(f"Recommended tools: \n{recommend_tools}")
-
+            
             module_name = ML_MODULE_MAP[task_type]
+            
             prompt = TOOL_USAGE_PROMPT.format(
                 user_requirement=plan.goal,
                 history_code=code_context,
@@ -201,6 +202,8 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
                 module_name=module_name,
                 tool_catalog=tool_catalog,
             )
+        
+        
         else:
             prompt = GENERATE_CODE_PROMPT.format(
                 user_requirement=plan.goal,
@@ -210,7 +213,7 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
                 special_prompt=special_prompt,
                 code_steps=code_steps,
             )
-
+        
         tool_config = create_func_config(CODE_GENERATOR_WITH_TOOLS)
         rsp = await self.llm.aask_code(prompt, **tool_config)
         context = [Message(content=prompt, role="user")]
diff --git a/metagpt/actions/write_code_steps.py b/metagpt/actions/write_code_steps.py
index efee96749..9e06bc91e 100644
--- a/metagpt/actions/write_code_steps.py
+++ b/metagpt/actions/write_code_steps.py
@@ -44,7 +44,7 @@ The output plan should following the subsequent principles:
 1.The plan is a rough checklist of steps outlining the entire program's structure.Try to keep the number of steps fewer than 5.
 2.The steps should be written concisely and at a high level, avoiding overly detailed implementation specifics.
 3.The execution of the plan happens sequentially, but the plan can incorporate conditional (if) and looping(loop) keywords for more complex structures.
-4.Follow the code logic to design and provide the code steps. You can analysis it step by step
+4.Design and provide code steps by following the code logic. Analyze the provided code step by step and reuse the imported library.
 
 Output the code steps in a JSON format, as shown in this example:
 ```json
@@ -101,11 +101,12 @@ class WriteCodeSteps(Action):
     def get_context(self, plan: Plan):
         user_requirement = plan.goal
         # select_task_keys = ['task_id', 'instruction', 'is_finished', 'code']
-        select_task_keys = ['task_id','instruction']
+        # select_task_keys = ['task_id','instruction']
         
         def process_task(task):
             task_dict = task.dict()
-            ptask = {k: task_dict[k] for k in task_dict if k in select_task_keys }
+            # ptask = {k: task_dict[k] for k in task_dict if k in select_task_keys }
+            ptask = f"task_id_{task_dict['task_id']}:{task_dict['instruction']}\n"
             return ptask
         
         
diff --git a/metagpt/prompts/ml_engineer.py b/metagpt/prompts/ml_engineer.py
index d11cbf453..2d2d3315a 100644
--- a/metagpt/prompts/ml_engineer.py
+++ b/metagpt/prompts/ml_engineer.py
@@ -231,8 +231,8 @@ for col in num_cols:
 # Constraints:
 - Prioritize using pre-defined tools for the same functionality.
 - Copy DataFrame before processing if needed.
-- If 'Code Steps' contains step done in 'Done Tasks', such as reading data, don't repeat it.
 """
+#- If 'Code Steps' contains step done in 'Done Tasks', such as reading data, don't repeat it.
 
 DATA_PREPROCESS_PROMPT = """
 The current task is about data preprocessing, please note the following:

From abbaa6afa95e7fcada42df8a299f1dd3a7cc97c5 Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Wed, 13 Dec 2023 17:03:56 +0800
Subject: [PATCH 30/49] refine prompt

---
 metagpt/prompts/ml_engineer.py | 36 ++++++++++++++++++++--------------
 1 file changed, 21 insertions(+), 15 deletions(-)

diff --git a/metagpt/prompts/ml_engineer.py b/metagpt/prompts/ml_engineer.py
index 2d2d3315a..f2412c35b 100644
--- a/metagpt/prompts/ml_engineer.py
+++ b/metagpt/prompts/ml_engineer.py
@@ -155,46 +155,51 @@ PRINT_DATA_COLUMNS = {
 
 GENERATE_CODE_PROMPT = """
 # Background
-Assist in completing [{user_requirement}] in a Jupyter notebook.
+As a data scientist, you need to help user to achieve their goal [{user_requirement}] step-by-step in an continuous Jupyter notebook.
 
-## Task Progress
-### Done Tasks
+## Done Tasks
 ```python
 {history_code}
 ```end
 
-### Current Task
+## Current Task
 {current_task}
 
-## Latest Data Info
+# Latest Data Info
+Latest data info after previous tasks:
 {column_info}
 
 # Task
-Fully implement 'Current Task', ensuring all necessary steps are covered without repeating code from 'Done Tasks'. Specifically, {special_prompt}
+Write complete code for 'Current Task'. And avoid duplicating code from 'Done Tasks', such as repeated import of packages, reading data, etc.
+Specifically, {special_prompt}
 
 # Code Steps:
 Follow steps below when you writing code if it's convenient.
 {code_steps}
+
+# Constraints:
+- Ensure the output new code is executable in the same Jupyter notebook with previous tasks code have been executed.
 """
 
 TOOL_USAGE_PROMPT = """
 # Background
-Assist in completing [{user_requirement}] in a Jupyter notebook.
+As a data scientist, you need to help user to achieve their goal [{user_requirement}] step-by-step in an continuous Jupyter notebook.
 
-## Task Progress
-### Done Tasks
+## Done Tasks
 ```python
 {history_code}
 ```end
 
-### Current Task
+## Current Task
 {current_task}
 
-## Latest Data Info
+# Latest Data Info
+Latest data info after previous tasks:
 {column_info}
 
 # Task
-Fully implement 'Current Task', ensuring all necessary steps are covered without repeating code from 'Done Tasks'. Specifically, {special_prompt}
+Write complete code for 'Current Task'. And avoid duplicating code from 'Done Tasks', such as repeated import of packages, reading data, etc.
+Specifically, {special_prompt}
 
 # Code Steps:
 Follow steps below when you writing code if it's convenient.
@@ -205,11 +210,11 @@ Follow steps below when you writing code if it's convenient.
 - You can freely combine the use of any other public packages, like sklearn, numpy, pandas, etc..
 
 # Available Tools:
-Each Class tool is described in JSON format. When you call it, import the tool from `{module_name}` first.
+Each Class tool is described in JSON format. When you call a tool, import the tool from `{module_name}` first.
 {tool_catalog}
 
 # Output Example:
-For "fill missing value and handle outliers", the output code be like when there are training data and test data:
+when current task is "fill missing value and handle outliers", and their are training data and test data, the output code be like:
 ```python
 # Tools used: ['FillMissingValue']
 from metagpt.tools.functions.libs.data_preprocess import FillMissingValue
@@ -229,8 +234,9 @@ for col in num_cols:
 ```end
 
 # Constraints:
+- Ensure the output new code is executable in the same Jupyter notebook with previous tasks code have been executed.
 - Prioritize using pre-defined tools for the same functionality.
-- Copy DataFrame before processing if needed.
+- Always copy the DataFrame before processing it and use the copy to process.
 """
 #- If 'Code Steps' contains step done in 'Done Tasks', such as reading data, don't repeat it.
 

From 4423524734b15fdb9ca8aafb5eefa823d70ba671 Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Wed, 13 Dec 2023 18:11:54 +0800
Subject: [PATCH 31/49] fix schema

---
 .../tools/functions/schemas/feature_engineering.yml    | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/metagpt/tools/functions/schemas/feature_engineering.yml b/metagpt/tools/functions/schemas/feature_engineering.yml
index 4f2a7100d..3ba9e863b 100644
--- a/metagpt/tools/functions/schemas/feature_engineering.yml
+++ b/metagpt/tools/functions/schemas/feature_engineering.yml
@@ -53,17 +53,17 @@ PolynomialExpansion:
 
 CatCount:
   type: class
-  description: "Add value counts of categorical columns as new features."
+  description: "Add value counts of a categorical column as new feature."
   methods:
     __init__:
       description: "Initialize self."
       parameters:
         properties:
-          cols:
-            type: list
-            description: "Columns for value counts."
+          col:
+            type: str
+            description: "Column for value counts."
         required:
-          - cols
+          - col
     fit:
       description: "Fit the CatCount model."
       parameters:

From e59bab73b06985fd02cc955002372909a0c571aa Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Wed, 13 Dec 2023 19:36:02 +0800
Subject: [PATCH 32/49] refine prompt

---
 metagpt/prompts/ml_engineer.py | 31 ++++++++++++++++++++++++++-----
 metagpt/roles/ml_engineer.py   |  9 +--------
 2 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/metagpt/prompts/ml_engineer.py b/metagpt/prompts/ml_engineer.py
index f2412c35b..05d8db8e9 100644
--- a/metagpt/prompts/ml_engineer.py
+++ b/metagpt/prompts/ml_engineer.py
@@ -174,11 +174,29 @@ Write complete code for 'Current Task'. And avoid duplicating code from 'Done Ta
 Specifically, {special_prompt}
 
 # Code Steps:
-Follow steps below when you writing code if it's convenient.
+Strictly follow steps below when you writing code if it's convenient.
 {code_steps}
 
+# Output Example:
+when current task is "train a lightgbm model on training data", and their are two steps in 'Code Steps', the code be like:
+```python
+# Step 1: check data type and convert to numeric
+ojb_cols = train.select_dtypes(include='object').columns.tolist()
+
+for col in obj_cols:
+    encoder = LabelEncoder()
+    train[col] = encoder.fit_transform(train[col])
+    test[col] = test[col].apply(lambda x: x if x in encoder.classes_ else 'unknown')
+    test[col] = encoder.transform(test[col])
+
+# Step 2: train lightgbm model
+model = LGBMClassifier()
+model.fit(train, y_train)
+```end
+
 # Constraints:
 - Ensure the output new code is executable in the same Jupyter notebook with previous tasks code have been executed.
+- The output code should contain all steps implemented in 'Code Steps'.
 """
 
 TOOL_USAGE_PROMPT = """
@@ -202,7 +220,7 @@ Write complete code for 'Current Task'. And avoid duplicating code from 'Done Ta
 Specifically, {special_prompt}
 
 # Code Steps:
-Follow steps below when you writing code if it's convenient.
+Strictly follow steps below when you writing code if it's convenient.
 {code_steps}
 
 # Capabilities
@@ -214,8 +232,9 @@ Each Class tool is described in JSON format. When you call a tool, import the to
 {tool_catalog}
 
 # Output Example:
-when current task is "fill missing value and handle outliers", and their are training data and test data, the output code be like:
+when current task is "do data preprocess, like fill missing value, handle outliers, etc.", and their are two steps in 'Code Steps', the code be like:
 ```python
+# Step 1: fill missing value
 # Tools used: ['FillMissingValue']
 from metagpt.tools.functions.libs.data_preprocess import FillMissingValue
 
@@ -227,6 +246,7 @@ fill_missing_value.fit(train_processed)
 train_processed = fill_missing_value.transform(train_processed)
 test_processed = fill_missing_value.transform(test_processed)
 
+# Step 2: handle outliers
 for col in num_cols:
     low, high = train_processed[col].quantile([0.01, 0.99])
     train_processed[col] = train_processed[col].clip(low, high)
@@ -235,8 +255,9 @@ for col in num_cols:
 
 # Constraints:
 - Ensure the output new code is executable in the same Jupyter notebook with previous tasks code have been executed.
-- Prioritize using pre-defined tools for the same functionality.
+- Always prioritize using pre-defined tools for the same functionality.
 - Always copy the DataFrame before processing it and use the copy to process.
+- The output code should contain all steps implemented correctly in 'Code Steps'.
 """
 #- If 'Code Steps' contains step done in 'Done Tasks', such as reading data, don't repeat it.
 
@@ -266,7 +287,7 @@ The current task is about training a model, please ensure high performance:
 
 MODEL_EVALUATE_PROMPT = """
 The current task is about evaluating a model, please note the following:
-- Ensure that the evaluated data is same processed as the training data.
+- Ensure that the evaluated data is same processed as the training data. If not, remember use object in 'Done Tasks' to transform the data.
 - Use trained model from previous task result directly, do not mock or reload model yourself.
 """
 
diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index c735eb983..6a2a9e2b0 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -32,13 +32,6 @@ STRUCTURAL_CONTEXT = """
 {tasks}
 ## Current Task
 {current_task}
-## Packages Installed
-scikit-learn
-pandas
-numpy
-lightgbm
-xgboost
-catboost
 """
 
 
@@ -212,7 +205,7 @@ class MLEngineer(Role):
             result, success = await self.execute_code.run(code)
             # truncated the result
             print(truncate(result))
-            
+
             self.working_memory.add(
                 Message(content=truncate(remove_escape_and_color_codes(result)), role="user", cause_by=ExecutePyCode)
             )

From 7e6e493499c41c91c56a19a2ebc7ecb329ab6f5f Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Wed, 13 Dec 2023 19:36:31 +0800
Subject: [PATCH 33/49] refine prompt

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

diff --git a/metagpt/actions/debug_code.py b/metagpt/actions/debug_code.py
index 53ca2f54d..58d006a08 100644
--- a/metagpt/actions/debug_code.py
+++ b/metagpt/actions/debug_code.py
@@ -47,7 +47,7 @@ REFLECTION_PROMPT = """
                        [runtime Error]
                        {runtime_result}
 
-                        Analysis the error step by step, provide me improve method and code. Remember to follow [context] requirement.
+                        Analysis the error step by step, provide me improve method and code. Remember to follow [context] rerquirement. Don't forget write code for steps behind the error step.
                         [reflection on previous impl]:
                         xxx
 

From cfb577d6747ba7dca7cea92b7199494a66eb3dfb Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Wed, 13 Dec 2023 20:10:17 +0800
Subject: [PATCH 34/49] rollback config

---
 config/config.yaml | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/config/config.yaml b/config/config.yaml
index 694251f17..17605307a 100644
--- a/config/config.yaml
+++ b/config/config.yaml
@@ -5,7 +5,7 @@
 ## The official OPENAI_API_BASE is https://api.openai.com/v1
 ## If the official OPENAI_API_BASE is not available, we recommend using the [openai-forward](https://github.com/beidongjiedeguang/openai-forward).
 ## Or, you can configure OPENAI_PROXY to access official OPENAI_API_BASE.
-#OPENAI_API_BASE: "https://api.openai.com/v1"
+OPENAI_API_BASE: "https://api.openai.com/v1"
 #OPENAI_PROXY: "http://127.0.0.1:8118"
 #OPENAI_API_KEY: "YOUR_API_KEY"   # set the value to sk-xxx if you host the openai interface for open llm model
 OPENAI_API_MODEL: "gpt-4"
@@ -24,13 +24,12 @@ RPM: 10
 
 #### if AZURE, check https://github.com/openai/openai-cookbook/blob/main/examples/azure/chat.ipynb
 #### You can use ENGINE or DEPLOYMENT mode
-OPENAI_API_TYPE: "azure"
-OPENAI_API_BASE: "https://deepwisdom.openai.azure.com/"
-OPENAI_API_KEY: "02ae6058d09849c691176befeae2107c"
-#OPENAI_API_VERSION: "2023-05-15"
-OPENAI_API_VERSION: "2023-07-01-preview"
-DEPLOYMENT_ID: "GPT-4"
-OPENAI_API_ENGINE: "gpt-4"
+#OPENAI_API_TYPE: "azure"
+#OPENAI_API_BASE: "YOUR_AZURE_ENDPOINT"
+#OPENAI_API_KEY: "YOUR_AZURE_API_KEY"
+#OPENAI_API_VERSION: "YOUR_AZURE_API_VERSION"
+#DEPLOYMENT_NAME: "YOUR_DEPLOYMENT_NAME"
+#DEPLOYMENT_ID: "YOUR_DEPLOYMENT_ID"
 
 #### if zhipuai from `https://open.bigmodel.cn`. You can set here or export API_KEY="YOUR_API_KEY"
 # ZHIPUAI_API_KEY: "YOUR_API_KEY"
@@ -88,7 +87,7 @@ SD_T2I_API: "/sdapi/v1/txt2img"
 MODEL_FOR_RESEARCHER_SUMMARY: gpt-3.5-turbo
 MODEL_FOR_RESEARCHER_REPORT: gpt-3.5-turbo-16k
 
-### choose the engine for mermaid conversion, 
+### choose the engine for mermaid conversion,
 # default is nodejs, you can change it to playwright,pyppeteer or ink
 # MERMAID_ENGINE: nodejs
 

From 8b0b5eeb804402f6a5329b92cdcb6da9e387d59d Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Wed, 13 Dec 2023 20:14:10 +0800
Subject: [PATCH 35/49] fix conflict

---
 metagpt/actions/write_code_steps.py |  1 -
 metagpt/roles/ml_engineer.py        | 48 ++++++++++++-----------------
 2 files changed, 19 insertions(+), 30 deletions(-)

diff --git a/metagpt/actions/write_code_steps.py b/metagpt/actions/write_code_steps.py
index 9e06bc91e..3c08adc19 100644
--- a/metagpt/actions/write_code_steps.py
+++ b/metagpt/actions/write_code_steps.py
@@ -120,6 +120,5 @@ class WriteCodeSteps(Action):
         context = STRUCTURAL_CONTEXT.format(
             user_requirement=user_requirement, tasks=tasks, codes=codes, current_task=current_task
         )
-        print(context)
         # print(context)
         return context
diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index 26dfdbc67..8ab3ac981 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -4,30 +4,26 @@ from datetime import datetime
 
 import fire
 
-from metagpt.roles import Role
-from metagpt.schema import Message, Plan
-from metagpt.memory import Memory
-from metagpt.logs import logger
 from metagpt.actions import Action
-from metagpt.actions.write_plan import WritePlan, update_plan_from_rsp, precheck_update_plan_from_rsp
-from metagpt.actions.write_analysis_code import WriteCodeByGenerate, WriteCodeWithTools
-from metagpt.actions.ml_da_action import AskReview, SummarizeAnalysis, Reflect, ReviewConst
+from metagpt.actions.debug_code import DebugCode
 from metagpt.actions.execute_code import ExecutePyCode
-from metagpt.roles.kaggle_manager import DownloadData, SubmitResult
-from metagpt.prompts.ml_engineer import STRUCTURAL_CONTEXT
+from metagpt.actions.ml_da_action import AskReview, SummarizeAnalysis, Reflect, ReviewConst
+from metagpt.actions.write_analysis_code import WriteCodeByGenerate, WriteCodeWithTools
 from metagpt.actions.write_code_steps import WriteCodeSteps
 from metagpt.actions.write_plan import WritePlan
+from metagpt.actions.write_plan import update_plan_from_rsp, precheck_update_plan_from_rsp
 from metagpt.const import DATA_PATH, PROJECT_ROOT
 from metagpt.logs import logger
+from metagpt.memory import Memory
+from metagpt.prompts.ml_engineer import STRUCTURAL_CONTEXT
 from metagpt.prompts.ml_engineer import (
-    GEN_DATA_DESC_PROMPT,
     UPDATE_DATA_COLUMNS,
     PRINT_DATA_COLUMNS
 )
 from metagpt.roles import Role
+from metagpt.roles.kaggle_manager import DownloadData, SubmitResult
 from metagpt.schema import Message, Plan
-from metagpt.utils.common import CodeParser, remove_comments, create_func_config
-from metagpt.actions.debug_code import DebugCode
+from metagpt.utils.common import remove_comments, create_func_config
 from metagpt.utils.save_code import save_code_file
 
 
@@ -103,9 +99,10 @@ class MLEngineer(Role):
                 self.plan.finish_current_task()
                 self.working_memory.clear()
 
-                success, new_code = await self._update_data_columns()
-                if success:
-                    task.code = task.code + "\n\n" + new_code
+                if self.use_tools:
+                    success, new_code = await self._update_data_columns()
+                    if success:
+                        task.code = task.code + "\n\n" + new_code
 
                 confirmed_and_more = (ReviewConst.CONTINUE_WORD[0] in review.lower()
                     and review.lower() not in ReviewConst.CONTINUE_WORD[0])  # "confirm, ... (more content, such as changing downstream tasks)"
@@ -134,9 +131,6 @@ class MLEngineer(Role):
         save_code_file(name=project_record, code_context=self.execute_code.nb, file_format="ipynb")
         return rsp
 
-        time = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
-        self.execute_code.save_notebook(f"{DATA_PATH}/notebooks/ml_{time}.ipynb")
-
     async def _update_data_columns(self):
         rsp = await UpdateDataColumns().run(self.plan)
         is_update, code = rsp["is_update"], rsp["code"]
@@ -159,12 +153,6 @@ class MLEngineer(Role):
         success = False
         debug_context = []
 
-        finished_tasks = self.plan.get_finished_tasks()
-        code_context = [task.code for task in finished_tasks]
-        code_result = [task.result for task in finished_tasks]
-        code_context = "\n\n".join(code_context)
-        code_result = "\n\n".join(code_result)
-
         while not success and counter < max_retry:
             context = self.get_useful_memories()
 
@@ -272,16 +260,18 @@ class MLEngineer(Role):
         self.working_memory.add(Message(content=reflection, role="assistant"))
         self.working_memory.add(Message(content=Reflect.REWRITE_PLAN_INSTRUCTION, role="user"))
 
-    def get_useful_memories(self, task_exclude_field: set = None) -> List[Message]:
+    def get_useful_memories(self, task_exclude_field=None) -> List[Message]:
         """find useful memories only to reduce context length and improve performance"""
         # TODO dataset description , code steps
+        if task_exclude_field is None:
+            task_exclude_field = {'code_steps'}
         user_requirement = self.plan.goal
         data_desc = self.plan.context
         tasks = [task.dict(exclude=task_exclude_field) for task in self.plan.tasks]
-        for task in tasks:
-            # Shorten the context as we don't need code steps after we get the codes.
-            # This doesn't affect current_task below, which should hold the code steps
-            task.pop("code_steps")
+        # for task in tasks:
+        #     # Shorten the context as we don't need code steps after we get the codes.
+        #     # This doesn't affect current_task below, which should hold the code steps
+        #     task.pop("code_steps")
         tasks = json.dumps(tasks, indent=4, ensure_ascii=False)
         current_task = self.plan.current_task.json() if self.plan.current_task else {}
         context = STRUCTURAL_CONTEXT.format(

From 7744815c5ff8f61eb90ccee07555c9f7207182bd Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Wed, 13 Dec 2023 20:32:49 +0800
Subject: [PATCH 36/49] fix conflict

---
 requirements.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/requirements.txt b/requirements.txt
index 9b75fd200..2328de2a1 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -45,6 +45,7 @@ wrapt==1.15.0
 websocket-client==0.58.0
 zhipuai==1.0.7
 rich==13.6.0
+nbclient==0.9.0
 nbformat==5.9.2
 ipython==8.17.2
 ipykernel==6.27.0

From edd6987a1c4738f27fb1936fa701441145b96869 Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Wed, 13 Dec 2023 20:41:32 +0800
Subject: [PATCH 37/49] drop old tool definition

---
 metagpt/tools/functions/__init__.py           |   3 -
 metagpt/tools/functions/libs/ml_model.py      | 196 ------------------
 metagpt/tools/functions/register/__init__.py  |   6 -
 metagpt/tools/functions/register/register.py  |  78 -------
 metagpt/tools/functions/schemas/base.py       | 100 ---------
 .../functions/schemas/data_preprocess.py      |  67 ------
 .../functions/schemas/feature_engineering.py  | 110 ----------
 metagpt/tools/functions/schemas/ml_model.py   |  55 -----
 8 files changed, 615 deletions(-)
 delete mode 100644 metagpt/tools/functions/libs/ml_model.py
 delete mode 100644 metagpt/tools/functions/register/__init__.py
 delete mode 100644 metagpt/tools/functions/register/register.py
 delete mode 100644 metagpt/tools/functions/schemas/base.py
 delete mode 100644 metagpt/tools/functions/schemas/data_preprocess.py
 delete mode 100644 metagpt/tools/functions/schemas/feature_engineering.py
 delete mode 100644 metagpt/tools/functions/schemas/ml_model.py

diff --git a/metagpt/tools/functions/__init__.py b/metagpt/tools/functions/__init__.py
index 30ee10827..a0a43f507 100644
--- a/metagpt/tools/functions/__init__.py
+++ b/metagpt/tools/functions/__init__.py
@@ -4,6 +4,3 @@
 # @Author  : lidanyang
 # @File    : __init__.py
 # @Desc    :
-from metagpt.tools.functions.register.register import registry
-import metagpt.tools.functions.libs.feature_engineering
-import metagpt.tools.functions.libs.data_preprocess
diff --git a/metagpt/tools/functions/libs/ml_model.py b/metagpt/tools/functions/libs/ml_model.py
deleted file mode 100644
index b669de2c1..000000000
--- a/metagpt/tools/functions/libs/ml_model.py
+++ /dev/null
@@ -1,196 +0,0 @@
-from sklearn.model_selection import train_test_split
-from sklearn.preprocessing import LabelEncoder
-
-from sklearn.linear_model import LogisticRegression
-from sklearn.ensemble import RandomForestClassifier
-from sklearn.ensemble import GradientBoostingClassifier
-
-
-from sklearn.linear_model import LinearRegression
-from sklearn.ensemble import RandomForestRegressor
-from sklearn.ensemble import GradientBoostingRegressor
-
-from metagpt.tools.functions import registry
-from metagpt.tools.functions.schemas.ml_model import *
-
-
-#########
-## 分类 ##
-#########
-
-
-@registry.register("classification_model", LogisticRegressionClassification)
-def logistic_regression_classification(df, label, test_size=0.2, penalty='l2', dual=False):
-    nonnumeric_columns = [col for col in df if df[col].dtype == 'object']
-    for col in nonnumeric_columns:
-        df[col] = LabelEncoder().fit_transform(df[col])
-    df = df.fillna(0)
-
-    features = [col for col in df if col != label]
-    x, y = df[features], df[label]
-    tr_x, te_x, tr_y, te_y = train_test_split(x, y, test_size=test_size, random_state=1)
-
-    model = LogisticRegression(penalty=penalty, dual=dual)
-    model.fit(tr_x, tr_y, )
-    te_pred_prob = model.predict_proba(te_x)
-
-    res = {
-        'te_pred_prob': te_pred_prob
-    }
-    return res
-
-
-@registry.register("classification_model", RandomForestClassification)
-def random_forest_classification(df, label, test_size=0.2, n_estimators=100, criterion='gini'):
-    nonnumeric_columns = [col for col in df if df[col].dtype == 'object']
-    for col in nonnumeric_columns:
-        df[col] = LabelEncoder().fit_transform(df[col])
-    df = df.fillna(0)
-
-    features = [col for col in df if col != label]
-    x, y = df[features], df[label]
-    tr_x, te_x, tr_y, te_y = train_test_split(x, y, test_size=test_size, random_state=1)
-    model = RandomForestClassifier(n_estimators=n_estimators, criterion=criterion)
-    model.fit(tr_x, tr_y, )
-    te_pred_prob = model.predict_proba(te_x)
-
-    res = {
-        'te_pred_prob': te_pred_prob
-    }
-    return res
-
-
-@registry.register("classification_model", GradientBoostingClassification)
-def gradient_boosting_classification(df, label, test_size=0.2, n_estimators=100, learning_rate=0.1):
-    nonnumeric_columns = [col for col in df if df[col].dtype == 'object']
-    for col in nonnumeric_columns:
-        df[col] = LabelEncoder().fit_transform(df[col])
-    df = df.fillna(0)
-
-    features = [col for col in df if col != label]
-    x, y = df[features], df[label]
-    tr_x, te_x, tr_y, te_y = train_test_split(x, y, test_size=test_size, random_state=1)
-    model = GradientBoostingClassifier(n_estimators=n_estimators, learning_rate=learning_rate)
-    model.fit(tr_x, tr_y, )
-    te_pred_prob = model.predict_proba(te_x)
-
-    res = {
-        'te_pred_prob': te_pred_prob
-    }
-    return res
-
-
-
-#########
-## 回归 ##
-#########
-
-
-@registry.register("regression_model", LinearRegressionRegression)
-def linear_regression(df, label, test_size=0.2, ):
-    nonnumeric_columns = [col for col in df if df[col].dtype == 'object']
-    for col in nonnumeric_columns:
-        df[col] = LabelEncoder().fit_transform(df[col])
-    df = df.fillna(0)
-
-    features = [col for col in df if col != label]
-    x, y = df[features], df[label]
-    tr_x, te_x, tr_y, te_y = train_test_split(x, y, test_size=test_size, random_state=1)
-
-    model = LinearRegression()
-    model.fit(tr_x, tr_y, )
-    te_pred_prob = model.predict(te_x)
-
-    res = {
-        'te_pred_prob': te_pred_prob
-    }
-    return res
-
-
-@registry.register("regression_model", RandomForestRegression)
-def random_forest_regression(df, label, test_size=0.2, n_estimators=100, criterion='squared_error'):
-    nonnumeric_columns = [col for col in df if df[col].dtype == 'object']
-    for col in nonnumeric_columns:
-        df[col] = LabelEncoder().fit_transform(df[col])
-    df = df.fillna(0)
-
-    features = [col for col in df if col != label]
-    x, y = df[features], df[label]
-    tr_x, te_x, tr_y, te_y = train_test_split(x, y, test_size=test_size, random_state=1)
-    model = RandomForestRegressor(n_estimators=n_estimators, criterion=criterion)
-    model.fit(tr_x, tr_y, )
-    te_pred_prob = model.predict(te_x)
-
-    res = {
-        'te_pred_prob': te_pred_prob
-    }
-    return res
-
-
-@registry.register("regression_model", GradientBoostingRegression)
-def gradient_boosting_regression(df, label, test_size=0.2, n_estimators=100, learning_rate=0.1):
-    nonnumeric_columns = [col for col in df if df[col].dtype == 'object']
-    for col in nonnumeric_columns:
-        df[col] = LabelEncoder().fit_transform(df[col])
-    df = df.fillna(0)
-
-    features = [col for col in df if col != label]
-    x, y = df[features], df[label]
-    tr_x, te_x, tr_y, te_y = train_test_split(x, y, test_size=test_size, random_state=1)
-    model = GradientBoostingRegressor(n_estimators=n_estimators, learning_rate=learning_rate)
-    model.fit(tr_x, tr_y, )
-    te_pred_prob = model.predict(te_x)
-
-    res = {
-        'te_pred_prob': te_pred_prob
-    }
-    return res
-
-
-if __name__ == '__main__':
-    def run():
-        from sklearn.datasets import load_iris
-        loader = load_iris(as_frame=True)
-        df = loader['data']
-        df['target'] = loader['target']
-
-        df[df.columns[0]] = df[df.columns[0]].astype(str)
-        df[df.columns[1]] = df[df.columns[1]].astype(int)
-        df['target'] = df['target'].astype(str)
-
-        print(df)
-        print('####'*5)
-        res = logistic_regression_classification(df, 'target', test_size=0.25, penalty='l2', dual=False)
-        print(res['te_pred_prob'])
-
-        print('####'*5)
-        res = random_forest_classification(df, 'target', test_size=0.25, n_estimators=100, criterion='gini')
-        print(res['te_pred_prob'])
-
-        print('####'*5)
-        res = gradient_boosting_classification(df, 'target', test_size=0.25, n_estimators=100, learning_rate=0.1)
-        print(res['te_pred_prob'])
-
-        from sklearn.datasets import make_regression
-        import pandas as pd
-        loader = make_regression()
-        df = pd.DataFrame(loader[0])
-        df['target'] = loader[1]
-
-        df[df.columns[0]] = df[df.columns[0]].astype(str)
-        df[df.columns[1]] = df[df.columns[1]].astype(int)
-        # df['target'] = df['target'].astype(str)
-
-        print(df)
-        print('####' * 5)
-        res = linear_regression(df, 'target', test_size=0.25, )
-        print(res['te_pred_prob'])
-
-        print('####' * 5)
-        res = random_forest_regression(df, 'target', test_size=0.25, n_estimators=100, criterion='squared_error')
-        print(res['te_pred_prob'])
-
-        print('####' * 5)
-        res = gradient_boosting_regression(df, 'target', test_size=0.25, n_estimators=100, learning_rate=0.1)
-        print(res['te_pred_prob'])
-    run()
\ No newline at end of file
diff --git a/metagpt/tools/functions/register/__init__.py b/metagpt/tools/functions/register/__init__.py
deleted file mode 100644
index c80872750..000000000
--- a/metagpt/tools/functions/register/__init__.py
+++ /dev/null
@@ -1,6 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-# @Time    : 2023/11/16 16:37
-# @Author  : lidanyang
-# @File    : __init__.py
-# @Desc    :
diff --git a/metagpt/tools/functions/register/register.py b/metagpt/tools/functions/register/register.py
deleted file mode 100644
index 0731e31c0..000000000
--- a/metagpt/tools/functions/register/register.py
+++ /dev/null
@@ -1,78 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-# @Time    : 2023/11/16 16:38
-# @Author  : lidanyang
-# @File    : register.py
-# @Desc    :
-import inspect
-from typing import Type, Optional, Callable, Dict, Union, List
-
-from metagpt.tools.functions.schemas.base import ToolSchema
-
-
-class FunctionRegistry:
-    def __init__(self):
-        self.functions: Dict[str, Dict[str, Dict]] = {}
-
-    @staticmethod
-    def _check_param_consistency(func_params, schema):
-        param_names = set(func_params.keys())
-        schema_names = set(schema["parameters"]["properties"].keys())
-
-        if param_names != schema_names:
-            raise ValueError("Function parameters do not match schema properties")
-
-    def register(self, module: str, tool_schema: Type[ToolSchema]) -> Callable:
-        def wrapper(func: Callable) -> Callable:
-            module_registry = self.functions.setdefault(module, {})
-
-            if func.__name__ in module_registry:
-                raise ValueError(f"Function {func.__name__} is already registered in {module}")
-
-            func_params = inspect.signature(func).parameters
-
-            schema = tool_schema.schema()
-            schema["name"] = func.__name__
-
-            self._check_param_consistency(func_params, schema)
-
-            module_registry[func.__name__] = {
-                "func": func,
-                "schema": schema,
-            }
-            return func
-
-        return wrapper
-
-    def get(self, module: str, name: str) -> Optional[Union[Callable, Dict]]:
-        """Get function by module and name"""
-        module_registry = self.functions.get(module, {})
-        return module_registry.get(name)
-
-    def get_by_name(self, name: str) -> Optional[Dict]:
-        """Get function by name"""
-        for module_registry in self.functions.values():
-            if name in module_registry:
-                return module_registry.get(name, {})
-
-    def get_all_by_module(self, module: str) -> Optional[Dict]:
-        """Get all functions by module"""
-        return self.functions.get(module, {})
-
-    def get_schema(self, module: str, name: str) -> Optional[Dict]:
-        """Get schema by module and name"""
-        module_registry = self.functions.get(module, {})
-        return module_registry.get(name, {}).get("schema")
-
-    def get_schemas(self, module: str, names: List[str]) -> List[Dict]:
-        """Get schemas by module and names"""
-        module_registry = self.functions.get(module, {})
-        return [module_registry.get(name, {}).get("schema") for name in names]
-
-    def get_all_schema_by_module(self, module: str) -> List[Dict]:
-        """Get all schemas by module"""
-        module_registry = self.functions.get(module, {})
-        return [v.get("schema") for v in module_registry.values()]
-
-
-registry = FunctionRegistry()
diff --git a/metagpt/tools/functions/schemas/base.py b/metagpt/tools/functions/schemas/base.py
deleted file mode 100644
index aef604c8d..000000000
--- a/metagpt/tools/functions/schemas/base.py
+++ /dev/null
@@ -1,100 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-# @Time    : 2023/11/16 16:34
-# @Author  : lidanyang
-# @File    : base.py
-# @Desc    : Build base class to generate schema for tool
-from typing import Any, List, Optional, get_type_hints
-
-
-class NoDefault:
-    """
-    A class to represent a missing default value.
-
-    This is used to distinguish between a default value of None and a missing default value.
-    """
-    pass
-
-
-def tool_field(
-    description: str, default: Any = NoDefault(), enum: Optional[List[Any]] = None, **kwargs
-):
-    """
-    Create a field for a tool parameter.
-
-    Args:
-        description (str): A description of the field.
-        default (Any, optional): The default value for the field. Defaults to None.
-        enum (Optional[List[Any]], optional): A list of possible values for the field. Defaults to None.
-        **kwargs: Additional keyword arguments.
-
-    Returns:
-        dict: A dictionary representing the field with provided attributes.
-    """
-    field_info = {
-        "description": description,
-        "default": default,
-        "enum": enum,
-    }
-    field_info.update(kwargs)
-    return field_info
-
-
-class ToolSchema:
-    @staticmethod
-    def format_type(type_hint):
-        """
-        Format a type hint into a string representation.
-
-        Args:
-            type_hint (type): The type hint to format.
-
-        Returns:
-            str: A string representation of the type hint.
-        """
-        if isinstance(type_hint, type):
-            # Handle built-in types separately
-            if type_hint.__module__ == "builtins":
-                return type_hint.__name__
-            else:
-                return f"{type_hint.__module__}.{type_hint.__name__}"
-        elif hasattr(type_hint, "__origin__") and hasattr(type_hint, "__args__"):
-            # Handle generic types (like List[int])
-            origin_type = ToolSchema.format_type(type_hint.__origin__)
-            args_type = ", ".join(
-                [ToolSchema.format_type(t) for t in type_hint.__args__]
-            )
-            return f"{origin_type}[{args_type}]"
-        else:
-            return str(type_hint)
-
-    @classmethod
-    def schema(cls):
-        """
-        Generate a schema dictionary for the class.
-
-        The schema includes the class name, description, and information about
-        each class parameter based on type hints and field definitions.
-
-        Returns:
-            dict: A dictionary representing the schema of the class.
-        """
-        schema = {
-            "name": cls.__name__,
-            "description": cls.__doc__,
-            "parameters": {"type": "object", "properties": {}, "required": []},
-        }
-        type_hints = get_type_hints(cls)
-        for attr, type_hint in type_hints.items():
-            value = getattr(cls, attr, None)
-            if isinstance(value, dict):
-                # Process each attribute that is defined using the field function
-                prop_info = {k: v for k, v in value.items() if v is not None or k == "default"}
-                if isinstance(prop_info["default"], NoDefault):
-                    del prop_info["default"]
-                prop_info["type"] = ToolSchema.format_type(type_hint)
-                schema["parameters"]["properties"][attr] = prop_info
-                # Check for required fields
-                if "default" not in prop_info:
-                    schema["parameters"]["required"].append(attr)
-        return schema
diff --git a/metagpt/tools/functions/schemas/data_preprocess.py b/metagpt/tools/functions/schemas/data_preprocess.py
deleted file mode 100644
index 16b97aeac..000000000
--- a/metagpt/tools/functions/schemas/data_preprocess.py
+++ /dev/null
@@ -1,67 +0,0 @@
-
-import pandas as pd
-
-from metagpt.tools.functions.schemas.base import tool_field, ToolSchema
-
-
-class FillMissingValue(ToolSchema):
-    """Completing missing values with simple strategies"""
-    df: pd.DataFrame = tool_field(description="input dataframe")
-    features: list = tool_field(description="columns to be processed")
-    strategy: str = tool_field(
-        description="the imputation strategy",
-        default='mean',
-        enum=['mean', 'median', 'most_frequent', 'constant']
-    )
-    fill_value: int = tool_field(
-        description="fill_value is used to replace all occurrences of missing_values", default=None)
-
-
-class SplitBins(ToolSchema):
-    """Bin continuous data into intervals and return the bin identifier encoded as an integer value"""
-    df: pd.DataFrame = tool_field(description="input dataframe")
-    features: list = tool_field(description="columns to be processed")
-    strategy: str = tool_field(description="Strategy used to define the widths of the bins", default='quantile')
-
-
-class MinMaxScale(ToolSchema):
-    """Transform features by scaling each feature to a range, witch is (0, 1)"""
-    df: pd.DataFrame = tool_field(description="input dataframe")
-    features: list = tool_field(description="columns to be processed")
-
-
-class StandardScale(ToolSchema):
-    """Standardize features by removing the mean and scaling to unit variance"""
-    df: pd.DataFrame = tool_field(description="input dataframe")
-    features: list = tool_field(description="columns to be processed")
-
-
-class LogTransform(ToolSchema):
-    """Performs a logarithmic transformation on the specified columns"""
-    df: pd.DataFrame = tool_field(description="input dataframe")
-    features: list = tool_field(description="columns to be processed")
-
-
-class MaxAbsScale(ToolSchema):
-    """Scale each feature by its maximum absolute value"""
-    df: pd.DataFrame = tool_field(description="input dataframe")
-    features: list = tool_field(description="columns to be processed")
-
-
-class RobustScale(ToolSchema):
-    """Scale features using statistics that are robust to outliers, the quantile_range is (25.0, 75.0)"""
-    df: pd.DataFrame = tool_field(description="input dataframe")
-    features: list = tool_field(description="columns to be processed")
-
-
-class OrdinalEncode(ToolSchema):
-    """Encode categorical features as an integer array"""
-    df: pd.DataFrame = tool_field(description="input dataframe")
-    features: list = tool_field(description="columns to be processed")
-
-
-class OneHotEncoding(ToolSchema):
-    """Apply one-hot encoding to specified categorical columns, the original columns will be dropped."""
-
-    df: pd.DataFrame = tool_field(description="DataFrame to process.")
-    cols: list = tool_field(description="Categorical columns to be one-hot encoded and dropped.")
diff --git a/metagpt/tools/functions/schemas/feature_engineering.py b/metagpt/tools/functions/schemas/feature_engineering.py
deleted file mode 100644
index 5c89d9b16..000000000
--- a/metagpt/tools/functions/schemas/feature_engineering.py
+++ /dev/null
@@ -1,110 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-# @Time    : 2023/11/17 10:34
-# @Author  : lidanyang
-# @File    : feature_engineering.py
-# @Desc    : Schema for feature engineering functions
-from typing import List
-
-import pandas as pd
-
-from metagpt.tools.functions.schemas.base import ToolSchema, tool_field
-
-
-class PolynomialExpansion(ToolSchema):
-    """Add polynomial and interaction features from selected numeric columns, excluding the bias column."""
-
-    df: pd.DataFrame = tool_field(description="DataFrame to process.")
-    cols: list = tool_field(description="Columns for polynomial expansion.")
-    degree: int = tool_field(description="Degree of polynomial features.", default=2)
-
-
-class FrequencyEncoding(ToolSchema):
-    """Add value counts of categorical columns as new features."""
-
-    df: pd.DataFrame = tool_field(description="DataFrame to process.")
-    cols: list = tool_field(description="Categorical columns to be frequency encoded.")
-
-
-class TargetMeanEncoder(ToolSchema):
-    """Encodes a categorical column by the mean of the label column, and adds the result as a new feature."""
-
-    df: pd.DataFrame = tool_field(description="DataFrame to process.")
-    col: str = tool_field(description="Column to be mean encoded.")
-    label: str = tool_field(description="Predicted label column.")
-
-
-class KFoldTargetMeanEncoder(ToolSchema):
-    """Adds a new feature to the DataFrame by k-fold mean encoding of a categorical column using the label column."""
-    df: pd.DataFrame = tool_field(description="DataFrame to process.")
-    col: str = tool_field(description="Column to be k-fold mean encoded.")
-    label: str = tool_field(description="Predicted label column.")
-    n_splits: int = tool_field(description="Number of splits for K-fold.", default=5)
-    random_state: int = tool_field(description="Random seed.", default=2021)
-
-
-class CatCross(ToolSchema):
-    """Add pairwise crossed features and convert them to numerical features."""
-
-    df: pd.DataFrame = tool_field(description="DataFrame to process.")
-    cols: list = tool_field(description="Columns to be pairwise crossed.")
-    max_cat_num: int = tool_field(
-        description="Maximum unique categories per crossed feature.", default=100
-    )
-
-
-class GroupStat(ToolSchema):
-    """Aggregate specified column in a DataFrame grouped by another column, adding new features named '<agg_col>_<agg_func>_by_<group_col>'."""
-
-    df: pd.DataFrame = tool_field(description="DataFrame to process.")
-    group_col: str = tool_field(description="Column used for grouping.")
-    agg_col: str = tool_field(description="Column on which aggregation is performed.")
-    agg_funcs: list = tool_field(
-        description="""List of aggregation functions to apply, such as ['mean', 'std'].
-                    Each function must be supported by pandas."""
-    )
-
-
-class ExtractTimeComps(ToolSchema):
-    """Extract and add specific time components as new features from a designated time column."""
-
-    df: pd.DataFrame = tool_field(description="DataFrame to process.")
-    time_col: str = tool_field(
-        description="The name of the column containing time data."
-    )
-    time_comps: List[str] = tool_field(
-        description="""List of time components to extract.
-        Each component must be in ['year', 'month', 'day', 'hour', 'dayofweek', 'is_weekend']."""
-    )
-
-
-class FeShiftByTime(ToolSchema):
-    """Shift column values based on specified time intervals and add the resulting new features to the DataFrame. New features are named in the format of '<group_col>_<shift_col>_lag_<period>_<freq>'."""
-
-    df: pd.DataFrame = tool_field(description="DataFrame to process.")
-    time_col: str = tool_field(description="Column for time-based shifting.")
-    group_col: str = tool_field(description="Column for grouping before shifting.")
-    shift_col: str = tool_field(description="Column to shift.")
-    periods: list = tool_field(description="Time intervals for shifting.")
-    freq: str = tool_field(
-        description="Frequency unit for time intervals (e.g., 'D', 'M').",
-        enum=["D", "M", "Y", "W", "H"],
-    )
-
-
-class FeRollingByTime(ToolSchema):
-    """Calculate rolling statistics for a DataFrame column over time intervals."""
-
-    df: pd.DataFrame = tool_field(description="DataFrame to process.")
-    time_col: str = tool_field(description="Column for time-based rolling.")
-    group_col: str = tool_field(description="Column for grouping before rolling.")
-    rolling_col: str = tool_field(description="Column for rolling calculations.")
-    periods: list = tool_field(description="Window sizes for rolling.")
-    freq: str = tool_field(
-        description="Frequency unit for time windows (e.g., 'D', 'M').",
-        enum=["D", "M", "Y", "W", "H"],
-    )
-    agg_funcs: list = tool_field(
-        description="""List of aggregation functions for rolling, like ['mean', 'std'].
-        Each function must be in ['mean', 'std', 'min', 'max', 'median', 'sum', 'count']."""
-    )
diff --git a/metagpt/tools/functions/schemas/ml_model.py b/metagpt/tools/functions/schemas/ml_model.py
deleted file mode 100644
index 9268156af..000000000
--- a/metagpt/tools/functions/schemas/ml_model.py
+++ /dev/null
@@ -1,55 +0,0 @@
-import pandas as pd
-
-from metagpt.tools.functions.schemas.base import tool_field, ToolSchema
-
-
-class LogisticRegressionClassification(ToolSchema):
-    """Logistic Regression (aka logit, MaxEnt) classifier"""
-    df: pd.DataFrame = tool_field(description="input dataframe")
-    label: str = tool_field(description="target name")
-    test_size: float = tool_field(description="The proportion of the test set to all the data", default=0.2)
-    penalty: str = tool_field(description="Specify the norm of the penalty", default="l2")
-    dual: bool = tool_field(description="Dual (constrained) or primal (regularized) formulation", default="l2")
-
-
-class RandomForestClassification(ToolSchema):
-    """random forest is a meta estimator that fits a number of decision tree classifiers on various sub-samples of the dataset and uses averaging to improve the predictive accuracy and control over-fitting"""
-    df: pd.DataFrame = tool_field(description="input dataframe")
-    label: str = tool_field(description="target name")
-    test_size: float = tool_field(description="The proportion of the test set to all the data", default=0.2)
-    n_estimators: int = tool_field(description="The number of trees in the forest", default=100)
-    criterion: str = tool_field(description="The function to measure the quality of a split", default="gini")
-
-
-class GradientBoostingClassification(ToolSchema):
-    """Gradient Boosting for classification.This algorithm builds an additive model in a forward stage-wise fashion"""
-    df: pd.DataFrame = tool_field(description="input dataframe")
-    label: str = tool_field(description="target name")
-    test_size: float = tool_field(description="The proportion of the test set to all the data", default=0.2)
-    n_estimators: int = tool_field(description="The number of boosting stages to perform", default=100)
-    learning_rate: float = tool_field(description="Learning rate shrinks the contribution of each tree by learning_rate", default=0.1)
-
-
-class LinearRegressionRegression(ToolSchema):
-    """Ordinary least squares Linear Regression."""
-    df: pd.DataFrame = tool_field(description="input dataframe")
-    label: str = tool_field(description="target name")
-    test_size: float = tool_field(description="The proportion of the test set to all the data", default=0.2)
-
-
-class RandomForestRegression(ToolSchema):
-    """random forest is a meta estimator that fits a number of decision tree on various sub-samples of the dataset and uses averaging to improve the predictive accuracy and control over-fitting"""
-    df: pd.DataFrame = tool_field(description="input dataframe")
-    label: str = tool_field(description="target name")
-    test_size: float = tool_field(description="The proportion of the test set to all the data", default=0.2)
-    n_estimators: int = tool_field(description="The number of trees in the forest", default=100)
-    criterion: str = tool_field(description="The function to measure the quality of a split", default="squared_error")
-
-
-class GradientBoostingRegression(ToolSchema):
-    """Gradient Boosting for regression.This estimator builds an additive model in a forward stage-wise fashion"""
-    df: pd.DataFrame = tool_field(description="input dataframe")
-    label: str = tool_field(description="target name")
-    test_size: float = tool_field(description="The proportion of the test set to all the data", default=0.2)
-    n_estimators: int = tool_field(description="The number of boosting stages to perform", default=100)
-    learning_rate: float = tool_field(description="Learning rate shrinks the contribution of each tree by learning_rate", default=0.1)

From 2a3f23ec62ebca8329c2748179d731025a685d0a Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Thu, 14 Dec 2023 10:32:58 +0800
Subject: [PATCH 38/49] fix unittest

---
 .../actions/test_write_analysis_code.py       | 33 ++++++++-----------
 1 file changed, 13 insertions(+), 20 deletions(-)

diff --git a/tests/metagpt/actions/test_write_analysis_code.py b/tests/metagpt/actions/test_write_analysis_code.py
index 661202115..1a568cdcd 100644
--- a/tests/metagpt/actions/test_write_analysis_code.py
+++ b/tests/metagpt/actions/test_write_analysis_code.py
@@ -31,22 +31,15 @@ async def test_tool_recommendation():
     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",
-        },
-    ]
+    available_tools = {
+        "fill_missing_value": "Completing missing values with simple strategies",
+        "split_bins": "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"]
+    assert len(tools) == 1
+    assert tools[0] == ["fill_missing_value"]
 
 
 @pytest.mark.asyncio
@@ -57,7 +50,7 @@ async def test_write_code_with_tools():
         "1": Task(
                 task_id="1",
                 instruction="随机生成一个pandas DataFrame数据集",
-                task_type="unknown",
+                task_type="other",
                 dependent_task_ids=[],
                 code="""
                 import pandas as pd
@@ -75,6 +68,10 @@ async def test_write_code_with_tools():
                 instruction="对数据集进行数据清洗",
                 task_type="data_preprocess",
                 dependent_task_ids=["1"],
+                code_steps="""
+                {"Step 1": "对数据集进行去重",
+                "Step 2": "对数据集进行缺失值处理"}
+                """
             ),
     }
     plan = Plan(
@@ -83,13 +80,9 @@ async def test_write_code_with_tools():
         task_map=task_map,
         current_task_id="2",
     )
-    task_guide = """
-    step 1: 对数据集进行去重
-    step 2: 对数据集进行缺失值处理
-    """
-    data_desc = "None"
+    column_info = ""
 
-    code = await write_code.run(messages, plan, task_guide, data_desc)
+    code = await write_code.run(messages, plan, column_info)
     assert len(code) > 0
     print(code)
 

From d84e9cae2c8dfc5345edb253f59ca1f0901cacab Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Thu, 14 Dec 2023 10:34:15 +0800
Subject: [PATCH 39/49] fix conflict

---
 metagpt/actions/write_analysis_code.py |  6 ++----
 metagpt/roles/ml_engineer.py           | 12 ++++++------
 2 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/metagpt/actions/write_analysis_code.py b/metagpt/actions/write_analysis_code.py
index 2c45281f9..6970fb4f0 100644
--- a/metagpt/actions/write_analysis_code.py
+++ b/metagpt/actions/write_analysis_code.py
@@ -24,8 +24,8 @@ from metagpt.utils.common import create_func_config, remove_comments
 
 
 class BaseWriteAnalysisCode(Action):
-    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."""  # 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. 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 process_msg(self, prompt: Union[str, List[Dict], Message, List[Message]], system_msg: str = None):
         default_system_msg = system_msg or self.DEFAULT_SYSTEM_MSG
@@ -201,8 +201,6 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
                 module_name=module_name,
                 tool_catalog=tool_catalog,
             )
-
-
         else:
             prompt = GENERATE_CODE_PROMPT.format(
                 user_requirement=plan.goal,
diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index 8f06a541c..0b76711f4 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -159,12 +159,12 @@ class MLEngineer(Role):
             # print("*" * 10)
             # breakpoint()
             if counter > 0:
-                improve_code = await DebugCode().run(plan=self.plan.current_task.instruction,
-                                                     # finished_code=code_context,
-                                                     # finished_code_result=code_result,
-                                                     code=code,
-                                                     runtime_result=self.working_memory.get(),
-                                                     context=debug_context)
+                improve_code = await DebugCode().run(
+                    plan=self.plan.current_task.instruction,
+                    code=code,
+                    runtime_result=self.working_memory.get(),
+                    context=debug_context
+                )
 
             if improve_code != "":
                 code = improve_code

From 44334c0c9aa6b8a0d6314d3b24623d9633ce7c2d Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Thu, 14 Dec 2023 10:59:42 +0800
Subject: [PATCH 40/49] drop old schema import

---
 metagpt/tools/functions/libs/data_preprocess.py     | 7 +++++--
 metagpt/tools/functions/libs/feature_engineering.py | 2 +-
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/metagpt/tools/functions/libs/data_preprocess.py b/metagpt/tools/functions/libs/data_preprocess.py
index fa70bf8fc..ec3580889 100644
--- a/metagpt/tools/functions/libs/data_preprocess.py
+++ b/metagpt/tools/functions/libs/data_preprocess.py
@@ -1,4 +1,5 @@
 import numpy as np
+import pandas as pd
 from sklearn.impute import SimpleImputer
 from sklearn.preprocessing import LabelEncoder
 from sklearn.preprocessing import MaxAbsScaler
@@ -9,7 +10,6 @@ from sklearn.preprocessing import RobustScaler
 from sklearn.preprocessing import StandardScaler
 
 from metagpt.tools.functions.libs.base import MLProcess
-from metagpt.tools.functions.schemas.data_preprocess import *
 
 
 class FillMissingValue(MLProcess):
@@ -141,7 +141,10 @@ def get_column_info(df: pd.DataFrame) -> dict:
     for i in df.columns:
         nan_freq = float("%.2g" % (df[i].isna().mean() * 100))
         n_unique = df[i].nunique()
-        data.append([i, df[i].dtype, nan_freq, n_unique])
+        data_type = str(df[i].dtype).replace("dtype('", "").replace("')", "")
+        if data_type == "O":
+            data_type = "object"
+        data.append([i, data_type, nan_freq, n_unique])
 
     samples = pd.DataFrame(
         data,
diff --git a/metagpt/tools/functions/libs/feature_engineering.py b/metagpt/tools/functions/libs/feature_engineering.py
index de54e4db0..1ec2b9675 100644
--- a/metagpt/tools/functions/libs/feature_engineering.py
+++ b/metagpt/tools/functions/libs/feature_engineering.py
@@ -7,6 +7,7 @@
 import itertools
 
 import numpy as np
+import pandas as pd
 from dateutil.relativedelta import relativedelta
 from joblib import Parallel, delayed
 from pandas.api.types import is_numeric_dtype
@@ -15,7 +16,6 @@ from sklearn.model_selection import KFold
 from sklearn.preprocessing import PolynomialFeatures, KBinsDiscretizer
 
 from metagpt.tools.functions.libs.base import MLProcess
-from metagpt.tools.functions.schemas.feature_engineering import *
 
 
 class PolynomialExpansion(MLProcess):

From 5940c8d908b12d8c99cc03305dec4fcf8bcc3dd8 Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Thu, 14 Dec 2023 12:56:01 +0800
Subject: [PATCH 41/49] remove old comments

---
 metagpt/roles/ml_engineer.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index 0b76711f4..51faf1e0d 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -261,14 +261,12 @@ class MLEngineer(Role):
         """find useful memories only to reduce context length and improve performance"""
         # TODO dataset description , code steps
         if task_exclude_field is None:
+            # Shorten the context as we don't need code steps after we get the codes.
+            # This doesn't affect current_task below, which should hold the code steps
             task_exclude_field = {'code_steps'}
         user_requirement = self.plan.goal
         data_desc = self.plan.context
         tasks = [task.dict(exclude=task_exclude_field) for task in self.plan.tasks]
-        # for task in tasks:
-        #     # Shorten the context as we don't need code steps after we get the codes.
-        #     # This doesn't affect current_task below, which should hold the code steps
-        #     task.pop("code_steps")
         tasks = json.dumps(tasks, indent=4, ensure_ascii=False)
         current_task = self.plan.current_task.json() if self.plan.current_task else {}
         context = STRUCTURAL_CONTEXT.format(

From ef6e4a1b77a21cefeb165301dd1d47b5c273fdbb Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Thu, 14 Dec 2023 13:46:27 +0800
Subject: [PATCH 42/49] debug only when use_tools

---
 metagpt/roles/ml_engineer.py | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index 51faf1e0d..3755e7bac 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -147,7 +147,6 @@ class MLEngineer(Role):
         )
 
         counter = 0
-        improve_code = ""
         success = False
         debug_context = []
 
@@ -158,17 +157,14 @@ class MLEngineer(Role):
             # print(context)
             # print("*" * 10)
             # breakpoint()
-            if counter > 0:
-                improve_code = await DebugCode().run(
+            if counter > 0 and self.use_tools:
+                code = await DebugCode().run(
                     plan=self.plan.current_task.instruction,
                     code=code,
                     runtime_result=self.working_memory.get(),
                     context=debug_context
                 )
-
-            if improve_code != "":
-                code = improve_code
-                logger.info(f"new code \n{improve_code}")
+                logger.info(f"new code \n{code}")
                 cause_by = DebugCode
             elif not self.use_tools or self.plan.current_task.task_type == "other":
                 logger.info("Write code with pure generation")

From 97f707784bd8558b3bbd138d9380af55bb85f9a4 Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Thu, 14 Dec 2023 13:56:23 +0800
Subject: [PATCH 43/49] reformat

---
 metagpt/actions/debug_code.py | 124 ++++++++++++++++++----------------
 1 file changed, 64 insertions(+), 60 deletions(-)

diff --git a/metagpt/actions/debug_code.py b/metagpt/actions/debug_code.py
index 58d006a08..3e1705d8e 100644
--- a/metagpt/actions/debug_code.py
+++ b/metagpt/actions/debug_code.py
@@ -1,57 +1,56 @@
 from typing import Dict, List, Union, Tuple, Optional, Any
 
-from metagpt.actions import Action
 from metagpt.logs import logger
 from metagpt.schema import Message, Plan
 from metagpt.utils.common import CodeParser, create_func_config
 from metagpt.actions.write_analysis_code import BaseWriteAnalysisCode
 
-DEBUG_REFLECTION_EXAMPLE = '''Example 1:
-                           [previous impl]:
-                           ```python
-                           def add(a: int, b: int) -> int:
-                               """
-                               Given integers a and b, return the total value of a and b.
-                               """
-                               return a - b
-                           ```
+DEBUG_REFLECTION_EXAMPLE = '''
+Example 1:
+[previous impl]:
+```python
+def add(a: int, b: int) -> int:
+   """
+   Given integers a and b, return the total value of a and b.
+   """
+   return a - b
+```
 
-                           [runtime Error]:
-                           Tested passed:
+[runtime Error]:
+Tested passed:
 
-                           Tests failed:
-                           assert add(1, 2) == 3 # output: -1
-                           assert add(1, 2) == 4 # output: -1
+Tests failed:
+assert add(1, 2) == 3 # output: -1
+assert add(1, 2) == 4 # output: -1
 
-                           [reflection on previous impl]:
-                           The implementation failed the test cases where the input integers are 1 and 2. The issue arises because the code does not add the two integers together, but instead subtracts the second integer from the first. To fix this issue, we should change the operator from `-` to `+` in the return statement. This will ensure that the function returns the correct output for the given input.
+[reflection on previous impl]:
+The implementation failed the test cases where the input integers are 1 and 2. The issue arises because the code does not add the two integers together, but instead subtracts the second integer from the first. To fix this issue, we should change the operator from `-` to `+` in the return statement. This will ensure that the function returns the correct output for the given input.
 
-                           [improved impl]:
-                           ```python
-                           def add(a: int, b: int) -> int:
-                               """
-                               Given integers a and b, return the total value of a and b.
-                               """
-                               return a + b
-                           ```
-                           '''
+[improved impl]:
+```python
+def add(a: int, b: int) -> int:
+   """
+   Given integers a and b, return the total value of a and b.
+   """
+   return a + b
+```
+'''
 
 REFLECTION_PROMPT = """
-                       Here is an example for you.
-                       {debug_example}
-                       [context]
-                       {context}
-                       
-                       [previous impl]
-                       {code}
-                       [runtime Error]
-                       {runtime_result}
+Here is an example for you.
+{debug_example}
+[context]
+{context}
 
-                        Analysis the error step by step, provide me improve method and code. Remember to follow [context] rerquirement. Don't forget write code for steps behind the error step.
-                        [reflection on previous impl]:
-                        xxx
+[previous impl]
+{code}
+[runtime Error]
+{runtime_result}
 
-                       """
+Analysis the error step by step, provide me improve method and code. Remember to follow [context] rerquirement. Don't forget write code for steps behind the error step.
+[reflection on previous impl]:
+xxx
+"""
 
 CODE_REFLECTION = {
     "name": "execute_reflection_code",
@@ -85,10 +84,10 @@ class DebugCode(BaseWriteAnalysisCode):
     name: str = "debugcode"
     context: Optional[str] = None
     llm: None
-    
+
     def __init__(self, **kwargs: Any):
         super().__init__(**kwargs)
-    
+
     async def run_reflection(
         self,
         # goal,
@@ -100,23 +99,26 @@ class DebugCode(BaseWriteAnalysisCode):
     ) -> dict:
         info = []
         # finished_code_and_result = finished_code + "\n [finished results]\n\n" + finished_code_result
-        reflection_prompt = REFLECTION_PROMPT.format(debug_example=DEBUG_REFLECTION_EXAMPLE,
-                                                     context=context,
-                                                     # goal=goal,
-                                                     # finished_code=finished_code_and_result,
-                                                     code=code,
-                                                     runtime_result=runtime_result
-                                                     )
+        reflection_prompt = REFLECTION_PROMPT.format(
+            debug_example=DEBUG_REFLECTION_EXAMPLE,
+            context=context,
+            # goal=goal,
+            # finished_code=finished_code_and_result,
+            code=code,
+            runtime_result=runtime_result,
+        )
         system_prompt = "You are an AI Python assistant. You will be given your previous implementation code of a task, runtime error results, and a hint to change the implementation appropriately. Write your full implementation "
         info.append(Message(role="system", content=system_prompt))
         info.append(Message(role="user", content=reflection_prompt))
-        
+
         # msg = messages_to_str(info)
         # resp = await self.llm.aask(msg=msg)
-        resp = await self.llm.aask_code(messages=info, **create_func_config(CODE_REFLECTION))
+        resp = await self.llm.aask_code(
+            messages=info, **create_func_config(CODE_REFLECTION)
+        )
         logger.info(f"reflection is {resp}")
         return resp
-    
+
     # async def rewrite_code(self, reflection: str = "", context: List[Message] = None) -> str:
     #     """
     #     根据reflection重写代码
@@ -131,14 +133,16 @@ class DebugCode(BaseWriteAnalysisCode):
     #     resp = await self.llm.aask(msg=msg)
     #     improv_code = CodeParser.parse_code(block=None, text=resp)
     #     return improv_code
-    
-    async def run(self,
-                  context: List[Message] = None,
-                  plan: str = "",
-                  # finished_code: str = "",
-                  # finished_code_result: str = "",
-                  code: str = "",
-                  runtime_result: str = "") -> str:
+
+    async def run(
+        self,
+        context: List[Message] = None,
+        plan: str = "",
+        # finished_code: str = "",
+        # finished_code_result: str = "",
+        code: str = "",
+        runtime_result: str = "",
+    ) -> str:
         """
         根据当前运行代码和报错信息进行reflection和纠错
         """
@@ -152,5 +156,5 @@ class DebugCode(BaseWriteAnalysisCode):
         )
         # 根据reflection结果重写代码
         # improv_code = await self.rewrite_code(reflection, context=context)
-        improv_code = reflection['improved_impl']
+        improv_code = reflection["improved_impl"]
         return improv_code

From 2da141abbe43fa2c046a8f4bbdb0edc9325b03d3 Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Thu, 14 Dec 2023 13:57:39 +0800
Subject: [PATCH 44/49] recover code

---
 metagpt/tools/web_browser_engine.py | 2 +-
 metagpt/utils/__init__.py           | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/metagpt/tools/web_browser_engine.py b/metagpt/tools/web_browser_engine.py
index 7228ae9cf..453d87f31 100644
--- a/metagpt/tools/web_browser_engine.py
+++ b/metagpt/tools/web_browser_engine.py
@@ -7,7 +7,7 @@ from typing import Any, Callable, Coroutine, Literal, overload
 
 from metagpt.config import CONFIG
 from metagpt.tools import WebBrowserEngineType
-# from metagpt.utils.parse_html import WebPage
+from metagpt.utils.parse_html import WebPage
 
 
 class WebBrowserEngine:
diff --git a/metagpt/utils/__init__.py b/metagpt/utils/__init__.py
index 86cac50db..f13175cf8 100644
--- a/metagpt/utils/__init__.py
+++ b/metagpt/utils/__init__.py
@@ -6,7 +6,7 @@
 @File    : __init__.py
 """
 
-# from metagpt.utils.read_document import read_docx
+from metagpt.utils.read_document import read_docx
 from metagpt.utils.singleton import Singleton
 from metagpt.utils.token_counter import (
     TOKEN_COSTS,
@@ -16,7 +16,7 @@ from metagpt.utils.token_counter import (
 
 
 __all__ = [
-    # "read_docx",
+    "read_docx",
     "Singleton",
     "TOKEN_COSTS",
     "count_message_tokens",

From 234ffdab355f729ddfc385aa20bb6676e314174d Mon Sep 17 00:00:00 2001
From: Zhou <1359698378@qq.com>
Date: Thu, 14 Dec 2023 14:37:27 +0800
Subject: [PATCH 45/49] remove old typing-extensions version

---
 requirements.txt | 1 -
 1 file changed, 1 deletion(-)

diff --git a/requirements.txt b/requirements.txt
index 1d1bc95a1..2328de2a1 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -35,7 +35,6 @@ tqdm==4.64.0
 # webdriver_manager<3.9
 anthropic==0.3.6
 typing-inspect==0.8.0
-typing_extensions==4.5.0
 libcst==1.0.1
 qdrant-client==1.4.0
 pytest-mock==3.11.1

From 48d542d383bdb4bd80da68c546d3a553d8c543ed Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Thu, 14 Dec 2023 15:54:02 +0800
Subject: [PATCH 46/49] recover code

---
 metagpt/actions/execute_code.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/metagpt/actions/execute_code.py b/metagpt/actions/execute_code.py
index c5ed8964e..36e01ed0e 100644
--- a/metagpt/actions/execute_code.py
+++ b/metagpt/actions/execute_code.py
@@ -157,11 +157,6 @@ class ExecutePyCode(ExecuteCode, Action):
 
         return code, language
 
-    def save_notebook(self, path: str):
-        path = Path(path)
-        path.parent.mkdir(parents=True, exist_ok=True)
-        nbformat.write(self.nb, path)
-
     async def run(self, code: Union[str, Dict, Message], language: str = "python") -> Tuple[str, bool]:
         code, language = self._process_code(code, language)
 

From 2fe9f2b9cfed79677c11b16e34c3944d09b68df2 Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Fri, 15 Dec 2023 10:06:46 +0800
Subject: [PATCH 47/49] remove old comments

---
 .../tools/functions/libs/data_preprocess.py   | 50 -------------------
 1 file changed, 50 deletions(-)

diff --git a/metagpt/tools/functions/libs/data_preprocess.py b/metagpt/tools/functions/libs/data_preprocess.py
index ec3580889..8c70462ee 100644
--- a/metagpt/tools/functions/libs/data_preprocess.py
+++ b/metagpt/tools/functions/libs/data_preprocess.py
@@ -151,53 +151,3 @@ def get_column_info(df: pd.DataFrame) -> dict:
         columns=["Column_name", "Data_type", "NaN_Frequency(%)", "N_unique"],
     )
     return samples.to_dict(orient='list')
-#
-#
-# if __name__ == '__main__':
-#     def run():
-#         V = {
-#             'a': [-1, 2, 3, 6, 5, 4],
-#             'b': [1.1, 2.2, 3.3, 6.6, 5.5, 4.4],
-#             'c': ['aa', 'bb', 'cc', 'dd', 'ee', 'ff'],
-#             'd': [1, None, 3, None, 5, 4],
-#             'e': [1.1, np.NAN, 3.3, None, 5.5, 4.4],
-#             'f': ['aa', np.NAN, 'cc', None, '', 'ff'],
-#
-#         }
-#
-#         df = pd.DataFrame(V)
-#         print(df.dtypes)
-#
-#         numeric_features = ['a', 'b', 'd', 'e']
-#         numeric_features_wo_miss = ['a', 'b', ]
-#         categorial_features = ['c', 'f']
-#
-#         df_ = fill_missing_value(df.copy(), numeric_features)
-#         print(df_)
-#         df_ = fill_missing_value(df.copy(), categorial_features, strategy='constant', fill_value='hehe')
-#         print(df_)
-#
-#         df_ = fill_missing_value(df.copy(), numeric_features, strategy='constant', fill_value=999)
-#         print(df_)
-#
-#         # df_ = label_encode(df.copy(), numeric_features + categorial_features, )
-#         # print(df_)
-#
-#         df_ = split_bins(df.copy(), numeric_features_wo_miss, strategy='quantile')
-#         print(df_)
-#
-#         df_ = min_max_scale(df.copy(), numeric_features, )
-#         print(df_)
-#
-#         df_ = standard_scale(df.copy(), numeric_features, )
-#         print(df_)
-#
-#         df_ = log_transform(df.copy(), numeric_features, )
-#         print(df_)
-#
-#         df_ = max_abs_scale(df.copy(), numeric_features, )
-#         print(df_)
-#
-#         df_ = robust_scale(df.copy(), numeric_features, )
-#         print(df_)
-#     run()
\ No newline at end of file

From 27b59a67daa91318f48615dea0e8bef722592d1e Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Mon, 18 Dec 2023 10:33:17 +0800
Subject: [PATCH 48/49] recover code

---
 .gitignore | 2 --
 1 file changed, 2 deletions(-)

diff --git a/.gitignore b/.gitignore
index f2ccde1d1..2f1250b93 100644
--- a/.gitignore
+++ b/.gitignore
@@ -148,8 +148,6 @@ allure-results
 .DS_Store
 .vscode
 
-# Config
-config/config.yaml
 
 log.txt
 docs/scripts/set_env.sh

From d6566019b0c71e376b6aa27b85a9e54ee96e88ab Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Mon, 18 Dec 2023 10:34:38 +0800
Subject: [PATCH 49/49] recover code

---
 .gitignore | 1 -
 1 file changed, 1 deletion(-)

diff --git a/.gitignore b/.gitignore
index 2f1250b93..9b679d48a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -148,7 +148,6 @@ allure-results
 .DS_Store
 .vscode
 
-
 log.txt
 docs/scripts/set_env.sh
 key.yaml