feat: add actions: plan, write_code_function, pycode_executor.

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__ = [

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!")

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__)

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")

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
+...
+"""

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)])

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"

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"

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)