ag2 0.9.1__py3-none-any.whl → 0.9.1.post0__py3-none-any.whl

autogen/coding/jupyter/docker_jupyter_server.py

@@ -0,0 +1,167 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Portions derived from https://github.com/microsoft/autogen are under the MIT License.
+# SPDX-License-Identifier: MIT
+from __future__ import annotations
+
+import atexit
+import io
+import logging
+import secrets
+import sys
+import uuid
+from pathlib import Path
+from types import TracebackType
+from typing import Optional
+
+import docker
+
+from ...doc_utils import export_module
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+from ..docker_commandline_code_executor import _wait_for_ready
+from .base import JupyterConnectable, JupyterConnectionInfo
+from .import_utils import require_jupyter_kernel_gateway_installed
+from .jupyter_client import JupyterClient
+
+
+@require_jupyter_kernel_gateway_installed()
+@export_module("autogen.coding.jupyter")
+class DockerJupyterServer(JupyterConnectable):
+    DEFAULT_DOCKERFILE = """FROM quay.io/jupyter/docker-stacks-foundation
+
+SHELL ["/bin/bash", "-o", "pipefail", "-c"]
+
+USER ${NB_UID}
+RUN mamba install --yes jupyter_kernel_gateway ipykernel && \
+    mamba clean --all -f -y && \
+    fix-permissions "${CONDA_DIR}" && \
+    fix-permissions "/home/${NB_USER}"
+
+ENV TOKEN="UNSET"
+CMD python -m jupyter kernelgateway --KernelGatewayApp.ip=0.0.0.0 \
+    --KernelGatewayApp.port=8888 \
+    --KernelGatewayApp.auth_token="${TOKEN}" \
+    --JupyterApp.answer_yes=true \
+    --JupyterWebsocketPersonality.list_kernels=true
+
+EXPOSE 8888
+
+WORKDIR "${HOME}"
+"""
+
+    class GenerateToken:
+        pass
+
+    def __init__(
+        self,
+        *,
+        custom_image_name: Optional[str] = None,
+        container_name: Optional[str] = None,
+        auto_remove: bool = True,
+        stop_container: bool = True,
+        docker_env: dict[str, str] = {},
+        token: str | GenerateToken = GenerateToken(),
+    ):
+        """Start a Jupyter kernel gateway server in a Docker container.
+
+        Args:
+            custom_image_name (Optional[str], optional): Custom image to use. If this is None,
+                then the bundled image will be built and used. The default image is based on
+                quay.io/jupyter/docker-stacks-foundation and extended to include jupyter_kernel_gateway
+            container_name (Optional[str], optional): Name of the container to start.
+                A name will be generated if None.
+            auto_remove (bool, optional): If true the Docker container will be deleted
+                when it is stopped.
+            stop_container (bool, optional): If true the container will be stopped,
+                either by program exit or using the context manager
+            docker_env (Dict[str, str], optional): Extra environment variables to pass
+                to the running Docker container.
+            token (Union[str, GenerateToken], optional): Token to use for authentication.
+                If GenerateToken is used, a random token will be generated. Empty string
+                will be unauthenticated.
+        """
+        if container_name is None:
+            container_name = f"autogen-jupyterkernelgateway-{uuid.uuid4()}"
+
+        client = docker.from_env()
+        if custom_image_name is None:
+            image_name = "autogen-jupyterkernelgateway"
+            # Make sure the image exists
+            try:
+                client.images.get(image_name)
+            except docker.errors.ImageNotFound:
+                # Build the image
+                # Get this script directory
+                here = Path(__file__).parent
+                dockerfile = io.BytesIO(self.DEFAULT_DOCKERFILE.encode("utf-8"))
+                logging.info(f"Image {image_name} not found. Building it now.")
+                client.images.build(path=here, fileobj=dockerfile, tag=image_name)
+                logging.info(f"Image {image_name} built successfully.")
+        else:
+            image_name = custom_image_name
+            # Check if the image exists
+            try:
+                client.images.get(image_name)
+            except docker.errors.ImageNotFound:
+                raise ValueError(f"Custom image {image_name} does not exist")
+
+        if isinstance(token, DockerJupyterServer.GenerateToken):
+            self._token = secrets.token_hex(32)
+        else:
+            self._token = token
+
+        # Run the container
+        env = {"TOKEN": self._token}
+        env.update(docker_env)
+        container = client.containers.run(
+            image_name,
+            detach=True,
+            auto_remove=auto_remove,
+            environment=env,
+            publish_all_ports=True,
+            name=container_name,
+        )
+        _wait_for_ready(container)
+        container_ports = container.ports
+        self._port = int(container_ports["8888/tcp"][0]["HostPort"])
+        self._container_id = container.id
+
+        def cleanup() -> None:
+            try:
+                inner_container = client.containers.get(container.id)
+                inner_container.stop()
+            except docker.errors.NotFound:
+                pass
+
+            atexit.unregister(cleanup)
+
+        if stop_container:
+            atexit.register(cleanup)
+
+        self._cleanup_func = cleanup
+        self._stop_container = stop_container
+
+    @property
+    def connection_info(self) -> JupyterConnectionInfo:
+        return JupyterConnectionInfo(host="127.0.0.1", use_https=False, port=self._port, token=self._token)
+
+    def stop(self) -> None:
+        self._cleanup_func()
+
+    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()

