diff --git a/metagpt/actions/mi/execute_nb_code.py b/metagpt/actions/mi/execute_nb_code.py index 3a64a8bec..f6a8defbd 100644 --- a/metagpt/actions/mi/execute_nb_code.py +++ b/metagpt/actions/mi/execute_nb_code.py @@ -9,7 +9,6 @@ from __future__ import annotations import asyncio import base64 import re -import traceback from typing import Literal, Tuple import nbformat @@ -92,17 +91,17 @@ class ExecuteNbCode(Action): else: cell["outputs"].append(new_output(output_type="stream", name="stdout", text=str(output))) - def parse_outputs(self, outputs: list[str]) -> str: + def parse_outputs(self, outputs: list[str], keep_len: int = 2000) -> Tuple[bool, str]: """Parses the outputs received from notebook execution.""" assert isinstance(outputs, list) - parsed_output = "" - + parsed_output, is_success = [], True for i, output in enumerate(outputs): + output_text = "" if output["output_type"] == "stream" and not any( tag in output["text"] for tag in ["| INFO | metagpt", "| ERROR | metagpt", "| WARNING | metagpt", "DEBUG"] ): - parsed_output += output["text"] + output_text = output["text"] elif output["output_type"] == "display_data": if "image/png" in output["data"]: self.show_bytes_figure(output["data"]["image/png"], self.interaction) @@ -111,8 +110,22 @@ class ExecuteNbCode(Action): f"{i}th output['data'] from nbclient outputs dont have image/png, continue next output ..." ) elif output["output_type"] == "execute_result": - parsed_output += output["data"]["text/plain"] - return parsed_output + output_text = output["data"]["text/plain"] + elif output["output_type"] == "error": + output_text, is_success = "\n".join(output["traceback"]), False + + # handle coroutines that are not executed asynchronously + if output_text.strip().startswith(" Tuple[str, bool]: """ @@ -174,14 +187,7 @@ class ExecuteNbCode(Action): # run code cell_index = len(self.nb.cells) - 1 - success, error_message = await self.run_cell(self.nb.cells[-1], cell_index) - - if not success: - return truncate(remove_escape_and_color_codes(error_message), is_success=success) - - # code success - outputs = self.parse_outputs(self.nb.cells[-1].outputs) - outputs, success = truncate(remove_escape_and_color_codes(outputs), is_success=success) + success, outputs = await self.run_cell(self.nb.cells[-1], cell_index) if "!pip" in code: success = False @@ -197,24 +203,6 @@ class ExecuteNbCode(Action): raise ValueError(f"Only support for language: python, markdown, but got {language}, ") -def truncate(result: str, keep_len: int = 2000, is_success: bool = True): - """对于超出keep_len个字符的result: 执行失败的代码, 展示result后keep_len个字符; 执行成功的代码, 展示result前keep_len个字符。""" - if is_success: - desc = f"Executed code successfully. Truncated to show only first {keep_len} characters\n" - else: - desc = f"Executed code failed, please reflect the cause of bug and then debug. Truncated to show only last {keep_len} characters\n" - - if result.strip().startswith(" keep_len: - result = result[-keep_len:] if not is_success else result[:keep_len] - return desc + result, is_success - - return result, is_success - - def remove_escape_and_color_codes(input_str: str): # 使用正则表达式去除jupyter notebook输出结果中的转义字符和颜色代码 # Use regular expressions to get rid of escape characters and color codes in jupyter notebook output. diff --git a/tests/metagpt/actions/mi/test_execute_nb_code.py b/tests/metagpt/actions/mi/test_execute_nb_code.py index 3059ad3ae..2ecfbd2a2 100644 --- a/tests/metagpt/actions/mi/test_execute_nb_code.py +++ b/tests/metagpt/actions/mi/test_execute_nb_code.py @@ -1,6 +1,6 @@ import pytest -from metagpt.actions.mi.execute_nb_code import ExecuteNbCode, truncate +from metagpt.actions.mi.execute_nb_code import ExecuteNbCode @pytest.mark.asyncio @@ -54,21 +54,6 @@ async def test_plotting_code(): assert is_success -def test_truncate(): - # 代码执行成功 - output, is_success = truncate("hello world", 5, True) - assert "Truncated to show only first 5 characters\nhello" in output - assert is_success - # 代码执行失败 - output, is_success = truncate("hello world", 5, False) - assert "Truncated to show only last 5 characters\nworld" in output - assert not is_success - # 异步 - output, is_success = truncate("