ag2 0.3.2b2__py3-none-any.whl

autogen/coding/jupyter/jupyter_code_executor.py

@@ -0,0 +1,161 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/autogen-ai
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Original portions of this file are derived from https://github.com/microsoft/autogen under the MIT License.
+# SPDX-License-Identifier: MIT
+import base64
+import json
+import os
+import re
+import sys
+import uuid
+from pathlib import Path
+from types import TracebackType
+from typing import Any, ClassVar, List, Optional, Type, Union
+
+from autogen.coding.utils import silence_pip
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+
+from ...agentchat.agent import LLMAgent
+from ..base import CodeBlock, CodeExecutor, CodeExtractor, IPythonCodeResult
+from ..markdown_code_extractor import MarkdownCodeExtractor
+from .base import JupyterConnectable, JupyterConnectionInfo
+from .jupyter_client import JupyterClient
+
+
+class JupyterCodeExecutor(CodeExecutor):
+    def __init__(
+        self,
+        jupyter_server: Union[JupyterConnectable, JupyterConnectionInfo],
+        kernel_name: str = "python3",
+        timeout: int = 60,
+        output_dir: Union[Path, str] = Path("."),
+    ):
+        """(Experimental) A code executor class that executes code statefully using
+        a Jupyter server supplied to this class.
+
+        Each execution is stateful and can access variables created from previous
+        executions in the same session.
+
+        Args:
+            jupyter_server (Union[JupyterConnectable, JupyterConnectionInfo]): The Jupyter server to use.
+            timeout (int): The timeout for code execution, by default 60.
+            kernel_name (str): The kernel name to use. Make sure it is installed.
+                By default, it is "python3".
+            output_dir (str): The directory to save output files, by default ".".
+        """
+        if timeout < 1:
+            raise ValueError("Timeout must be greater than or equal to 1.")
+
+        if isinstance(output_dir, str):
+            output_dir = Path(output_dir)
+
+        if not output_dir.exists():
+            raise ValueError(f"Output directory {output_dir} does not exist.")
+
+        if isinstance(jupyter_server, JupyterConnectable):
+            self._connection_info = jupyter_server.connection_info
+        elif isinstance(jupyter_server, JupyterConnectionInfo):
+            self._connection_info = jupyter_server
+        else:
+            raise ValueError("jupyter_server must be a JupyterConnectable or JupyterConnectionInfo.")
+
+        self._jupyter_client = JupyterClient(self._connection_info)
+        available_kernels = self._jupyter_client.list_kernel_specs()
+        if kernel_name not in available_kernels["kernelspecs"]:
+            raise ValueError(f"Kernel {kernel_name} is not installed.")
+
+        self._kernel_id = self._jupyter_client.start_kernel(kernel_name)
+        self._kernel_name = kernel_name
+        self._jupyter_kernel_client = self._jupyter_client.get_kernel_client(self._kernel_id)
+        self._timeout = timeout
+        self._output_dir = output_dir
+
+    @property
+    def code_extractor(self) -> CodeExtractor:
+        """(Experimental) Export a code extractor that can be used by an agent."""
+        return MarkdownCodeExtractor()
+
+    def execute_code_blocks(self, code_blocks: List[CodeBlock]) -> IPythonCodeResult:
+        """(Experimental) Execute a list of code blocks and return the result.
+
+        This method executes a list of code blocks as cells in the Jupyter kernel.
+        See: https://jupyter-client.readthedocs.io/en/stable/messaging.html
+        for the message protocol.
+
+        Args:
+            code_blocks (List[CodeBlock]): A list of code blocks to execute.
+
+        Returns:
+            IPythonCodeResult: The result of the code execution.
+        """
+        self._jupyter_kernel_client.wait_for_ready()
+        outputs = []
+        output_files = []
+        for code_block in code_blocks:
+            code = silence_pip(code_block.code, code_block.language)
+            result = self._jupyter_kernel_client.execute(code, timeout_seconds=self._timeout)
+            if result.is_ok:
+                outputs.append(result.output)
+                for data in result.data_items:
+                    if data.mime_type == "image/png":
+                        path = self._save_image(data.data)
+                        outputs.append(f"Image data saved to {path}")
+                        output_files.append(path)
+                    elif data.mime_type == "text/html":
+                        path = self._save_html(data.data)
+                        outputs.append(f"HTML data saved to {path}")
+                        output_files.append(path)
+                    else:
+                        outputs.append(json.dumps(data.data))
+            else:
+                return IPythonCodeResult(
+                    exit_code=1,
+                    output=f"ERROR: {result.output}",
+                )
+
+        return IPythonCodeResult(
+            exit_code=0, output="\n".join([str(output) for output in outputs]), output_files=output_files
+        )
+
+    def restart(self) -> None:
+        """(Experimental) Restart a new session."""
+        self._jupyter_client.restart_kernel(self._kernel_id)
+        self._jupyter_kernel_client = self._jupyter_client.get_kernel_client(self._kernel_id)
+
+    def _save_image(self, image_data_base64: str) -> str:
+        """Save image data to a file."""
+        image_data = base64.b64decode(image_data_base64)
+        # Randomly generate a filename.
+        filename = f"{uuid.uuid4().hex}.png"
+        path = os.path.join(self._output_dir, filename)
+        with open(path, "wb") as f:
+            f.write(image_data)
+        return os.path.abspath(path)
+
+    def _save_html(self, html_data: str) -> str:
+        """Save html data to a file."""
+        # Randomly generate a filename.
+        filename = f"{uuid.uuid4().hex}.html"
+        path = os.path.join(self._output_dir, filename)
+        with open(path, "w") as f:
+            f.write(html_data)
+        return os.path.abspath(path)
+
+    def stop(self) -> None:
+        """Stop the kernel."""
+        self._jupyter_client.delete_kernel(self._kernel_id)
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(
+        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+    ) -> None:
+        self.stop()

