From 09e3a8e1fa98a6d03507f6f4db9ba2a0064d15fd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= Date: Mon, 11 Mar 2024 17:00:44 +0800 Subject: [PATCH] refine run_cell and parse_outputs. --- metagpt/actions/mi/execute_nb_code.py | 36 ++++++++----------- .../actions/mi/test_execute_nb_code.py | 5 +-- 2 files changed, 17 insertions(+), 24 deletions(-) diff --git a/metagpt/actions/mi/execute_nb_code.py b/metagpt/actions/mi/execute_nb_code.py index 632f0076c..217fc8ddc 100644 --- a/metagpt/actions/mi/execute_nb_code.py +++ b/metagpt/actions/mi/execute_nb_code.py @@ -96,11 +96,12 @@ class ExecuteNbCode(Action): assert isinstance(outputs, list) 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"] ): - ioutput, is_success = remove_escape_and_color_codes(output["text"]), True + 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) @@ -108,28 +109,22 @@ class ExecuteNbCode(Action): logger.info( f"{i}th output['data'] from nbclient outputs dont have image/png, continue next output ..." ) - ioutput, is_success = "", True elif output["output_type"] == "execute_result": - no_escape_color_output = remove_escape_and_color_codes(output["data"]["text/plain"]) - ioutput, is_success = no_escape_color_output, True + output_text = output["data"]["text/plain"] elif output["output_type"] == "error": - no_escape_color_output = remove_escape_and_color_codes("\n".join(output["traceback"])) - ioutput, is_success = no_escape_color_output, False + output_text, is_success = "\n".join(output["traceback"]), False # handle coroutines that are not executed asynchronously - if ioutput.strip().startswith(" keep_len and is_success: - prefix = f"Executed code successfully. Truncated to show only first {keep_len} characters\n" - ioutput = prefix + ioutput[:keep_len] - elif len(ioutput) > keep_len and not is_success: - prefix = f"Executed code failed, please reflect the cause of bug and then debug. Truncated to show only last {keep_len} characters\n" - ioutput = prefix + ioutput[-keep_len:] + output_text = remove_escape_and_color_codes(output_text) + # The valid information of the exception is at the end, + # the valid information of Normal output is at the begining. + output_text = output_text[:keep_len] if is_success else output_text[-keep_len:] - parsed_output.append(ioutput) + parsed_output.append(output_text) return is_success, ",".join(parsed_output) def show_bytes_figure(self, image_base64: str, interaction_type: Literal["ipython", None]): @@ -164,7 +159,7 @@ class ExecuteNbCode(Action): """ try: await self.nb_client.async_execute_cell(cell, cell_index) - return True, "" + return self.parse_outputs(self.nb.cells[-1].outputs) except CellTimeoutError: assert self.nb_client.km is not None await self.nb_client.km.interrupt_kernel() @@ -175,7 +170,7 @@ class ExecuteNbCode(Action): await self.reset() return False, "DeadKernelError" except Exception: - return False, "" + return self.parse_outputs(self.nb.cells[-1].outputs) async def run(self, code: str, language: Literal["python", "markdown"] = "python") -> Tuple[str, bool]: """ @@ -192,10 +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) - success, outputs = self.parse_outputs(self.nb.cells[-1].outputs) - if error_message: - outputs = error_message + outputs + success, outputs = await self.run_cell(self.nb.cells[-1], cell_index) if "!pip" in code: success = False diff --git a/tests/metagpt/actions/mi/test_execute_nb_code.py b/tests/metagpt/actions/mi/test_execute_nb_code.py index 4b90289ea..2ecfbd2a2 100644 --- a/tests/metagpt/actions/mi/test_execute_nb_code.py +++ b/tests/metagpt/actions/mi/test_execute_nb_code.py @@ -68,7 +68,7 @@ async def test_run_code_text(): executor = ExecuteNbCode() message, success = await executor.run(code='print("This is a code!")', language="python") assert success - assert message == "This is a code!\n" + assert "This is a code!" in message message, success = await executor.run(code="# This is a code!", language="markdown") assert success assert message == "# This is a code!" @@ -118,10 +118,11 @@ async def test_parse_outputs(): import pandas as pd df = pd.DataFrame({'ID': [1,2,3], 'NAME': ['a', 'b', 'c']}) print(df.columns) + print(f"columns num:{len(df.columns)}") print(df['DUMMPY_ID']) """ output, is_success = await executor.run(code) assert not is_success assert "Index(['ID', 'NAME'], dtype='object')" in output - assert "Executed code failed," in output assert "KeyError: 'DUMMPY_ID'" in output + assert "columns num:2" in output