From 89f1dce936f238d1931341c1ff79abce1555fc57 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Tue, 21 Nov 2023 18:38:41 +0800
Subject: [PATCH 01/11] feat: add actions: plan, write_code_function,
 pycode_executor.

---
 metagpt/actions/__init__.py                 |   6 +
 metagpt/actions/code_executor.py            | 173 ++++++++++++++++++++
 metagpt/actions/plan.py                     |  20 +++
 metagpt/actions/write_code_v2.py            |  36 ++++
 metagpt/prompts/plan.py                     |   7 +
 metagpt/schema.py                           |   1 +
 tests/metagpt/actions/test_code_executor.py |  58 +++++++
 tests/metagpt/actions/test_plan.py          |  12 ++
 tests/metagpt/actions/test_write_code_v2.py |  22 +++
 9 files changed, 335 insertions(+)
 create mode 100644 metagpt/actions/code_executor.py
 create mode 100644 metagpt/actions/plan.py
 create mode 100644 metagpt/actions/write_code_v2.py
 create mode 100644 metagpt/prompts/plan.py
 create mode 100644 tests/metagpt/actions/test_code_executor.py
 create mode 100644 tests/metagpt/actions/test_plan.py
 create mode 100644 tests/metagpt/actions/test_write_code_v2.py

diff --git a/metagpt/actions/__init__.py b/metagpt/actions/__init__.py
index b004bd58e..d7afae2fe 100644
--- a/metagpt/actions/__init__.py
+++ b/metagpt/actions/__init__.py
@@ -23,6 +23,9 @@ from metagpt.actions.write_code_review import WriteCodeReview
 from metagpt.actions.write_prd import WritePRD
 from metagpt.actions.write_prd_review import WritePRDReview
 from metagpt.actions.write_test import WriteTest
+from metagpt.actions.code_executor import PyCodeExecutor
+from metagpt.actions.write_code_v2 import WriteCode as WriteCodeFunction
+from metagpt.actions.plan import Plan
 
 
 class ActionType(Enum):
@@ -45,6 +48,9 @@ class ActionType(Enum):
     COLLECT_LINKS = CollectLinks
     WEB_BROWSE_AND_SUMMARIZE = WebBrowseAndSummarize
     CONDUCT_RESEARCH = ConductResearch
