Make Error Handling consistent with jupyter notebook.

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 _:

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,45 +72,55 @@ 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)
             await asyncio.wait_for(wait_for_ready_task, timeout=self._timeout)
 
             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."""

python/packages/autogen-ext/tests/code_executors/test_jupyter_code_executor.py

@@ -12,24 +12,63 @@
 @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"),
         ]
         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())