Add tests for images and html

python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_client.py

@@ -167,7 +167,7 @@ async def execute(self, code: str) -> ExecutionResult:
                         match data_type:
                             case "text/plain":
                                 text_output.append(data)
-                            case data if data.startswith("image/") or data == "text/html":
+                            case type if type.startswith("image/") or type == "text/html":
                                 data_output.append(self.ExecutionResult.DataItem(mime_type=data_type, data=data))
                             case _:
                                 text_output.append(json.dumps(data))

python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_code_executor.py

@@ -90,11 +90,12 @@ async def execute_code_blocks(
                         match data.mime_type:
                             case "image/png":
                                 path = self._save_image(data.data)
-                                outputs.append(f"Image data saved to {path}")
                                 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)
-                                outputs.append(f"HTML data saved to {path}")
                                 output_files.append(path)
                             case _:
                                 outputs.append(json.dumps(data.data))

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

@@ -1,4 +1,5 @@
 import asyncio
+import inspect
 from pathlib import Path
 
 import pytest
@@ -75,7 +76,6 @@ async def test_commandline_code_executor_timeout(tmp_path: Path) -> None:
 @pytest.mark.asyncio
 async def test_commandline_code_executor_cancellation(tmp_path: Path) -> None:
     with LocalJupyterServer() as server, JupyterCodeExecutor(server, output_dir=tmp_path) as executor:
-        executor = JupyterCodeExecutor(server=server, output_dir=tmp_path)
         code_blocks = [CodeBlock(code="import time; time.sleep(10); print('hello world!')", language="python")]
 
         cancellation_token = CancellationToken()
@@ -86,3 +86,54 @@ async def test_commandline_code_executor_cancellation(tmp_path: Path) -> None:
 
         with pytest.raises(asyncio.CancelledError):
             await code_result_coroutine
+
+
+@pytest.mark.asyncio
+async def test_execute_code_with_image_output(tmp_path: Path) -> None:
+    with LocalJupyterServer() as server, JupyterCodeExecutor(server, output_dir=tmp_path) as executor:
+        code_blocks = [
+            CodeBlock(
+                code=inspect.cleandoc("""
+                    from PIL import Image, ImageDraw
+                    img = Image.new("RGB", (100, 100), color="white")
+                    draw = ImageDraw.Draw(img)
+                    draw.rectangle((10, 10, 90, 90), outline="black", fill="blue")
+                    display(img)
+                """),
+                language="python",
+            )
+        ]
+
+        code_result = await executor.execute_code_blocks(code_blocks, CancellationToken())
+
+        assert len(code_result.output_files) == 1
+        assert code_result == JupyterCodeResult(
+            exit_code=0,
+            output="<PIL.Image.Image image mode=RGB size=100x100>",
+            output_files=code_result.output_files,
+        )
+        assert code_result.output_files[0].parent == tmp_path
+
+
+@pytest.mark.asyncio
+async def test_execute_code_with_html_output(tmp_path: Path) -> None:
+    with LocalJupyterServer() as server, JupyterCodeExecutor(server, output_dir=tmp_path) as executor:
+        code_blocks = [
+            CodeBlock(
+                code=inspect.cleandoc("""
+                    from IPython.core.display import HTML
+                    HTML("<div style='color:blue'>Hello, HTML world!</div>")
+                """),
+                language="python",
+            )
+        ]
+
+        code_result = await executor.execute_code_blocks(code_blocks, CancellationToken())
+
+        assert len(code_result.output_files) == 1
+        assert code_result == JupyterCodeResult(
+            exit_code=0,
+            output="<IPython.core.display.HTML object>",
+            output_files=code_result.output_files,
+        )
+        assert code_result.output_files[0].parent == tmp_path