+    PYCODE_EXECUTOR = PyCodeExecutor
+    WRITE_CODE_FUNCTION = WriteCodeFunction
+    PLAN = Plan
 
 
 __all__ = [
diff --git a/metagpt/actions/code_executor.py b/metagpt/actions/code_executor.py
new file mode 100644
index 000000000..c05c00c9c
--- /dev/null
+++ b/metagpt/actions/code_executor.py
@@ -0,0 +1,173 @@
+# -*- encoding: utf-8 -*-
+"""
+@Date    :   2023/11/17 14:22:15
+@Author  :   orange-crow
+@File    :   code_executor.py
+"""
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Dict, List, Tuple, Union
+
+import nbformat
+from nbclient import NotebookClient
+from nbformat.v4 import new_code_cell, new_output
+from rich.console import Console
+from rich.syntax import Syntax
+
+from metagpt.actions import Action
+from metagpt.schema import Message
+
+
+class CodeExecutor(ABC):
+    @abstractmethod
+    async def build(self):
+        """build code executor"""
+        ...
+
+    @abstractmethod
+    async def run(self, code: str):
+        """run code"""
+        ...
+
+    @abstractmethod
+    async def terminate(self):
+        """terminate executor"""
+        ...
+
+    @abstractmethod
+    async def reset(self):
+        """reset executor"""
+        ...
+
+
+class PyCodeExecutor(CodeExecutor, Action):
+    """execute code, return result to llm, and display it."""
+
+    def __init__(self, name: str = "python_executor", context=None, llm=None):
+        super().__init__(name, context, llm)
+        self.nb = nbformat.v4.new_notebook()
+        self.nb_client = NotebookClient(self.nb)
+        self.console = Console()
+        self.interaction = "ipython" if self.is_ipython() else "terminal"
+
+    async def build(self):
+        if self.nb_client.kc is None or not await self.nb_client.kc.is_alive():
+            self.nb_client.create_kernel_manager()
+            self.nb_client.start_new_kernel()
+            self.nb_client.start_new_kernel_client()
+
+    async def terminate(self):
+        """kill NotebookClient"""
+        await self.nb_client._async_cleanup_kernel()
+
+    async def reset(self):
+        """reset NotebookClient"""
+        await self.terminate()
+        self.nb_client = NotebookClient(self.nb)
+
+    def add_code_cell(self, code):
+        self.nb.cells.append(new_code_cell(source=code))
+
+    def _display(self, code, language: str = "python"):
+        if language == "python":
+            code = Syntax(code, "python", theme="paraiso-dark", line_numbers=True)
+            self.console.print("\n")
+            self.console.print(code)
+
+    def add_output_to_cell(self, cell, output):
+        if "outputs" not in cell:
+            cell["outputs"] = []
+        # TODO: show figures
+        else:
+            cell["outputs"].append(new_output(output_type="stream", name="stdout", text=str(output)))
+
+    def parse_outputs(self, outputs: List) -> str:
+        assert isinstance(outputs, list)
+        parsed_output = {"text": [], "image": []}
+
+        # empty outputs: such as 'x=1\ny=2'
+        if not outputs:
+            return parsed_output
+
+        for output in outputs:
+            if output["output_type"] == "stream":
+                parsed_output["text"].append(output["text"])
+            elif output["output_type"] == "display_data":
+                self.show_bytes_figure(output["data"]["image/png"], self.interaction)
+                parsed_output["image"].append(output["data"]["image/png"])
+        return str(parsed_output)
+
+    def show_bytes_figure(self, image_base64: str, interaction_type: str = "ipython"):
+        import base64
+
+        image_bytes = base64.b64decode(image_base64)
+        if interaction_type == "ipython":
+            from IPython.display import Image, display
+
+            display(Image(data=image_bytes))
+        else:
+            import io
+
+            from PIL import Image
+
+            image = Image.open(io.BytesIO(image_bytes))
+            image.show()
+
+    def is_ipython(self) -> bool:
+        try:
+            # 如果在Jupyter Notebook中运行，__file__ 变量不存在
+            from IPython import get_ipython
+
+            if get_ipython() is not None and "IPKernelApp" in get_ipython().config:
+                return True
+            else:
+                return False
+        except NameError:
+            # 如果在Python脚本中运行，__file__ 变量存在
+            return False
+
+    def _process_code(self, code: Union[str, Dict, Message], language: str = None) -> Tuple:
+        if isinstance(code, str) and Path(code).suffix in (".py", ".txt"):
+            code = Path(code).read_text(encoding="utf-8")
+            return code, language
+
+        if isinstance(code, str):
+            return code, language
+
+        if isinstance(code, dict):
+            assert "code" in code
+            assert "language" in code
+            code, language = code["code"], code["language"]
+        elif isinstance(code, Message):
+            assert hasattr(code, "language")
+            code, language = code.content, code.language
+        else:
+            raise ValueError(f"Not support code type {type(code).__name__}.")
+
+        return code, language
+
+    async def run(self, code: Union[str, Dict, Message], language: str = "python") -> Message:
+        code, language = self._process_code(code, language)
+
+        self._display(code, language)
+
+        if language == "python":
+            # add code to the notebook
+            self.add_code_cell(code=code)
+            try:
+                # build code executor
+                await self.build()
+                # run code
+                # TODO: add max_tries for run code.
+                cell_index = len(self.nb.cells) - 1
+                await self.nb_client.async_execute_cell(self.nb.cells[-1], cell_index)
+                return Message(
+                    self.parse_outputs(self.nb.cells[-1].outputs), state="done", sent_from=self.__class__.__name__
+                )
+            except Exception as e:
+                # FIXME: CellExecutionError is hard to read. for example `1\0` raise ZeroDivisionError:
+                #  CellExecutionError('An error occurred while executing the following cell:\n------------------\nz=1/0\n------------------\n\n\n\x1b[0;31m---------------------------------------------------------------------------\x1b[0m\n\x1b[0;31mZeroDivisionError\x1b[0m                         Traceback (most recent call last)\nCell \x1b[0;32mIn[1], line 1\x1b[0m\n\x1b[0;32m----> 1\x1b[0m z\x1b[38;5;241m=\x1b[39m\x1b[38;5;241;43m1\x1b[39;49m\x1b[38;5;241;43m/\x1b[39;49m\x1b[38;5;241;43m0\x1b[39;49m\n\n\x1b[0;31mZeroDivisionError\x1b[0m: division by zero\n')
+                return Message(e, state="error", sent_from=self.__class__.__name__)
+        else:
+            # TODO: markdown
+            raise NotImplementedError(f"Not support this code type : {language}, Only support code!")
diff --git a/metagpt/actions/plan.py b/metagpt/actions/plan.py
new file mode 100644
index 000000000..d46783ba2
--- /dev/null
+++ b/metagpt/actions/plan.py
@@ -0,0 +1,20 @@
+# -*- encoding: utf-8 -*-
+"""
+@Date    :   2023/11/20 11:24:03
+@Author  :   orange-crow
+@File    :   plan.py
+"""
+from metagpt.actions import Action
+from metagpt.prompts.plan import TASK_PLAN_SYSTEM_MSG
+from metagpt.schema import Message
+
+
+class Plan(Action):
+    def __init__(self, llm=None):
+        super().__init__("", None, llm)
+
+    async def run(self, prompt: str, role: str = None, system_msg: str = None) -> str:
+        if role:
+            system_msg = TASK_PLAN_SYSTEM_MSG.format(role=role)
+        rsp = await self._aask(system_msg + prompt)
+        return Message(rsp, role="assistant", sent_from=self.__class__.__name__)
diff --git a/metagpt/actions/write_code_v2.py b/metagpt/actions/write_code_v2.py
new file mode 100644
index 000000000..335e70dc0
--- /dev/null
+++ b/metagpt/actions/write_code_v2.py
@@ -0,0 +1,36 @@
+# -*- encoding: utf-8 -*-
+"""
+@Date    :   2023/11/20 13:19:39
+@Author  :   orange-crow
+@File    :   write_code_v2.py
+"""
+from typing import Dict, List, Union
+
+from metagpt.actions import Action
+from metagpt.schema import Message
+
+
+class WriteCode(Action):
+    """Use openai function to generate code."""
+
+    def __init__(self, name: str = "", context=None, llm=None) -> str:
+        super().__init__(name, context, llm)
+
+    def process_msg(self, prompt: Union[str, List[Dict], Message, List[Message]], system_msg: str = None):
+        if isinstance(prompt, str):
+            return system_msg + prompt if system_msg else prompt
+
+        if isinstance(prompt, Message):
+            prompt.content = system_msg + prompt.content if system_msg else prompt.content
+            return prompt
+
+        if isinstance(prompt, list) and system_msg:
+            prompt.insert(0, {"role": "system", "content": system_msg})
+        return prompt
+
+    async def run(
+        self, prompt: Union[str, List[Dict], Message, List[Message]], system_msg: str = None, **kwargs
+    ) -> Dict:
+        prompt = self.process_msg(prompt, system_msg)
+        code_content = await self.llm.aask_code(prompt, **kwargs)
+        return Message(content=code_content, role="assistant")
diff --git a/metagpt/prompts/plan.py b/metagpt/prompts/plan.py
new file mode 100644
index 000000000..c4b056ab0
--- /dev/null
+++ b/metagpt/prompts/plan.py
@@ -0,0 +1,7 @@
+TASK_PLAN_SYSTEM_MSG = """You are a {role}. Write a plan with single digits steps. make sure others can understand what you are doing.
+Example:
+# plan
+1. ...\n\n
+2. ...\n\n
+...
+"""
diff --git a/metagpt/schema.py b/metagpt/schema.py
index bdca093c2..4bada005a 100644
--- a/metagpt/schema.py
+++ b/metagpt/schema.py
@@ -30,6 +30,7 @@ class Message:
     sent_from: str = field(default="")
     send_to: str = field(default="")
     restricted_to: str = field(default="")
+    state: str = None  # None, done, todo, doing, error
 
     def __str__(self):
         # prefix = '-'.join([self.role, str(self.cause_by)])
diff --git a/tests/metagpt/actions/test_code_executor.py b/tests/metagpt/actions/test_code_executor.py
new file mode 100644
index 000000000..d1833b48c
--- /dev/null
+++ b/tests/metagpt/actions/test_code_executor.py
@@ -0,0 +1,58 @@
+import pytest
+
+from metagpt.actions import PyCodeExecutor
+from metagpt.schema import Message
+
+
+@pytest.mark.asyncio
+async def test_code_running():
+    pi = PyCodeExecutor()
+    output = await pi.run("print('hello world!')")
+    assert output.state == "done"
+    output = await pi.run({"code": "print('hello world!')", "language": "python"})
+    assert output.state == "done"
+    code_msg = Message("print('hello world!')")
+    setattr(code_msg, "language", "python")
+    output = await pi.run(code_msg)
+    assert output.state == "done"
+
+
+@pytest.mark.asyncio
+async def test_split_code_running():
+    pi = PyCodeExecutor()
+    output = await pi.run("x=1\ny=2")
+    output = await pi.run("z=x+y")
+    output = await pi.run("assert z==3")
+    assert output.state == "done"
+
+
+@pytest.mark.asyncio
+async def test_execute_error():
+    pi = PyCodeExecutor()
+    output = await pi.run("z=1/0")
+    assert output.state == "error"
+
+
+@pytest.mark.asyncio
+async def test_plotting_code():
+    pi = PyCodeExecutor()
+    code = """
+    import numpy as np
+    import matplotlib.pyplot as plt
+
+    # 生成随机数据
+    random_data = np.random.randn(1000)  # 生成1000个符合标准正态分布的随机数
+
+    # 绘制直方图
+    plt.hist(random_data, bins=30, density=True, alpha=0.7, color='blue', edgecolor='black')
+
+    # 添加标题和标签
+    plt.title('Histogram of Random Data')
+    plt.xlabel('Value')
+    plt.ylabel('Frequency')
+
+    # 显示图形
+    plt.show()
+    """
+    output = await pi.run(code)
+    assert output.state == "done"
diff --git a/tests/metagpt/actions/test_plan.py b/tests/metagpt/actions/test_plan.py
new file mode 100644
index 000000000..35f8f20cc
--- /dev/null
+++ b/tests/metagpt/actions/test_plan.py
@@ -0,0 +1,12 @@
+import pytest
+
+from metagpt.actions.plan import Plan
+
+
+@pytest.mark.asyncio
+async def test_plan():
+    p = Plan()
+    task_desc = """Here’s some background information on Cyclistic, a bike-sharing company designing a marketing strategy aimed at converting casual riders into annual members: So far, Cyclistic’s marketing strategy has relied on building general awareness and engaging a wide range of consumers. group. One way to help achieve these goals is the flexibility of its pricing plans: one-way passes, full-day passes, and annual memberships. Customers who purchase a one-way or full-day pass are known as recreational riders. Customers purchasing an annual membership are Cyclistic members. I will provide you with a data sheet that records user behavior: '/Users/vicis/Downloads/202103-divvy-tripdata.csv"""
+    rsp = await p.run(task_desc, role="data analyst")
+    assert len(rsp.content) > 0
+    assert rsp.sent_from == "Plan"
diff --git a/tests/metagpt/actions/test_write_code_v2.py b/tests/metagpt/actions/test_write_code_v2.py
new file mode 100644
index 000000000..929407051
--- /dev/null
+++ b/tests/metagpt/actions/test_write_code_v2.py
@@ -0,0 +1,22 @@
+import pytest
+
+from metagpt.actions.write_code_v2 import WriteCode
+
+
+@pytest.mark.asyncio
+async def test_write_code():
+    coder = WriteCode()
+    code = await coder.run("Write a hello world code.")
+    assert "language" in code.content
+    assert "code" in code.content
+    print(code)
+
+
+@pytest.mark.asyncio
+async def test_write_code_by_list_prompt():
+    coder = WriteCode()
+    msg = ["a=[1,2,5,10,-10]", "写出求a中最大值的代码python"]
+    code = await coder.run(msg)
+    assert "language" in code.content
+    assert "code" in code.content
+    print(code)

From 50f64ca934d13072910989c20769e740920f7d7d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Tue, 21 Nov 2023 19:14:52 +0800
Subject: [PATCH 02/11] doc: add rich.

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

diff --git a/requirements.txt b/requirements.txt
index f0169d7fa..53176bd0a 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -45,3 +45,4 @@ semantic-kernel==0.3.13.dev0
 wrapt==1.15.0
 websocket-client==0.58.0
 zhipuai==1.0.7
+rich==13.6.0
\ No newline at end of file

From fa400a0b0d43d59afdea037430cc7fac57e34634 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Tue, 21 Nov 2023 19:15:16 +0800
Subject: [PATCH 03/11] chore: rename WriteCode -> WriteCodeFunction.

---
 metagpt/actions/__init__.py                                 | 2 +-
 .../actions/{write_code_v2.py => write_code_function.py}    | 2 +-
 .../{test_write_code_v2.py => test_write_code_function.py}  | 6 +++---
 3 files changed, 5 insertions(+), 5 deletions(-)
 rename metagpt/actions/{write_code_v2.py => write_code_function.py} (97%)
 rename tests/metagpt/actions/{test_write_code_v2.py => test_write_code_function.py} (78%)

diff --git a/metagpt/actions/__init__.py b/metagpt/actions/__init__.py
index d7afae2fe..d0163c24e 100644
--- a/metagpt/actions/__init__.py
+++ b/metagpt/actions/__init__.py
@@ -24,7 +24,7 @@ from metagpt.actions.write_prd import WritePRD
 from metagpt.actions.write_prd_review import WritePRDReview
 from metagpt.actions.write_test import WriteTest
 from metagpt.actions.code_executor import PyCodeExecutor
-from metagpt.actions.write_code_v2 import WriteCode as WriteCodeFunction
+from metagpt.actions.write_code_function import WriteCodeFunction
 from metagpt.actions.plan import Plan
 
 
diff --git a/metagpt/actions/write_code_v2.py b/metagpt/actions/write_code_function.py
similarity index 97%
rename from metagpt/actions/write_code_v2.py
rename to metagpt/actions/write_code_function.py
index 335e70dc0..2d943176a 100644
--- a/metagpt/actions/write_code_v2.py
+++ b/metagpt/actions/write_code_function.py
@@ -10,7 +10,7 @@ from metagpt.actions import Action
 from metagpt.schema import Message
 
 
-class WriteCode(Action):
+class WriteCodeFunction(Action):
     """Use openai function to generate code."""
 
     def __init__(self, name: str = "", context=None, llm=None) -> str:
diff --git a/tests/metagpt/actions/test_write_code_v2.py b/tests/metagpt/actions/test_write_code_function.py
similarity index 78%
rename from tests/metagpt/actions/test_write_code_v2.py
rename to tests/metagpt/actions/test_write_code_function.py
index 929407051..0e57b4ced 100644
--- a/tests/metagpt/actions/test_write_code_v2.py
+++ b/tests/metagpt/actions/test_write_code_function.py
@@ -1,11 +1,11 @@
 import pytest
 
-from metagpt.actions.write_code_v2 import WriteCode
+from metagpt.actions.write_code_function import WriteCodeFunction
 
 
 @pytest.mark.asyncio
 async def test_write_code():
-    coder = WriteCode()
+    coder = WriteCodeFunction()
     code = await coder.run("Write a hello world code.")
     assert "language" in code.content
     assert "code" in code.content
@@ -14,7 +14,7 @@ async def test_write_code():
 
 @pytest.mark.asyncio
 async def test_write_code_by_list_prompt():
-    coder = WriteCode()
+    coder = WriteCodeFunction()
     msg = ["a=[1,2,5,10,-10]", "写出求a中最大值的代码python"]
     code = await coder.run(msg)
     assert "language" in code.content

From 8ef05bb19f2ebbea8df60dc93813957e19cedbac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Wed, 22 Nov 2023 17:56:43 +0800
Subject: [PATCH 04/11] chore: prompt support Message type.

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

diff --git a/metagpt/actions/plan.py b/metagpt/actions/plan.py
index d46783ba2..ab3963c72 100644
--- a/metagpt/actions/plan.py
+++ b/metagpt/actions/plan.py
@@ -4,6 +4,8 @@
 @Author  :   orange-crow
 @File    :   plan.py
 """
+from typing import Union
+
 from metagpt.actions import Action
 from metagpt.prompts.plan import TASK_PLAN_SYSTEM_MSG
 from metagpt.schema import Message
@@ -13,8 +15,8 @@ class Plan(Action):
     def __init__(self, llm=None):
         super().__init__("", None, llm)
 
-    async def run(self, prompt: str, role: str = None, system_msg: str = None) -> str:
+    async def run(self, prompt: Union[str, Message], role: str = None, system_msg: str = None) -> str:
         if role:
             system_msg = TASK_PLAN_SYSTEM_MSG.format(role=role)
-        rsp = await self._aask(system_msg + prompt)
+        rsp = self._aask(system_msg + prompt.content) if isinstance(prompt, Message) else await self._aask(system_msg + prompt)
         return Message(rsp, role="assistant", sent_from=self.__class__.__name__)

From 9f5108a4643bf4b27f6abe1b7543c02be937ec4b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Wed, 22 Nov 2023 19:36:21 +0800
Subject: [PATCH 05/11] chore: return plan by list.

---
 metagpt/actions/plan.py            | 4 +++-
 metagpt/prompts/plan.py            | 5 +++--
 tests/metagpt/actions/test_plan.py | 1 +
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/metagpt/actions/plan.py b/metagpt/actions/plan.py
index ab3963c72..8bc575992 100644
--- a/metagpt/actions/plan.py
+++ b/metagpt/actions/plan.py
@@ -9,6 +9,7 @@ from typing import Union
 from metagpt.actions import Action
 from metagpt.prompts.plan import TASK_PLAN_SYSTEM_MSG
 from metagpt.schema import Message
+from metagpt.utils.common import CodeParser
 
 
 class Plan(Action):
@@ -19,4 +20,5 @@ class Plan(Action):
         if role:
             system_msg = TASK_PLAN_SYSTEM_MSG.format(role=role)
         rsp = self._aask(system_msg + prompt.content) if isinstance(prompt, Message) else await self._aask(system_msg + prompt)
-        return Message(rsp, role="assistant", sent_from=self.__class__.__name__)
+        plan = CodeParser.parse_code(None, rsp).split('\n\n')
+        return Message(plan, role="assistant", sent_from=self.__class__.__name__)
diff --git a/metagpt/prompts/plan.py b/metagpt/prompts/plan.py
index c4b056ab0..4d3add211 100644
--- a/metagpt/prompts/plan.py
+++ b/metagpt/prompts/plan.py
@@ -1,7 +1,8 @@
 TASK_PLAN_SYSTEM_MSG = """You are a {role}. Write a plan with single digits steps. make sure others can understand what you are doing.
-Example:
-# plan
+Example, must start with ```, and end with ```:
+```
 1. ...\n\n
 2. ...\n\n
 ...
+```
 """
diff --git a/tests/metagpt/actions/test_plan.py b/tests/metagpt/actions/test_plan.py
index 35f8f20cc..1b1b90513 100644
--- a/tests/metagpt/actions/test_plan.py
+++ b/tests/metagpt/actions/test_plan.py
@@ -10,3 +10,4 @@ async def test_plan():
     rsp = await p.run(task_desc, role="data analyst")
     assert len(rsp.content) > 0
     assert rsp.sent_from == "Plan"
+    print(rsp)

From 8a0a89e604241187fdd253851678c6a16a8a4bbd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Wed, 22 Nov 2023 20:33:57 +0800
Subject: [PATCH 06/11] fix: fix bug about message.

---
 metagpt/actions/write_code_function.py | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/metagpt/actions/write_code_function.py b/metagpt/actions/write_code_function.py
index 2d943176a..6fb7f535e 100644
--- a/metagpt/actions/write_code_function.py
+++ b/metagpt/actions/write_code_function.py
@@ -21,16 +21,31 @@ class WriteCodeFunction(Action):
             return system_msg + prompt if system_msg else prompt
 
         if isinstance(prompt, Message):
-            prompt.content = system_msg + prompt.content if system_msg else prompt.content
+            if isinstance(prompt.content, dict):
+                prompt.content = system_msg + str([(k, v) for k, v in prompt.content.items()])\
+                    if system_msg else prompt.content
+            else:
+                prompt.content = system_msg + prompt.content if system_msg else prompt.content
             return prompt
 
+        if isinstance(prompt, list):
+            _prompt = []
+            for msg in prompt:
+                if isinstance(msg, Message) and isinstance(msg.content, dict):
+                    msg.content = str([(k, v) for k, v in msg.content.items()])
+                if isinstance(msg, Message):
+                    msg = msg.to_dict()
+                _prompt.append(msg)
+            prompt = _prompt
+
         if isinstance(prompt, list) and system_msg:
-            prompt.insert(0, {"role": "system", "content": system_msg})
+            if system_msg not in prompt[0]['content']:
+                prompt[0]['content'] = system_msg + prompt[0]['content']
         return prompt
 
     async def run(
         self, prompt: Union[str, List[Dict], Message, List[Message]], system_msg: str = None, **kwargs
-    ) -> Dict:
+    ) -> Message:
         prompt = self.process_msg(prompt, system_msg)
         code_content = await self.llm.aask_code(prompt, **kwargs)
         return Message(content=code_content, role="assistant")

From d8ddf1fcb0a516269c8a041ba2ca1a36931af87e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Wed, 22 Nov 2023 20:35:47 +0800
Subject: [PATCH 07/11] add new test for list plan.

---
 .../actions/test_write_code_function.py       | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/tests/metagpt/actions/test_write_code_function.py b/tests/metagpt/actions/test_write_code_function.py
index 0e57b4ced..cac459380 100644
--- a/tests/metagpt/actions/test_write_code_function.py
+++ b/tests/metagpt/actions/test_write_code_function.py
@@ -1,6 +1,7 @@
 import pytest
 
 from metagpt.actions.write_code_function import WriteCodeFunction
+from metagpt.actions.code_executor import PyCodeExecutor
 
 
 @pytest.mark.asyncio
@@ -20,3 +21,21 @@ async def test_write_code_by_list_prompt():
     assert "language" in code.content
     assert "code" in code.content
     print(code)
+
+
+@pytest.mark.asyncio
+async def test_write_code_by_list_plan():
+    coder = WriteCodeFunction()
+    executor = PyCodeExecutor()
+    messages = []
+    plan = ["随机生成一个pandas DataFrame时间序列", "绘制这个时间序列的直方图", "求均值"]
+    for task in plan:
+        print(f"\n任务: {task}\n\n")
+        messages.append(task)
+        code = await coder.run(messages)
+        messages.append(code)
+        assert "language" in code.content
+        assert "code" in code.content
+        output = await executor.run(code)
+        print(f"\n[Output]: 任务{task}的执行结果是: \n{output}\n")
+        messages.append(output)

From 7b94c04f51b5cca61143ddad67dbc87a40fc2fc8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Wed, 22 Nov 2023 20:36:51 +0800
Subject: [PATCH 08/11] fix: return string in parse_outputs.

---
 metagpt/actions/code_executor.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/metagpt/actions/code_executor.py b/metagpt/actions/code_executor.py
index c05c00c9c..0b4f5171f 100644
--- a/metagpt/actions/code_executor.py
+++ b/metagpt/actions/code_executor.py
@@ -83,7 +83,7 @@ class PyCodeExecutor(CodeExecutor, Action):
 
     def parse_outputs(self, outputs: List) -> str:
         assert isinstance(outputs, list)
-        parsed_output = {"text": [], "image": []}
+        parsed_output = ""
 
         # empty outputs: such as 'x=1\ny=2'
         if not outputs:
@@ -91,11 +91,12 @@ class PyCodeExecutor(CodeExecutor, Action):
 
         for output in outputs:
             if output["output_type"] == "stream":
-                parsed_output["text"].append(output["text"])
+                parsed_output += output["text"]
             elif output["output_type"] == "display_data":
                 self.show_bytes_figure(output["data"]["image/png"], self.interaction)
-                parsed_output["image"].append(output["data"]["image/png"])
-        return str(parsed_output)
+            elif output["output_type"] == "execute_result":
+                parsed_output += output["data"]["text/plain"]
+        return parsed_output
 
     def show_bytes_figure(self, image_base64: str, interaction_type: str = "ipython"):
         import base64
@@ -139,8 +140,8 @@ class PyCodeExecutor(CodeExecutor, Action):
             assert "language" in code
             code, language = code["code"], code["language"]
         elif isinstance(code, Message):
-            assert hasattr(code, "language")
-            code, language = code.content, code.language
+            assert "language" in code.content
+            code, language = code.content["code"], code.content["language"]
         else:
             raise ValueError(f"Not support code type {type(code).__name__}.")
 

From a0b13c8e0ff4b4647585780f75644aff3b64471e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Thu, 23 Nov 2023 10:45:40 +0800
Subject: [PATCH 09/11] chore: change name.

---
 metagpt/actions/__init__.py                    |  8 ++++----
 .../{code_executor.py => execute_code.py}      | 15 ++++++++++-----
 metagpt/actions/{plan.py => write_plan.py}     |  2 +-
 ...t_code_executor.py => test_execute_code.py} | 11 +++++------
 .../actions/test_write_code_function.py        | 18 +++++++++---------
 .../{test_plan.py => test_write_plan.py}       |  6 +++---
 6 files changed, 32 insertions(+), 28 deletions(-)
 rename metagpt/actions/{code_executor.py => execute_code.py} (92%)
 rename metagpt/actions/{plan.py => write_plan.py} (97%)
 rename tests/metagpt/actions/{test_code_executor.py => test_execute_code.py} (87%)
 rename tests/metagpt/actions/{test_plan.py => test_write_plan.py} (88%)

diff --git a/metagpt/actions/__init__.py b/metagpt/actions/__init__.py
index d0163c24e..ba2170cbd 100644
--- a/metagpt/actions/__init__.py
+++ b/metagpt/actions/__init__.py
@@ -23,9 +23,9 @@ from metagpt.actions.write_code_review import WriteCodeReview
 from metagpt.actions.write_prd import WritePRD
 from metagpt.actions.write_prd_review import WritePRDReview
 from metagpt.actions.write_test import WriteTest
-from metagpt.actions.code_executor import PyCodeExecutor
+from metagpt.actions.execute_code import ExecutePyCode
 from metagpt.actions.write_code_function import WriteCodeFunction
-from metagpt.actions.plan import Plan
+from metagpt.actions.write_plan import WritePlan
 
 
 class ActionType(Enum):
@@ -48,9 +48,9 @@ class ActionType(Enum):
     COLLECT_LINKS = CollectLinks
     WEB_BROWSE_AND_SUMMARIZE = WebBrowseAndSummarize
     CONDUCT_RESEARCH = ConductResearch
-    PYCODE_EXECUTOR = PyCodeExecutor
+    EXECUTE_PYCODE = ExecutePyCode
     WRITE_CODE_FUNCTION = WriteCodeFunction
-    PLAN = Plan
+    WRITE_PLAN = WritePlan
 
 
 __all__ = [
diff --git a/metagpt/actions/code_executor.py b/metagpt/actions/execute_code.py
similarity index 92%
rename from metagpt/actions/code_executor.py
rename to metagpt/actions/execute_code.py
index 0b4f5171f..e80886c3e 100644
--- a/metagpt/actions/code_executor.py
+++ b/metagpt/actions/execute_code.py
@@ -18,7 +18,7 @@ from metagpt.actions import Action
 from metagpt.schema import Message
 
 
-class CodeExecutor(ABC):
+class ExecuteCode(ABC):
     @abstractmethod
     async def build(self):
         """build code executor"""
@@ -40,7 +40,7 @@ class CodeExecutor(ABC):
         ...
 
 
-class PyCodeExecutor(CodeExecutor, Action):
+class ExecutePyCode(ExecuteCode, Action):
     """execute code, return result to llm, and display it."""
 
     def __init__(self, name: str = "python_executor", context=None, llm=None):
@@ -128,6 +128,7 @@ class PyCodeExecutor(CodeExecutor, Action):
             return False
 
     def _process_code(self, code: Union[str, Dict, Message], language: str = None) -> Tuple:
+        language = language or 'python'
         if isinstance(code, str) and Path(code).suffix in (".py", ".txt"):
             code = Path(code).read_text(encoding="utf-8")
             return code, language
@@ -137,11 +138,15 @@ class PyCodeExecutor(CodeExecutor, Action):
 
         if isinstance(code, dict):
             assert "code" in code
-            assert "language" in code
+            if "language" not in code:
+                code['language'] = 'python'
             code, language = code["code"], code["language"]
         elif isinstance(code, Message):
-            assert "language" in code.content
-            code, language = code.content["code"], code.content["language"]
+            if isinstance(code.content, dict) and "language" not in code.content:
+                code.content["language"] = 'python'
+                code, language = code.content["code"], code.content["language"]
+            elif isinstance(code.content, str):
+                code, language = code.content, language
         else:
             raise ValueError(f"Not support code type {type(code).__name__}.")
 
diff --git a/metagpt/actions/plan.py b/metagpt/actions/write_plan.py
similarity index 97%
rename from metagpt/actions/plan.py
rename to metagpt/actions/write_plan.py
index 8bc575992..96d15cb84 100644
--- a/metagpt/actions/plan.py
+++ b/metagpt/actions/write_plan.py
@@ -12,7 +12,7 @@ from metagpt.schema import Message
 from metagpt.utils.common import CodeParser
 
 
-class Plan(Action):
+class WritePlan(Action):
     def __init__(self, llm=None):
         super().__init__("", None, llm)
 
diff --git a/tests/metagpt/actions/test_code_executor.py b/tests/metagpt/actions/test_execute_code.py
similarity index 87%
rename from tests/metagpt/actions/test_code_executor.py
rename to tests/metagpt/actions/test_execute_code.py
index d1833b48c..88c5adf18 100644
--- a/tests/metagpt/actions/test_code_executor.py
+++ b/tests/metagpt/actions/test_execute_code.py
@@ -1,25 +1,24 @@
 import pytest
 
-from metagpt.actions import PyCodeExecutor
+from metagpt.actions import ExecutePyCode
 from metagpt.schema import Message
 
 
 @pytest.mark.asyncio
 async def test_code_running():
-    pi = PyCodeExecutor()
+    pi = ExecutePyCode()
     output = await pi.run("print('hello world!')")
     assert output.state == "done"
     output = await pi.run({"code": "print('hello world!')", "language": "python"})
     assert output.state == "done"
     code_msg = Message("print('hello world!')")
-    setattr(code_msg, "language", "python")
     output = await pi.run(code_msg)
     assert output.state == "done"
 
 
 @pytest.mark.asyncio
 async def test_split_code_running():
-    pi = PyCodeExecutor()
+    pi = ExecutePyCode()
     output = await pi.run("x=1\ny=2")
     output = await pi.run("z=x+y")
     output = await pi.run("assert z==3")
@@ -28,14 +27,14 @@ async def test_split_code_running():
 
 @pytest.mark.asyncio
 async def test_execute_error():
-    pi = PyCodeExecutor()
+    pi = ExecutePyCode()
     output = await pi.run("z=1/0")
     assert output.state == "error"
 
 
 @pytest.mark.asyncio
 async def test_plotting_code():
-    pi = PyCodeExecutor()
+    pi = ExecutePyCode()
     code = """
     import numpy as np
     import matplotlib.pyplot as plt
diff --git a/tests/metagpt/actions/test_write_code_function.py b/tests/metagpt/actions/test_write_code_function.py
index cac459380..4ff1a63c4 100644
--- a/tests/metagpt/actions/test_write_code_function.py
+++ b/tests/metagpt/actions/test_write_code_function.py
@@ -1,13 +1,13 @@
 import pytest
 
 from metagpt.actions.write_code_function import WriteCodeFunction
-from metagpt.actions.code_executor import PyCodeExecutor
+from metagpt.actions.execute_code import ExecutePyCode
 
 
 @pytest.mark.asyncio
 async def test_write_code():
-    coder = WriteCodeFunction()
-    code = await coder.run("Write a hello world code.")
+    write_code = WriteCodeFunction()
+    code = await write_code.run("Write a hello world code.")
     assert "language" in code.content
     assert "code" in code.content
     print(code)
@@ -15,9 +15,9 @@ async def test_write_code():
 
 @pytest.mark.asyncio
 async def test_write_code_by_list_prompt():
-    coder = WriteCodeFunction()
+    write_code = WriteCodeFunction()
     msg = ["a=[1,2,5,10,-10]", "写出求a中最大值的代码python"]
-    code = await coder.run(msg)
+    code = await write_code.run(msg)
     assert "language" in code.content
     assert "code" in code.content
     print(code)
@@ -25,17 +25,17 @@ async def test_write_code_by_list_prompt():
 
 @pytest.mark.asyncio
 async def test_write_code_by_list_plan():
-    coder = WriteCodeFunction()
-    executor = PyCodeExecutor()
+    write_code = WriteCodeFunction()
+    execute_code = ExecutePyCode()
     messages = []
     plan = ["随机生成一个pandas DataFrame时间序列", "绘制这个时间序列的直方图", "求均值"]
     for task in plan:
         print(f"\n任务: {task}\n\n")
         messages.append(task)
-        code = await coder.run(messages)
+        code = await write_code.run(messages)
         messages.append(code)
         assert "language" in code.content
         assert "code" in code.content
-        output = await executor.run(code)
+        output = await execute_code.run(code)
         print(f"\n[Output]: 任务{task}的执行结果是: \n{output}\n")
         messages.append(output)
diff --git a/tests/metagpt/actions/test_plan.py b/tests/metagpt/actions/test_write_plan.py
similarity index 88%
rename from tests/metagpt/actions/test_plan.py
rename to tests/metagpt/actions/test_write_plan.py
index 1b1b90513..2bf200ab3 100644
--- a/tests/metagpt/actions/test_plan.py
+++ b/tests/metagpt/actions/test_write_plan.py
@@ -1,13 +1,13 @@
 import pytest
 
-from metagpt.actions.plan import Plan
+from metagpt.actions.write_plan import WritePlan
 
 
 @pytest.mark.asyncio
 async def test_plan():
-    p = Plan()
+    p = WritePlan()
     task_desc = """Here’s some background information on Cyclistic, a bike-sharing company designing a marketing strategy aimed at converting casual riders into annual members: So far, Cyclistic’s marketing strategy has relied on building general awareness and engaging a wide range of consumers. group. One way to help achieve these goals is the flexibility of its pricing plans: one-way passes, full-day passes, and annual memberships. Customers who purchase a one-way or full-day pass are known as recreational riders. Customers purchasing an annual membership are Cyclistic members. I will provide you with a data sheet that records user behavior: '/Users/vicis/Downloads/202103-divvy-tripdata.csv"""
     rsp = await p.run(task_desc, role="data analyst")
     assert len(rsp.content) > 0
-    assert rsp.sent_from == "Plan"
+    assert rsp.sent_from == "WritePlan"
     print(rsp)

From 3d18dfe2b582f16cf08f6b4e23eea56e85ee1c59 Mon Sep 17 00:00:00 2001
From: yzlin <yzlin@fuzhi.ai>
Date: Thu, 23 Nov 2023 21:59:25 +0800
Subject: [PATCH 10/11] pipeline first version

---
 metagpt/actions/execute_code.py        |   9 +-
 metagpt/actions/write_code_function.py |  22 +++--
 metagpt/actions/write_plan.py          |  46 ++++++++---
 metagpt/prompts/plan.py                |   8 --
 metagpt/roles/ml_engineer.py           | 110 +++++++++++++++++++++++++
 metagpt/schema.py                      | 109 ++++++++++++++++++++++++
 requirements.txt                       |   6 +-
 tests/metagpt/test_schema.py           |  85 +++++++++++++++++++
 8 files changed, 362 insertions(+), 33 deletions(-)
 delete mode 100644 metagpt/prompts/plan.py
 create mode 100644 metagpt/roles/ml_engineer.py

diff --git a/metagpt/actions/execute_code.py b/metagpt/actions/execute_code.py
index e80886c3e..7b16d559a 100644
--- a/metagpt/actions/execute_code.py
+++ b/metagpt/actions/execute_code.py
@@ -7,6 +7,7 @@
 from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Dict, List, Tuple, Union
+import traceback
 
 import nbformat
 from nbclient import NotebookClient
@@ -152,7 +153,7 @@ class ExecutePyCode(ExecuteCode, Action):
 
         return code, language
 
-    async def run(self, code: Union[str, Dict, Message], language: str = "python") -> Message:
+    async def run(self, code: Union[str, Dict, Message], language: str = "python") -> Tuple[str, bool]:
         code, language = self._process_code(code, language)
 
         self._display(code, language)
@@ -167,13 +168,11 @@ class ExecutePyCode(ExecuteCode, Action):
                 # TODO: add max_tries for run code.
                 cell_index = len(self.nb.cells) - 1
                 await self.nb_client.async_execute_cell(self.nb.cells[-1], cell_index)
-                return Message(
-                    self.parse_outputs(self.nb.cells[-1].outputs), state="done", sent_from=self.__class__.__name__
-                )
+                return self.parse_outputs(self.nb.cells[-1].outputs), True
             except Exception as e:
                 # FIXME: CellExecutionError is hard to read. for example `1\0` raise ZeroDivisionError:
                 #  CellExecutionError('An error occurred while executing the following cell:\n------------------\nz=1/0\n------------------\n\n\n\x1b[0;31m---------------------------------------------------------------------------\x1b[0m\n\x1b[0;31mZeroDivisionError\x1b[0m                         Traceback (most recent call last)\nCell \x1b[0;32mIn[1], line 1\x1b[0m\n\x1b[0;32m----> 1\x1b[0m z\x1b[38;5;241m=\x1b[39m\x1b[38;5;241;43m1\x1b[39;49m\x1b[38;5;241;43m/\x1b[39;49m\x1b[38;5;241;43m0\x1b[39;49m\n\n\x1b[0;31mZeroDivisionError\x1b[0m: division by zero\n')
-                return Message(e, state="error", sent_from=self.__class__.__name__)
+                return traceback.format_exc(), False
         else:
             # TODO: markdown
             raise NotImplementedError(f"Not support this code type : {language}, Only support code!")
diff --git a/metagpt/actions/write_code_function.py b/metagpt/actions/write_code_function.py
index 6fb7f535e..4ec565eb1 100644
--- a/metagpt/actions/write_code_function.py
+++ b/metagpt/actions/write_code_function.py
@@ -7,10 +7,20 @@
 from typing import Dict, List, Union
 
 from metagpt.actions import Action
-from metagpt.schema import Message
+from metagpt.schema import Message, Plan
 
+class BaseWriteAnalysisCode(Action):
 
-class WriteCodeFunction(Action):
+    async def run(self, context: List[Message], plan: Plan = None, task_guidance: 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_guidance (str, optional): suggested step breakdown for the current task. Defaults to "".
+        """
+
+class WriteCodeFunction(BaseWriteAnalysisCode):
     """Use openai function to generate code."""
 
     def __init__(self, name: str = "", context=None, llm=None) -> str:
@@ -44,8 +54,8 @@ class WriteCodeFunction(Action):
         return prompt
 
     async def run(
-        self, prompt: Union[str, List[Dict], Message, List[Message]], system_msg: str = None, **kwargs
-    ) -> Message:
-        prompt = self.process_msg(prompt, system_msg)
+        self, context: [List[Message]], plan: Plan = None, task_guidance: str = "", system_msg: str = None, **kwargs
+    ) -> str:
+        prompt = self.process_msg(context, system_msg)
         code_content = await self.llm.aask_code(prompt, **kwargs)
-        return Message(content=code_content, role="assistant")
+        return code_content
diff --git a/metagpt/actions/write_plan.py b/metagpt/actions/write_plan.py
index 96d15cb84..48cb1aad5 100644
--- a/metagpt/actions/write_plan.py
+++ b/metagpt/actions/write_plan.py
@@ -4,21 +4,41 @@
 @Author  :   orange-crow
 @File    :   plan.py
 """
-from typing import Union
+from typing import List
+import json
 
 from metagpt.actions import Action
-from metagpt.prompts.plan import TASK_PLAN_SYSTEM_MSG
-from metagpt.schema import Message
-from metagpt.utils.common import CodeParser
-
+from metagpt.schema import Message, Task
 
 class WritePlan(Action):
-    def __init__(self, llm=None):
-        super().__init__("", None, llm)
+    PROMPT_TEMPLATE = """
+    # Context:
+    __context__
+    # Current Plan:
+    __current_plan__
+    # Task:
+    Based on the context, write a plan or modify an existing plan of what you should do to achieve the goal. A plan consists of one to __max_tasks__ tasks.
+    If you are modifying an existing plan, carefully follow the instruction, don't make unnecessary changes.
+    Output a list of jsons following the format:
+    [
+        {
+            "task_id": str = "unique identifier for a task in plan, can be a ordinal",
+            "dependent_task_ids": list[str] = "ids of tasks prerequisite to this task",
+            "instruction": "what you should do in this task, one short phrase or sentence",
+        },
+        ...
+    ]
+    """
+    async def run(self, context: List[Message], current_plan: str = "", max_tasks: int = 5) -> str:
+        prompt = (
+            self.PROMPT_TEMPLATE.replace("__context__", "\n".join([str(ct) for ct in context]))
+            .replace("__current_plan__", current_plan).replace("__max_tasks__", str(max_tasks))
+        )
+        rsp = await self._aask(prompt)
+        return rsp
 
-    async def run(self, prompt: Union[str, Message], role: str = None, system_msg: str = None) -> str:
-        if role:
-            system_msg = TASK_PLAN_SYSTEM_MSG.format(role=role)
-        rsp = self._aask(system_msg + prompt.content) if isinstance(prompt, Message) else await self._aask(system_msg + prompt)
-        plan = CodeParser.parse_code(None, rsp).split('\n\n')
-        return Message(plan, role="assistant", sent_from=self.__class__.__name__)
+    @staticmethod
+    def rsp_to_tasks(rsp: str) -> List[Task]:
+        rsp = json.loads(rsp)
+        tasks = [Task(**task_config) for task_config in rsp]
+        return tasks
diff --git a/metagpt/prompts/plan.py b/metagpt/prompts/plan.py
deleted file mode 100644
index 4d3add211..000000000
--- a/metagpt/prompts/plan.py
+++ /dev/null
@@ -1,8 +0,0 @@
-TASK_PLAN_SYSTEM_MSG = """You are a {role}. Write a plan with single digits steps. make sure others can understand what you are doing.
-Example, must start with ```, and end with ```:
-```
-1. ...\n\n
-2. ...\n\n
-...
-```
-"""
diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
new file mode 100644
index 000000000..c795bda11
--- /dev/null
+++ b/metagpt/roles/ml_engineer.py
@@ -0,0 +1,110 @@
+from typing import Dict, List, Union
+import json
+import subprocess
+
+import fire
+
+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_code_function import WriteCodeFunction
+from metagpt.actions.execute_code import ExecutePyCode
+
+class AskReview(Action):
+
+    async def run(self, context: List[Message], plan: Plan = None):
+        prompt = "\n".join(
+            [f"{msg.cause_by() if msg.cause_by else 'Main Requirement'}: {msg.content}" for msg in context]
+        )
+
+        latest_action = context[-1].cause_by()
+
+        prompt += f"\nPlease review output from {latest_action}, " \
+            "provide feedback or type YES to continue with the process:\n"
+        rsp = input(prompt)
+        confirmed = "yes" in rsp.lower()
+        return rsp, confirmed
+
+
+class MLEngineer(Role):
+    def __init__(self, name="ABC", profile="MLEngineer"):
+        super().__init__(name=name, profile=profile)
+        self._set_react_mode(react_mode="plan_and_act")
+        self.plan = Plan()
+
+    async def _plan_and_act(self):
+
+        # 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 = 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
+                task.result = result
+                self.plan.finish_current_task()
+
+            else:
+                # update plan according to user's feedback and to take on changed tasks
+                await self._update_plan()
+
+    async def _write_and_exec_code(self, max_retry: int = 3):
+        counter = 0
+        success = False
+        while not success and counter < max_retry:
+            context = self.get_memories()
+
+            code = "print('abc')"
+            # code = await WriteCodeFunction().run(context=context)
+            # code = await WriteCodeWithOps.run(context, task, result)
+            self._rc.memory.add(Message(content=code, role="assistant", cause_by=WriteCodeFunction))
+
+            result, success = await ExecutePyCode().run(code)
+            self._rc.memory.add(Message(content=result, role="assistant", cause_by=ExecutePyCode))
+
+            # if not success:
+            #     await self._ask_review()
+
+            counter += 1
+
+        return code, result, success
+
+    async def _ask_review(self):
+        context = self.get_memories()
+        review, confirmed = await AskReview().run(context=context[-5:], plan=self.plan)
+        self._rc.memory.add(Message(content=review, role="assistant", cause_by=AskReview))
+        return confirmed
+
+    async def _update_plan(self, max_tasks: int = 3):
+        current_plan = str([task.json() for task in self.plan.tasks])
+        plan_confirmed = False
+        while not plan_confirmed:
+            context = self.get_memories()
+            rsp = await WritePlan().run(context, current_plan=current_plan, max_tasks=max_tasks)
+            self._rc.memory.add(Message(content=rsp, role="assistant", cause_by=WritePlan))
+            plan_confirmed = await self._ask_review()
+
+        tasks = WritePlan.rsp_to_tasks(rsp)
+        self.plan.add_tasks(tasks)
+
+
+if __name__ == "__main__":
+    # requirement = "create a normal distribution and visualize it"
+    requirement = "run some analysis on iris dataset"
+
+    async def main(requirement: str = requirement):
+        role = MLEngineer()
+        await role.run(requirement)
+
+    fire.Fire(main)
diff --git a/metagpt/schema.py b/metagpt/schema.py
index 4bada005a..3cd7d9730 100644
--- a/metagpt/schema.py
+++ b/metagpt/schema.py
@@ -73,6 +73,115 @@ class AIMessage(Message):
         super().__init__(content, 'assistant')
 
 
+class Task(BaseModel):
+    task_id: str = ""
+    dependent_task_ids: list[str] = [] # Tasks prerequisite to this Task
+    instruction: str = ""
+    task_type: str = ""
+    code: str = ""
+    result: str = ""
+    is_finished: bool = False
+
+
+class Plan(BaseModel):
+    tasks: list[Task] = []
+    task_map: dict[str, Task] = {}
+    current_task_id = ""
+
+    def _topological_sort(self, tasks: list[Task]):
+        task_map = {task.task_id: task for task in tasks}
+        dependencies = {task.task_id: set(task.dependent_task_ids) for task in tasks}
+        sorted_tasks = []
+        visited = set()
+
+        def visit(task_id):
+            if task_id in visited:
+                return
+            visited.add(task_id)
+            for dependent_id in dependencies.get(task_id, []):
+                visit(dependent_id)
+            sorted_tasks.append(task_map[task_id])
+
+        for task in tasks:
+            visit(task.task_id)
+
+        return sorted_tasks
+
+    def add_tasks(self, tasks: list[Task]):
+        """
+        Integrates new tasks into the existing plan, ensuring dependency order is maintained.
+        
+        This method performs two primary functions based on the current state of the task list:
+        1. If there are no existing tasks, it topologically sorts the provided tasks to ensure 
+        correct execution order based on dependencies, and sets these as the current tasks.
+        2. If there are existing tasks, it merges the new tasks with the existing ones. It maintains 
+        any common prefix of tasks (based on task_id and instruction) and appends the remainder 
+        of the new tasks. The current task is updated to the first unfinished task in this merged list.
+
+        Args:
+            tasks (list[Task]): A list of tasks (may be unordered) to add to the plan.
+
+        Returns:
+            None: The method updates the internal state of the plan but does not return anything.
+        """
+        if not tasks:
+            return
+
+        # Topologically sort the new tasks to ensure correct dependency order
+        new_tasks = self._topological_sort(tasks)
+
+        if not self.tasks:
+            # If there are no existing tasks, set the new tasks as the current tasks
+            self.tasks = new_tasks
+
+        else:
+            # Find the length of the common prefix between existing and new tasks
+            prefix_length = 0
+            for old_task, new_task in zip(self.tasks, new_tasks):
+                if old_task.task_id != new_task.task_id or old_task.instruction != new_task.instruction:
+                    break
+                prefix_length += 1
+
+            # Combine the common prefix with the remainder of the new tasks
+            final_tasks = self.tasks[:prefix_length] + new_tasks[prefix_length:]
+            self.tasks = final_tasks
+        
+        # Update current_task_id to the first unfinished task in the merged list
+        for task in self.tasks:
+            if not task.is_finished:
+                self.current_task_id = task.task_id
+                break
+
+        # Update the task map for quick access to tasks by ID
+        self.task_map = {task.task_id: task for task in self.tasks}
+
+    @property
+    def current_task(self) -> Task:
+        """Find current task to execute
+
+        Returns:
+            Task: the current task to be executed
+        """
+        return self.task_map.get(self.current_task_id, None)
+
+    def finish_current_task(self):
+        """Finish current task, set Task.is_finished=True, set current task to next task
+        """
+        if self.current_task_id:
+            current_task = self.current_task
+            current_task.is_finished = True
+            next_task_index = self.tasks.index(current_task) + 1
+            self.current_task_id = self.tasks[next_task_index].task_id if next_task_index < len(self.tasks) else None
+
+    def get_finished_tasks(self) -> list[Task]:
+        """return all finished tasks in correct linearized order
+
+        Returns:
+            list[Task]: list of finished tasks
+        """
+        return [task for task in self.tasks if task.is_finished]
+
+
 if __name__ == '__main__':
     test_content = 'test_message'
     msgs = [
diff --git a/requirements.txt b/requirements.txt
index 53176bd0a..c0f466457 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -45,4 +45,8 @@ semantic-kernel==0.3.13.dev0
 wrapt==1.15.0
 websocket-client==0.58.0
 zhipuai==1.0.7
-rich==13.6.0
\ No newline at end of file
+rich==13.6.0
+nbclient==0.9.0
+nbformat==5.9.2
+ipython==8.17.2
+ipykernel==6.27.0
\ No newline at end of file
diff --git a/tests/metagpt/test_schema.py b/tests/metagpt/test_schema.py
index 12666e0d3..6aae82006 100644
--- a/tests/metagpt/test_schema.py
+++ b/tests/metagpt/test_schema.py
@@ -6,6 +6,7 @@
 @File    : test_schema.py
 """
 from metagpt.schema import AIMessage, Message, SystemMessage, UserMessage
+from metagpt.schema import Task, Plan
 
 
 def test_messages():
@@ -19,3 +20,87 @@ def test_messages():
     text = str(msgs)
     roles = ['user', 'system', 'assistant', 'QA']
     assert all([i in text for i in roles])
+
+
+class TestPlan:
+    def test_add_tasks_ordering(self):
+        plan = Plan()
+
+        tasks = [
+            Task(task_id="1", dependent_task_ids=["2", "3"], instruction="Third"),
+            Task(task_id="2", instruction="First"),
+            Task(task_id="3", dependent_task_ids=["2"], instruction="Second")
+        ]  # 2 -> 3 -> 1
+        plan.add_tasks(tasks)
+
+        assert [task.task_id for task in plan.tasks] == ["2", "3", "1"]
+
+    def test_add_tasks_to_existing_no_common_prefix(self):
+        plan = Plan()
+
+        tasks = [
+            Task(task_id="1", dependent_task_ids=["2", "3"], instruction="Third"),
+            Task(task_id="2", instruction="First"),
+            Task(task_id="3", dependent_task_ids=["2"], instruction="Second", is_finished=True)
+        ]  # 2 -> 3 -> 1
+        plan.add_tasks(tasks)
+
+        new_tasks = [Task(task_id="3", instruction="")]
+        plan.add_tasks(new_tasks)
+
+        assert [task.task_id for task in plan.tasks] == ["3"]
+        assert not plan.tasks[0].is_finished  # must be the new unfinished task
+
+    def test_add_tasks_to_existing_with_common_prefix(self):
+        plan = Plan()
+
+        tasks = [
+            Task(task_id="1", dependent_task_ids=["2", "3"], instruction="Third"),
+            Task(task_id="2", instruction="First"),
+            Task(task_id="3", dependent_task_ids=["2"], instruction="Second")
+        ]  # 2 -> 3 -> 1
+        plan.add_tasks(tasks)
+        plan.finish_current_task()  # finish 2
+        plan.finish_current_task()  # finish 3
+
+        new_tasks = [
+            Task(task_id="4", dependent_task_ids=["3"], instruction="Third"),
+            Task(task_id="2", instruction="First"),
+            Task(task_id="3", dependent_task_ids=["2"], instruction="Second")
+        ]  # 2 -> 3 -> 4, so the common prefix is 2 -> 3, and these two should be obtained from the existing tasks
+        plan.add_tasks(new_tasks)
+
+        assert [task.task_id for task in plan.tasks] == ["2", "3", "4"]
+        assert plan.tasks[0].is_finished and plan.tasks[1].is_finished  # "2" and "3" should be the original finished one
+        assert plan.current_task_id == "4"
+
+    def test_current_task(self):
+        plan = Plan()
+        tasks = [
+            Task(task_id="1", dependent_task_ids=["2"], instruction="Second"),
+            Task(task_id="2", instruction="First")
+        ]
+        plan.add_tasks(tasks)
+        assert plan.current_task.task_id == "2"
+
+    def test_finish_task(self):
+        plan = Plan()
+        tasks = [
+            Task(task_id="1", instruction="First"),
+            Task(task_id="2", dependent_task_ids=["1"], instruction="Second")
+        ]
+        plan.add_tasks(tasks)
+        plan.finish_current_task()
+        assert plan.current_task.task_id == "2"
+
+    def test_finished_tasks(self):
+        plan = Plan()
+        tasks = [
+            Task(task_id="1", instruction="First"),
+            Task(task_id="2", dependent_task_ids=["1"], instruction="Second")
+        ]
+        plan.add_tasks(tasks)
+        plan.finish_current_task()
+        finished_tasks = plan.get_finished_tasks()
+        assert len(finished_tasks) == 1
+        assert finished_tasks[0].task_id == "1"

From fdef9c8547d743d41116d8bcf16fb3dd38b13e2d Mon Sep 17 00:00:00 2001
From: yzlin <yzlin@fuzhi.ai>
Date: Fri, 24 Nov 2023 14:05:11 +0800
Subject: [PATCH 11/11] add more components in pipeline

---
 metagpt/actions/__init__.py                   |  4 +-
 ...ode_function.py => write_analysis_code.py} | 22 ++++--
 metagpt/actions/write_plan.py                 |  6 +-
 metagpt/roles/ml_engineer.py                  | 70 +++++++++++++------
 metagpt/schema.py                             |  1 +
 ...unction.py => test_write_analysis_code.py} |  8 +--
 tests/metagpt/test_schema.py                  | 12 ++--
 7 files changed, 82 insertions(+), 41 deletions(-)
 rename metagpt/actions/{write_code_function.py => write_analysis_code.py} (76%)
 rename tests/metagpt/actions/{test_write_code_function.py => test_write_analysis_code.py} (86%)

diff --git a/metagpt/actions/__init__.py b/metagpt/actions/__init__.py
index ba2170cbd..5055ce276 100644
--- a/metagpt/actions/__init__.py
+++ b/metagpt/actions/__init__.py
@@ -24,7 +24,7 @@ from metagpt.actions.write_prd import WritePRD
 from metagpt.actions.write_prd_review import WritePRDReview
 from metagpt.actions.write_test import WriteTest
 from metagpt.actions.execute_code import ExecutePyCode
-from metagpt.actions.write_code_function import WriteCodeFunction
+from metagpt.actions.write_analysis_code import WriteCodeByGenerate
 from metagpt.actions.write_plan import WritePlan
 
 
@@ -49,7 +49,7 @@ class ActionType(Enum):
     WEB_BROWSE_AND_SUMMARIZE = WebBrowseAndSummarize
     CONDUCT_RESEARCH = ConductResearch
     EXECUTE_PYCODE = ExecutePyCode
-    WRITE_CODE_FUNCTION = WriteCodeFunction
+    WRITE_CODE_BY_GENERATE = WriteCodeByGenerate
     WRITE_PLAN = WritePlan
 
 
diff --git a/metagpt/actions/write_code_function.py b/metagpt/actions/write_analysis_code.py
similarity index 76%
rename from metagpt/actions/write_code_function.py
rename to metagpt/actions/write_analysis_code.py
index 4ec565eb1..84922ada4 100644
--- a/metagpt/actions/write_code_function.py
+++ b/metagpt/actions/write_analysis_code.py
@@ -11,17 +11,20 @@ from metagpt.schema import Message, Plan
 
 class BaseWriteAnalysisCode(Action):
 
-    async def run(self, context: List[Message], plan: Plan = None, task_guidance: str = ""):
+    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_guidance (str, optional): suggested step breakdown for the current task. Defaults to "".
+            task_guide (str, optional): suggested step breakdown for the current task. Defaults to "".
+        
+        Returns:
+            str: The code string.
         """
 
-class WriteCodeFunction(BaseWriteAnalysisCode):
-    """Use openai function to generate code."""
+class WriteCodeByGenerate(BaseWriteAnalysisCode):
+    """Write code fully by generation"""
 
     def __init__(self, name: str = "", context=None, llm=None) -> str:
         super().__init__(name, context, llm)
@@ -54,8 +57,15 @@ class WriteCodeFunction(BaseWriteAnalysisCode):
         return prompt
 
     async def run(
-        self, context: [List[Message]], plan: Plan = None, task_guidance: str = "", system_msg: str = None, **kwargs
+        self, context: [List[Message]], plan: Plan = None, task_guide: str = "", system_msg: str = None, **kwargs
     ) -> str:
         prompt = self.process_msg(context, system_msg)
         code_content = await self.llm.aask_code(prompt, **kwargs)
-        return code_content
+        return code_content["code"]
+
+
+class WriteCodeWithTools(BaseWriteAnalysisCode):
+    """Write code with help of local available tools. Choose tools first, then generate code to use the tools"""
+
+    async def run(self, context: List[Message], plan: Plan = None, task_guide: str = "") -> str:
+        return "print('abc')"
diff --git a/metagpt/actions/write_plan.py b/metagpt/actions/write_plan.py
index 48cb1aad5..e35ba7a92 100644
--- a/metagpt/actions/write_plan.py
+++ b/metagpt/actions/write_plan.py
@@ -9,6 +9,7 @@ import json
 
 from metagpt.actions import Action
 from metagpt.schema import Message, Task
+from metagpt.utils.common import CodeParser
 
 class WritePlan(Action):
     PROMPT_TEMPLATE = """
@@ -20,14 +21,16 @@ class WritePlan(Action):
     Based on the context, write a plan or modify an existing plan of what you should do to achieve the goal. A plan consists of one to __max_tasks__ tasks.
     If you are modifying an existing plan, carefully follow the instruction, don't make unnecessary changes.
     Output a list of jsons following the format:
+    ```json
     [
         {
-            "task_id": str = "unique identifier for a task in plan, can be a ordinal",
+            "task_id": str = "unique identifier for a task in plan, can be an ordinal",
             "dependent_task_ids": list[str] = "ids of tasks prerequisite to this task",
             "instruction": "what you should do in this task, one short phrase or sentence",
         },
         ...
     ]
+    ```
     """
     async def run(self, context: List[Message], current_plan: str = "", max_tasks: int = 5) -> str:
         prompt = (
@@ -35,6 +38,7 @@ class WritePlan(Action):
             .replace("__current_plan__", current_plan).replace("__max_tasks__", str(max_tasks))
         )
         rsp = await self._aask(prompt)
+        rsp = CodeParser.parse_code(block=None, text=rsp)
         return rsp
 
     @staticmethod
diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index c795bda11..480f6cecf 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -9,30 +9,41 @@ 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_code_function import WriteCodeFunction
+from metagpt.actions.write_analysis_code import WriteCodeByGenerate, WriteCodeWithTools
 from metagpt.actions.execute_code import ExecutePyCode
 
 class AskReview(Action):
 
     async def run(self, context: List[Message], plan: Plan = None):
-        prompt = "\n".join(
-            [f"{msg.cause_by() if msg.cause_by else 'Main Requirement'}: {msg.content}" for msg in context]
-        )
+        logger.info("Current overall plan:")
+        logger.info("\n".join([f"{task.task_id}: {task.instruction}" for task in plan.tasks]))
 
-        latest_action = context[-1].cause_by()
-
-        prompt += f"\nPlease review output from {latest_action}, " \
-            "provide feedback or type YES to continue with the process:\n"
+        logger.info("most recent context:")
+        # prompt = "\n".join(
+        #     [f"{msg.cause_by.__name__ if msg.cause_by else 'Main Requirement'}: {msg.content}" for msg in context]
+        # )
+        prompt = ""
+        latest_action = context[-1].cause_by.__name__
+        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"
         rsp = input(prompt)
-        confirmed = "yes" in rsp.lower()
+        confirmed = "confirm" in rsp.lower()
+
         return rsp, confirmed
 
+class WriteTaskGuide(Action):
+
+    async def run(self, task_instruction: str, data_desc: str = "") -> str:
+        return ""
 
 class MLEngineer(Role):
-    def __init__(self, name="ABC", profile="MLEngineer"):
-        super().__init__(name=name, profile=profile)
+    def __init__(self, name="ABC", profile="MLEngineer", goal=""):
+        super().__init__(name=name, profile=profile, goal=goal)
         self._set_react_mode(react_mode="plan_and_act")
-        self.plan = Plan()
+        self.plan = Plan(goal=goal)
+        self.use_tools = False
+        self.use_task_guide = False
 
     async def _plan_and_act(self):
 
@@ -60,18 +71,28 @@ class MLEngineer(Role):
                 await self._update_plan()
 
     async def _write_and_exec_code(self, max_retry: int = 3):
+
+        task_guide = await WriteTaskGuide().run(self.plan.current_task.instruction) if self.use_task_guide else ""
+
         counter = 0
         success = False
         while not success and counter < max_retry:
-            context = self.get_memories()
+            context = self.get_useful_memories()
 
-            code = "print('abc')"
-            # code = await WriteCodeFunction().run(context=context)
-            # code = await WriteCodeWithOps.run(context, task, result)
-            self._rc.memory.add(Message(content=code, role="assistant", cause_by=WriteCodeFunction))
+            if not self.use_tools:
+                # code = "print('abc')"
+                code = await WriteCodeByGenerate().run(context=context, plan=self.plan, task_guide=task_guide)
+                cause_by = WriteCodeByGenerate
+
+            else:
+                code = await WriteCodeWithTools().run(context=context, plan=self.plan, task_guide=task_guide)
+                cause_by = WriteCodeWithTools
+
+            self._rc.memory.add(Message(content=code, role="assistant", cause_by=cause_by))
 
             result, success = await ExecutePyCode().run(code)
-            self._rc.memory.add(Message(content=result, role="assistant", cause_by=ExecutePyCode))
+            print(result)
+            self._rc.memory.add(Message(content=result, role="user", cause_by=ExecutePyCode))
 
             # if not success:
             #     await self._ask_review()
@@ -81,16 +102,16 @@ class MLEngineer(Role):
         return code, result, success
 
     async def _ask_review(self):
-        context = self.get_memories()
+        context = self.get_useful_memories()
         review, confirmed = await AskReview().run(context=context[-5:], plan=self.plan)
-        self._rc.memory.add(Message(content=review, role="assistant", cause_by=AskReview))
+        self._rc.memory.add(Message(content=review, role="user", cause_by=AskReview))
         return confirmed
 
     async def _update_plan(self, max_tasks: int = 3):
         current_plan = str([task.json() for task in self.plan.tasks])
         plan_confirmed = False
         while not plan_confirmed:
-            context = self.get_memories()
+            context = self.get_useful_memories()
             rsp = await WritePlan().run(context, current_plan=current_plan, max_tasks=max_tasks)
             self._rc.memory.add(Message(content=rsp, role="assistant", cause_by=WritePlan))
             plan_confirmed = await self._ask_review()
@@ -98,13 +119,18 @@ class MLEngineer(Role):
         tasks = WritePlan.rsp_to_tasks(rsp)
         self.plan.add_tasks(tasks)
 
+    def get_useful_memories(self, current_task_memories: List[str] = []) -> List[Message]:
+        """find useful memories only to reduce context length and improve performance"""
+        memories = super().get_memories()
+        return memories
+
 
 if __name__ == "__main__":
     # requirement = "create a normal distribution and visualize it"
     requirement = "run some analysis on iris dataset"
 
     async def main(requirement: str = requirement):
-        role = MLEngineer()
+        role = MLEngineer(goal=requirement)
         await role.run(requirement)
 
     fire.Fire(main)
diff --git a/metagpt/schema.py b/metagpt/schema.py
index 3cd7d9730..e39f54a0c 100644
--- a/metagpt/schema.py
+++ b/metagpt/schema.py
@@ -84,6 +84,7 @@ class Task(BaseModel):
 
 
 class Plan(BaseModel):
+    goal: str
     tasks: list[Task] = []
     task_map: dict[str, Task] = {}
     current_task_id = ""
diff --git a/tests/metagpt/actions/test_write_code_function.py b/tests/metagpt/actions/test_write_analysis_code.py
similarity index 86%
rename from tests/metagpt/actions/test_write_code_function.py
rename to tests/metagpt/actions/test_write_analysis_code.py
index 4ff1a63c4..41c0479a9 100644
--- a/tests/metagpt/actions/test_write_code_function.py
+++ b/tests/metagpt/actions/test_write_analysis_code.py
@@ -1,12 +1,12 @@
 import pytest
 
-from metagpt.actions.write_code_function import WriteCodeFunction
+from metagpt.actions.write_analysis_code import WriteCodeByGenerate
 from metagpt.actions.execute_code import ExecutePyCode
 
 
 @pytest.mark.asyncio
 async def test_write_code():
-    write_code = WriteCodeFunction()
+    write_code = WriteCodeByGenerate()
     code = await write_code.run("Write a hello world code.")
     assert "language" in code.content
     assert "code" in code.content
@@ -15,7 +15,7 @@ async def test_write_code():
 
 @pytest.mark.asyncio
 async def test_write_code_by_list_prompt():
-    write_code = WriteCodeFunction()
+    write_code = WriteCodeByGenerate()
     msg = ["a=[1,2,5,10,-10]", "写出求a中最大值的代码python"]
     code = await write_code.run(msg)
     assert "language" in code.content
@@ -25,7 +25,7 @@ async def test_write_code_by_list_prompt():
 
 @pytest.mark.asyncio
 async def test_write_code_by_list_plan():
-    write_code = WriteCodeFunction()
+    write_code = WriteCodeByGenerate()
     execute_code = ExecutePyCode()
     messages = []
     plan = ["随机生成一个pandas DataFrame时间序列", "绘制这个时间序列的直方图", "求均值"]
diff --git a/tests/metagpt/test_schema.py b/tests/metagpt/test_schema.py
index 6aae82006..8f65d3785 100644
--- a/tests/metagpt/test_schema.py
+++ b/tests/metagpt/test_schema.py
@@ -24,7 +24,7 @@ def test_messages():
 
 class TestPlan:
     def test_add_tasks_ordering(self):
-        plan = Plan()
+        plan = Plan(goal="")
 
         tasks = [
             Task(task_id="1", dependent_task_ids=["2", "3"], instruction="Third"),
@@ -36,7 +36,7 @@ class TestPlan:
         assert [task.task_id for task in plan.tasks] == ["2", "3", "1"]
 
     def test_add_tasks_to_existing_no_common_prefix(self):
-        plan = Plan()
+        plan = Plan(goal="")
 
         tasks = [
             Task(task_id="1", dependent_task_ids=["2", "3"], instruction="Third"),
@@ -52,7 +52,7 @@ class TestPlan:
         assert not plan.tasks[0].is_finished  # must be the new unfinished task
 
     def test_add_tasks_to_existing_with_common_prefix(self):
-        plan = Plan()
+        plan = Plan(goal="")
 
         tasks = [
             Task(task_id="1", dependent_task_ids=["2", "3"], instruction="Third"),
@@ -75,7 +75,7 @@ class TestPlan:
         assert plan.current_task_id == "4"
 
     def test_current_task(self):
-        plan = Plan()
+        plan = Plan(goal="")
         tasks = [
             Task(task_id="1", dependent_task_ids=["2"], instruction="Second"),
             Task(task_id="2", instruction="First")
@@ -84,7 +84,7 @@ class TestPlan:
         assert plan.current_task.task_id == "2"
 
     def test_finish_task(self):
-        plan = Plan()
+        plan = Plan(goal="")
         tasks = [
             Task(task_id="1", instruction="First"),
             Task(task_id="2", dependent_task_ids=["1"], instruction="Second")
@@ -94,7 +94,7 @@ class TestPlan:
         assert plan.current_task.task_id == "2"
 
     def test_finished_tasks(self):
-        plan = Plan()
+        plan = Plan(goal="")
         tasks = [
             Task(task_id="1", instruction="First"),
             Task(task_id="2", dependent_task_ids=["1"], instruction="Second")