ag2 0.3.2__py3-none-any.whl

autogen/coding/jupyter/base.py

@@ -0,0 +1,32 @@
+# 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 dataclasses import dataclass
+from typing import Optional, Protocol, runtime_checkable
+
+
+@dataclass
+class JupyterConnectionInfo:
+    """(Experimental)"""
+
+    host: str
+    """`str` - Host of the Jupyter gateway server"""
+    use_https: bool
+    """`bool` - Whether to use HTTPS"""
+    port: Optional[int] = None
+    """`Optional[int]` - Port of the Jupyter gateway server. If None, the default port is used"""
+    token: Optional[str] = None
+    """`Optional[str]` - Token for authentication. If None, no token is used"""
+
+
+@runtime_checkable
+class JupyterConnectable(Protocol):
+    """(Experimental)"""
+
+    @property
+    def connection_info(self) -> JupyterConnectionInfo:
+        """Return the connection information for this connectable."""
+        pass

autogen/coding/jupyter/docker_jupyter_server.py

@@ -0,0 +1,164 @@
+# 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 io
+import logging
+import secrets
+import sys
+import uuid
+from pathlib import Path
+from types import TracebackType
+from typing import Dict, Optional, Type, Union
+
+import docker
+
+from ..docker_commandline_code_executor import _wait_for_ready
+
+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 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: Union[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 - 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 uuid
+from pathlib import Path
+from queue import Empty
+from typing import Any, ClassVar, List
+
+from jupyter_client import KernelManager  # type: ignore[attr-defined]
+from jupyter_client.kernelspec import KernelSpecManager
+from pydantic import BaseModel, Field, field_validator
+
+from ...agentchat.agent import LLMAgent
+from ..base import CodeBlock, CodeExtractor, IPythonCodeResult
+from ..markdown_code_extractor import MarkdownCodeExtractor
+
+__all__ = "EmbeddedIPythonCodeExecutor"
+
+
+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.
+
+    Args:
+        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 ".".
+    """
+
+    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:
+                if "-qqq" not in line:
+                    lines[i] = line.replace(match.group(0), match.group(0) + " -qqq")
+        return "\n".join(lines)

autogen/coding/jupyter/jupyter_client.py

@@ -0,0 +1,224 @@
+# 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 sys
+from dataclasses import dataclass
+from types import TracebackType
+from typing import Any, Dict, List, Optional, Type, cast
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+import datetime
+import json
+import uuid
+
+import requests
+import websocket
+from requests.adapters import HTTPAdapter, Retry
+from websocket import WebSocket
+
+from .base import JupyterConnectionInfo
+
+
+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()
+
+    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)
+
+
+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):
+        self._session_id: str = uuid.uuid4().hex
+        self._websocket: WebSocket = websocket
+
+    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
+        )