autogen/coding/jupyter/local_jupyter_server.py

@@ -0,0 +1,168 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/autogen-ai
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Original portions of this file are derived from https://github.com/microsoft/autogen under the MIT License.
+# SPDX-License-Identifier: MIT
+from __future__ import annotations
+
+import atexit
+import json
+import secrets
+import signal
+import socket
+import subprocess
+import sys
+from types import TracebackType
+from typing import Optional, Type, Union, cast
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+from .base import JupyterConnectable, JupyterConnectionInfo
+from .jupyter_client import JupyterClient
+
+
+class LocalJupyterServer(JupyterConnectable):
+    class GenerateToken:
+        pass
+
+    def __init__(
+        self,
+        ip: str = "127.0.0.1",
+        port: Optional[int] = None,
+        token: Union[str, GenerateToken] = GenerateToken(),
+        log_file: str = "jupyter_gateway.log",
+        log_level: str = "INFO",
+        log_max_bytes: int = 1048576,
+        log_backup_count: int = 3,
+    ):
+        """Runs a Jupyter Kernel Gateway server locally.
+
+        Args:
+            ip (str, optional): IP address to bind to. Defaults to "127.0.0.1".
+            port (Optional[int], optional): Port to use, if None it automatically selects a port. Defaults to None.
+            token (Union[str, GenerateToken], optional): Token to use for Jupyter server. By default will generate a token. Using None will use no token for authentication. Defaults to GenerateToken().
+            log_file (str, optional): File for Jupyter Kernel Gateway logs. Defaults to "jupyter_gateway.log".
+            log_level (str, optional): Level for Jupyter Kernel Gateway logs. Defaults to "INFO".
+            log_max_bytes (int, optional): Max logfile size. Defaults to 1048576.
+            log_backup_count (int, optional): Number of backups for rotating log. Defaults to 3.
+        """
+        # Remove as soon as https://github.com/jupyter-server/kernel_gateway/issues/398 is fixed
+        if sys.platform == "win32":
+            raise ValueError("LocalJupyterServer is not supported on Windows due to kernelgateway bug.")
+
+        # Check Jupyter gateway server is installed
+        try:
+            subprocess.run(
+                [sys.executable, "-m", "jupyter", "kernelgateway", "--version"],
+                check=True,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+            )
+        except subprocess.CalledProcessError:
+            raise ValueError(
+                "Jupyter gateway server is not installed. Please install it with `pip install jupyter_kernel_gateway`."
+            )
+
+        self.ip = ip
+
+        if isinstance(token, LocalJupyterServer.GenerateToken):
+            token = secrets.token_hex(32)
+
+        self.token = token
+        logging_config = {
+            "handlers": {
+                "file": {
+                    "class": "logging.handlers.RotatingFileHandler",
+                    "level": log_level,
+                    "maxBytes": log_max_bytes,
+                    "backupCount": log_backup_count,
+                    "filename": log_file,
+                }
+            },
+            "loggers": {"KernelGatewayApp": {"level": log_level, "handlers": ["file", "console"]}},
+        }
+
+        # Run Jupyter gateway server with detached subprocess
+        args = [
+            sys.executable,
+            "-m",
+            "jupyter",
+            "kernelgateway",
+            "--KernelGatewayApp.ip",
+            ip,
+            "--KernelGatewayApp.auth_token",
+            token,
+            "--JupyterApp.answer_yes",
+            "true",
+            "--JupyterApp.logging_config",
+            json.dumps(logging_config),
+            "--JupyterWebsocketPersonality.list_kernels",
+            "true",
+        ]
+        if port is not None:
+            args.extend(["--KernelGatewayApp.port", str(port)])
+            args.extend(["--KernelGatewayApp.port_retries", "0"])
+        self._subprocess = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+
+        # Satisfy mypy, we know this is not None because we passed PIPE
+        assert self._subprocess.stderr is not None
+        # Read stderr until we see "is available at" or the process has exited with an error
+        stderr = ""
+        while True:
+            result = self._subprocess.poll()
+            if result is not None:
+                stderr += self._subprocess.stderr.read()
+                raise ValueError(f"Jupyter gateway server failed to start with exit code: {result}. stderr:\n{stderr}")
+            line = self._subprocess.stderr.readline()
+            stderr += line
+
+            if "ERROR:" in line:
+                error_info = line.split("ERROR:")[1]
+                raise ValueError(f"Jupyter gateway server failed to start. {error_info}")
+
+            if "is available at" in line:
+                # We need to extract what port it settled on
+                # Example output:
+                #   Jupyter Kernel Gateway 3.0.0 is available at http://127.0.0.1:8890
+                if port is None:
+                    port = int(line.split(":")[-1])
+                self.port = port
+
+                break
+
+        # Poll the subprocess to check if it is still running
+        result = self._subprocess.poll()
+        if result is not None:
+            raise ValueError(
+                f"Jupyter gateway server failed to start. Please check the logs ({log_file}) for more information."
+            )
+
+        atexit.register(self.stop)
+
+    def stop(self) -> None:
+        if self._subprocess.poll() is None:
+            if sys.platform == "win32":
+                self._subprocess.send_signal(signal.CTRL_C_EVENT)
+            else:
+                self._subprocess.send_signal(signal.SIGINT)
+            self._subprocess.wait()
+
+    @property
+    def connection_info(self) -> JupyterConnectionInfo:
+        return JupyterConnectionInfo(host=self.ip, use_https=False, port=self.port, token=self.token)
+
+    def get_client(self) -> JupyterClient:
+        return JupyterClient(self.connection_info)
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(
+        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+    ) -> None:
+        self.stop()