autogen/coding/jupyter/embedded_ipython_code_executor.py

@@ -0,0 +1,182 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Portions derived from https://github.com/microsoft/autogen are under the MIT License.
+# SPDX-License-Identifier: MIT
+import base64
+import json
+import os
+import re
+import uuid
+from pathlib import Path
+from queue import Empty
+from typing import Any
+
+from pydantic import BaseModel, Field, field_validator
+
+from ...doc_utils import export_module
+from ...import_utils import optional_import_block, require_optional_import
+from ..base import CodeBlock, CodeExtractor, IPythonCodeResult
+from ..markdown_code_extractor import MarkdownCodeExtractor
+from .import_utils import require_jupyter_kernel_gateway_installed
+
+with optional_import_block():
+    from jupyter_client import KernelManager  # type: ignore[attr-defined]
+    from jupyter_client.kernelspec import KernelSpecManager
+
+__all__ = ["EmbeddedIPythonCodeExecutor"]
+
+
+@require_optional_import("jupyter_client", "jupyter-executor")
+@require_jupyter_kernel_gateway_installed()
+@export_module("autogen.coding.jupyter")
+class EmbeddedIPythonCodeExecutor(BaseModel):
+    """(Experimental) A code executor class that executes code statefully using an embedded
+    IPython kernel managed by this class.
+
+    **This will execute LLM generated code on the local machine.**
+
+    Each execution is stateful and can access variables created from previous
+    executions in the same session. The kernel must be installed before using
+    this class. The kernel can be installed using the following command:
+    `python -m ipykernel install --user --name {kernel_name}`
+    where `kernel_name` is the name of the kernel to install.
+    """
+
+    timeout: int = Field(default=60, ge=1, description="The timeout for code execution.")
+    kernel_name: str = Field(default="python3", description="The kernel name to use. Make sure it is installed.")
+    output_dir: str = Field(default=".", description="The directory to save output files.")
+
+    @field_validator("output_dir")
+    @classmethod
+    def _output_dir_must_exist(cls, value: str) -> str:
+        if not os.path.exists(value):
+            raise ValueError(f"Output directory {value} does not exist.")
+        return value
+
+    def __init__(self, **kwargs: Any):
+        super().__init__(**kwargs)
+        # Check if the kernel is installed.
+        if self.kernel_name not in KernelSpecManager().find_kernel_specs():
+            raise ValueError(
+                f"Kernel {self.kernel_name} is not installed. "
+                "Please first install it with "
+                f"`python -m ipykernel install --user --name {self.kernel_name}`."
+            )
+        self._kernel_manager = KernelManager(kernel_name=self.kernel_name)
+        self._kernel_manager.start_kernel()
+        self._kernel_client = self._kernel_manager.client()
+        self._kernel_client.start_channels()
+        self._timeout = self.timeout
+        self._kernel_name = self.kernel_name
+        self._output_dir = Path(self.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 an IPython kernel
+        managed by this class.
+        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._kernel_client.wait_for_ready()
+        outputs = []
+        output_files = []
+        for code_block in code_blocks:
+            code = self._process_code(code_block.code)
+            self._kernel_client.execute(code, store_history=True)
+            while True:
+                try:
+                    msg = self._kernel_client.get_iopub_msg(timeout=self._timeout)
+                    msg_type = msg["msg_type"]
+                    content = msg["content"]
+                    if msg_type in ["execute_result", "display_data"]:
+                        for data_type, data in content["data"].items():
+                            if data_type == "text/plain":
+                                # Output is a text.
+                                outputs.append(data)
+                            elif data_type.startswith("image/"):
+                                # Output is an image.
+                                path = self._save_image(data)
+                                outputs.append(f"Image data saved to {path}")
+                                output_files.append(path)
+                            elif data_type == "text/html":
+                                # Output is an html.
+                                path = self._save_html(data)
+                                outputs.append(f"HTML data saved to {path}")
+                                output_files.append(path)
+                            else:
+                                # Output raw data.
+                                outputs.append(json.dumps(data))
+                    elif msg_type == "stream":
+                        # Output is a text.
+                        outputs.append(content["text"])
+                    elif msg_type == "error":
+                        # Output is an error.
+                        return IPythonCodeResult(
+                            exit_code=1,
+                            output=f"ERROR: {content['ename']}: {content['evalue']}\n{content['traceback']}",
+                        )
+                    if msg_type == "status" and content["execution_state"] == "idle":
+                        break
+                # handle time outs.
+                except Empty:
+                    return IPythonCodeResult(
+                        exit_code=1,
+                        output=f"ERROR: Timeout waiting for output from code block: {code_block.code}",
+                    )
+        # We return the full 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._kernel_client.stop_channels()
+        self._kernel_manager.shutdown_kernel()
+        self._kernel_manager = KernelManager(kernel_name=self.kernel_name)
+        self._kernel_manager.start_kernel()
+        self._kernel_client = self._kernel_manager.client()
+        self._kernel_client.start_channels()
+
+    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 _process_code(self, code: str) -> str:
+        """Process code before execution."""
+        # Find lines that start with `! pip install` and make sure "-qqq" flag is added.
+        lines = code.split("\n")
+        for i, line in enumerate(lines):
+            # use regex to find lines that start with `! pip install` or `!pip install`.
+            match = re.search(r"^! ?pip install", line)
+            if match is not None and "-qqq" not in line:
+                lines[i] = line.replace(match.group(0), match.group(0) + " -qqq")
+        return "\n".join(lines)

autogen/coding/jupyter/import_utils.py

@@ -0,0 +1,82 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import subprocess
+from functools import lru_cache
+from logging import getLogger
+from typing import Callable, TypeVar
+
+from ...import_utils import patch_object
+
+logger = getLogger(__name__)
+
+__all__ = ["require_jupyter_kernel_gateway_installed", "skip_on_missing_jupyter_kernel_gateway"]
+
+
+@lru_cache
+def is_jupyter_kernel_gateway_installed() -> bool:
+    """Check if jupyter-kernel-gateway is installed."""
+    try:
+        subprocess.run(
+            ["jupyter", "kernelgateway", "--version"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            check=True,
+        )
+        return True
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        logger.warning(
+            "jupyter-kernel-gateway is required for JupyterCodeExecutor, please install it with `pip install ag2[jupyter-executor]`"
+        )
+        return False
+
+
+T = TypeVar("T")
+
+
+def require_jupyter_kernel_gateway_installed() -> Callable[[T], T]:
+    """Decorator that checks if jupyter-kernel-gateway is installed before function execution.
+
+    Returns:
+        Callable[[T], T]: A decorator function that either:
+            - Returns the original function unchanged if jupyter-kernel-gateway is installed
+            - Returns a patched version of the function that will raise a helpful error indicating the missing dependency when called
+    """
+    if is_jupyter_kernel_gateway_installed():
+
+        def decorator(o: T) -> T:
+            return o
+
+    else:
+
+        def decorator(o: T) -> T:
+            return patch_object(o, missing_modules={}, dep_target="jupyter-executor")
+
+    return decorator
+
+
+def skip_on_missing_jupyter_kernel_gateway() -> Callable[[T], T]:
+    """Decorator to skip a test if an optional module is missing"""
+    # Add pytest.mark.jupyter_executor decorator
+    mark_name = "jupyter_executor"
+
+    if is_jupyter_kernel_gateway_installed():
+
+        def decorator(o: T) -> T:
+            import pytest
+
+            pytest_mark_o = getattr(pytest.mark, mark_name)(o)
+            return pytest_mark_o  # type: ignore[no-any-return]
+
+    else:
+
+        def decorator(o: T) -> T:
+            import pytest
+
+            pytest_mark_o = getattr(pytest.mark, mark_name)(o)
+            return pytest.mark.skip(  # type: ignore[return-value,no-any-return]
+                reason="jupyter-kernel-gateway is required for JupyterCodeExecutor, please install it with `pip install ag2[jupyter-executor]`"
+            )(pytest_mark_o)
+
+    return decorator

autogen/coding/jupyter/jupyter_client.py

@@ -0,0 +1,231 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Portions derived from https://github.com/microsoft/autogen are under the MIT License.
+# SPDX-License-Identifier: MIT
+from __future__ import annotations
+
+import datetime
+import json
+import sys
+import uuid
+from dataclasses import dataclass
+from types import TracebackType
+from typing import Any, Optional, cast
+
+from ...doc_utils import export_module
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+
+import requests
+from requests.adapters import HTTPAdapter, Retry
+
+from ...import_utils import optional_import_block, require_optional_import
+from .base import JupyterConnectionInfo
+
+with optional_import_block():
+    import websocket
+    from websocket import WebSocket
+
+
+@export_module("autogen.coding.jupyter")
+class JupyterClient:
+    def __init__(self, connection_info: JupyterConnectionInfo):
+        """(Experimental) A client for communicating with a Jupyter gateway server.
+
+        Args:
+            connection_info (JupyterConnectionInfo): Connection information
+        """
+        self._connection_info = connection_info
+        self._session = requests.Session()
+        retries = Retry(total=5, backoff_factor=0.1)
+        self._session.mount("http://", HTTPAdapter(max_retries=retries))
+
+    def _get_headers(self) -> dict[str, str]:
+        if self._connection_info.token is None:
+            return {}
+        return {"Authorization": f"token {self._connection_info.token}"}
+
+    def _get_api_base_url(self) -> str:
+        protocol = "https" if self._connection_info.use_https else "http"
+        port = f":{self._connection_info.port}" if self._connection_info.port else ""
+        return f"{protocol}://{self._connection_info.host}{port}"
+
+    def _get_ws_base_url(self) -> str:
+        port = f":{self._connection_info.port}" if self._connection_info.port else ""
+        return f"ws://{self._connection_info.host}{port}"
+
+    def list_kernel_specs(self) -> dict[str, dict[str, str]]:
+        response = self._session.get(f"{self._get_api_base_url()}/api/kernelspecs", headers=self._get_headers())
+        return cast(dict[str, dict[str, str]], response.json())
+
+    def list_kernels(self) -> list[dict[str, str]]:
+        response = self._session.get(f"{self._get_api_base_url()}/api/kernels", headers=self._get_headers())
+        return cast(list[dict[str, str]], response.json())
+
+    def start_kernel(self, kernel_spec_name: str) -> str:
+        """Start a new kernel.
+
+        Args:
+            kernel_spec_name (str): Name of the kernel spec to start
+
+        Returns:
+            str: ID of the started kernel
+        """
+        response = self._session.post(
+            f"{self._get_api_base_url()}/api/kernels",
+            headers=self._get_headers(),
+            json={"name": kernel_spec_name},
+        )
+        return cast(str, response.json()["id"])
+
+    def delete_kernel(self, kernel_id: str) -> None:
+        response = self._session.delete(
+            f"{self._get_api_base_url()}/api/kernels/{kernel_id}", headers=self._get_headers()
+        )
+        response.raise_for_status()
+
+    def restart_kernel(self, kernel_id: str) -> None:
+        response = self._session.post(
+            f"{self._get_api_base_url()}/api/kernels/{kernel_id}/restart", headers=self._get_headers()
+        )
+        response.raise_for_status()
+
+    @require_optional_import("websocket", "jupyter-executor")
+    def get_kernel_client(self, kernel_id: str) -> JupyterKernelClient:
+        ws_url = f"{self._get_ws_base_url()}/api/kernels/{kernel_id}/channels"
+        ws = websocket.create_connection(ws_url, header=self._get_headers())
+        return JupyterKernelClient(ws)
+
+
+@require_optional_import("websocket", "jupyter-executor")
+class JupyterKernelClient:
+    """(Experimental) A client for communicating with a Jupyter kernel."""
+
+    @dataclass
+    class ExecutionResult:
+        @dataclass
+        class DataItem:
+            mime_type: str
+            data: str
+
+        is_ok: bool
+        output: str
+        data_items: list[DataItem]
+
+    def __init__(self, websocket: WebSocket):  # type: ignore[no-any-unimported]
+        self._session_id: str = uuid.uuid4().hex
+        self._websocket: WebSocket = websocket  # type: ignore[no-any-unimported]
+
+    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()
+
+    def stop(self) -> None:
+        self._websocket.close()
+
+    def _send_message(self, *, content: dict[str, Any], channel: str, message_type: str) -> str:
+        timestamp = datetime.datetime.now().isoformat()
+        message_id = uuid.uuid4().hex
+        message = {
+            "header": {
+                "username": "autogen",
+                "version": "5.0",
+                "session": self._session_id,
+                "msg_id": message_id,
+                "msg_type": message_type,
+                "date": timestamp,
+            },
+            "parent_header": {},
+            "channel": channel,
+            "content": content,
+            "metadata": {},
+            "buffers": {},
+        }
+        self._websocket.send_text(json.dumps(message))
+        return message_id
+
+    def _receive_message(self, timeout_seconds: Optional[float]) -> Optional[dict[str, Any]]:
+        self._websocket.settimeout(timeout_seconds)
+        try:
+            data = self._websocket.recv()
+            if isinstance(data, bytes):
+                data = data.decode("utf-8")
+            return cast(dict[str, Any], json.loads(data))
+        except websocket.WebSocketTimeoutException:
+            return None
+
+    def wait_for_ready(self, timeout_seconds: Optional[float] = None) -> bool:
+        message_id = self._send_message(content={}, channel="shell", message_type="kernel_info_request")
+        while True:
+            message = self._receive_message(timeout_seconds)
+            # This means we timed out with no new messages.
+            if message is None:
+                return False
+            if (
+                message.get("parent_header", {}).get("msg_id") == message_id
+                and message["msg_type"] == "kernel_info_reply"
+            ):
+                return True
+
+    def execute(self, code: str, timeout_seconds: Optional[float] = None) -> ExecutionResult:
+        message_id = self._send_message(
+            content={
+                "code": code,
+                "silent": False,
+                "store_history": True,
+                "user_expressions": {},
+                "allow_stdin": False,
+                "stop_on_error": True,
+            },
+            channel="shell",
+            message_type="execute_request",
+        )
+
+        text_output = []
+        data_output = []
+        while True:
+            message = self._receive_message(timeout_seconds)
+            if message is None:
+                return JupyterKernelClient.ExecutionResult(
+                    is_ok=False, output="ERROR: Timeout waiting for output from code block.", data_items=[]
+                )
+
+            # Ignore messages that are not for this execution.
+            if message.get("parent_header", {}).get("msg_id") != message_id:
+                continue
+
+            msg_type = message["msg_type"]
+            content = message["content"]
+            if msg_type in ["execute_result", "display_data"]:
+                for data_type, data in content["data"].items():
+                    if data_type == "text/plain":
+                        text_output.append(data)
+                    elif data_type.startswith("image/") or data_type == "text/html":
+                        data_output.append(self.ExecutionResult.DataItem(mime_type=data_type, data=data))
+                    else:
+                        text_output.append(json.dumps(data))
+            elif msg_type == "stream":
+                text_output.append(content["text"])
+            elif msg_type == "error":
+                # Output is an error.
+                return JupyterKernelClient.ExecutionResult(
+                    is_ok=False,
+                    output=f"ERROR: {content['ename']}: {content['evalue']}\n{content['traceback']}",
+                    data_items=[],
+                )
+            if msg_type == "status" and content["execution_state"] == "idle":
+                break
+
+        return JupyterKernelClient.ExecutionResult(
+            is_ok=True, output="\n".join([str(output) for output in text_output]), data_items=data_output
+        )