diff --git a/python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_client.py b/python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_client.py index 0832c3cd4ec9..b8ae4058c992 100644 --- a/python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_client.py @@ -176,7 +176,7 @@ async def execute(self, code: str) -> ExecutionResult: case "error": return JupyterKernelClient.ExecutionResult( is_ok=False, - output=f"ERROR: {content['ename']}: {content['evalue']}\n{content['traceback']}", + output="\n".join(content["traceback"]), data_items=[], ) case _: diff --git a/python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_code_executor.py b/python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_code_executor.py index 88a2b8d1ff0b..6078f9da214c 100644 --- a/python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_code_executor.py +++ b/python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_code_executor.py @@ -1,6 +1,7 @@ import asyncio import base64 import json +import re import sys import uuid from dataclasses import dataclass @@ -71,6 +72,9 @@ async def execute_code_blocks( asyncio.TimeoutError: Code execution timeouts asyncio.CancelledError: CancellationToken evoked during execution """ + if self._kernel_id is None: + raise ValueError("Kernel not running") + async with await self._jupyter_client.get_kernel_client(self._kernel_id) as kernel_client: wait_for_ready_task = asyncio.create_task(kernel_client.wait_for_ready()) cancellation_token.link_future(wait_for_ready_task) @@ -78,38 +82,45 @@ async def execute_code_blocks( outputs: list[str] = [] output_files: list[Path] = [] + exit_code = 0 + for code_block in code_blocks: code = silence_pip(code_block.code, code_block.language) execute_task = asyncio.create_task(kernel_client.execute(code)) cancellation_token.link_future(execute_task) result = await asyncio.wait_for(execute_task, timeout=self._timeout) - if result.is_ok: - outputs.append(result.output) - for data in result.data_items: - match data.mime_type: - case "image/png": - path = self._save_image(data.data) - output_files.append(path) - case "image/jpeg": - # TODO: Should this also be encoded? Images are encoded as both png and jpg - pass - case "text/html": - path = self._save_html(data.data) - output_files.append(path) - case _: - outputs.append(json.dumps(data.data)) - else: - return JupyterCodeResult(exit_code=1, output=f"ERROR: {result.output}", output_files=[]) - - return JupyterCodeResult( - exit_code=0, output="\n".join([output for output in outputs]), output_files=output_files - ) + # Clean ansi escape sequences + result.output = re.sub(r"\x1b\[[0-9;]*[A-Za-z]", "", result.output) + outputs.append(result.output) + + if not result.is_ok: + exit_code = 1 + break + + for data in result.data_items: + match data.mime_type: + case "image/png": + path = self._save_image(data.data) + output_files.append(path) + case "image/jpeg": + # TODO: Should this also be encoded? Images are encoded as both png and jpg + pass + case "text/html": + path = self._save_html(data.data) + output_files.append(path) + case _: + outputs.append(json.dumps(data.data)) + + return JupyterCodeResult(exit_code=exit_code, output="\n".join(outputs), output_files=output_files) async def restart(self) -> None: """Restart the code executor.""" - self._jupyter_client.restart_kernel(self._kernel_id) - self._jupyter_kernel_client = self._jupyter_client.get_kernel_client(self._kernel_id) + if self._kernel_id is None: + self.start() + else: + self._jupyter_client.restart_kernel(self._kernel_id) + self._jupyter_kernel_client = self._jupyter_client.get_kernel_client(self._kernel_id) def start(self) -> None: """Start the kernel.""" diff --git a/python/packages/autogen-ext/tests/code_executors/test_jupyter_code_executor.py b/python/packages/autogen-ext/tests/code_executors/test_jupyter_code_executor.py index 2fc53eeef9e5..43430cb878e2 100644 --- a/python/packages/autogen-ext/tests/code_executors/test_jupyter_code_executor.py +++ b/python/packages/autogen-ext/tests/code_executors/test_jupyter_code_executor.py @@ -12,12 +12,33 @@ @pytest.mark.asyncio async def test_execute_code(tmp_path: Path) -> None: with LocalJupyterServer() as server, JupyterCodeExecutor(server, output_dir=tmp_path) as executor: - # Test single code block. code_blocks = [CodeBlock(code="import sys; print('hello world!')", language="python")] code_result = await executor.execute_code_blocks(code_blocks, CancellationToken()) assert code_result == JupyterCodeResult(exit_code=0, output="hello world!\n", output_files=[]) - # Test multiple code blocks. + +@pytest.mark.asyncio +async def test_execute_code_error(tmp_path: Path) -> None: + with LocalJupyterServer() as server, JupyterCodeExecutor(server, output_dir=tmp_path) as executor: + code_blocks = [CodeBlock(code="print(undefined_variable)", language="python")] + code_result = await executor.execute_code_blocks(code_blocks, CancellationToken()) + assert code_result == JupyterCodeResult( + exit_code=1, + output=inspect.cleandoc(""" + --------------------------------------------------------------------------- + NameError Traceback (most recent call last) + Cell In[1], line 1 + ----> 1 print(undefined_variable) + + NameError: name 'undefined_variable' is not defined + """), + output_files=[], + ) + + +@pytest.mark.asyncio +async def test_execute_multiple_code_blocks(tmp_path: Path) -> None: + with LocalJupyterServer() as server, JupyterCodeExecutor(server, output_dir=tmp_path) as executor: code_blocks = [ CodeBlock(code="import sys; print('hello world!')", language="python"), CodeBlock(code="a = 100 + 100; print(a)", language="python"), @@ -25,11 +46,29 @@ async def test_execute_code(tmp_path: Path) -> None: code_result = await executor.execute_code_blocks(code_blocks, CancellationToken()) assert code_result == JupyterCodeResult(exit_code=0, output="hello world!\n\n200\n", output_files=[]) - # Test running code. - file_lines = ["import sys", "print('hello world!')", "a = 100 + 100", "print(a)"] - code_blocks = [CodeBlock(code="\n".join(file_lines), language="python")] + +@pytest.mark.asyncio +async def test_execute_multiple_code_blocks_error(tmp_path: Path) -> None: + with LocalJupyterServer() as server, JupyterCodeExecutor(server, output_dir=tmp_path) as executor: + code_blocks = [ + CodeBlock(code="import sys; print('hello world!')", language="python"), + CodeBlock(code="a = 100 + 100; print(a); print(undefined_variable)", language="python"), + ] code_result = await executor.execute_code_blocks(code_blocks, CancellationToken()) - assert code_result == JupyterCodeResult(exit_code=0, output="hello world!\n200\n", output_files=[]) + assert code_result == JupyterCodeResult( + exit_code=1, + output=inspect.cleandoc(""" + hello world! + + --------------------------------------------------------------------------- + NameError Traceback (most recent call last) + Cell In[2], line 1 + ----> 1 a = 100 + 100; print(a); print(undefined_variable) + + NameError: name 'undefined_variable' is not defined + """), + output_files=[], + ) @pytest.mark.asyncio @@ -48,7 +87,7 @@ async def test_execute_code_after_stop(tmp_path: Path) -> None: await asyncio.sleep(1) executor.stop() - with pytest.raises(websockets.exceptions.InvalidStatus): + with pytest.raises(ValueError): code_blocks = [CodeBlock(code="import sys; print('hello world!')", language="python")] await executor.execute_code_blocks(code_blocks, CancellationToken())