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)