CoACT initialize (#292)

mm_agents/coact/autogen/coding/jupyter/local_jupyter_server.py

@@ -0,0 +1,172 @@
+# 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 json
+import secrets
+import signal
+import subprocess
+import sys
+from types import TracebackType
+from typing import Optional
+
+from ...doc_utils import export_module
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+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 LocalJupyterServer(JupyterConnectable):
+    class GenerateToken:
+        pass
+
+    def __init__(
+        self,
+        ip: str = "127.0.0.1",
+        port: Optional[int] = None,
+        token: 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,
+                capture_output=True,
+                text=True,
+            )
+        except subprocess.CalledProcessError:
+            raise ValueError(
+                "Jupyter gateway server is not installed. Please install it with `pip install jupyter_kernel_gateway`."
+            )
+
+        self.ip: str = ip
+
+        if isinstance(token, LocalJupyterServer.GenerateToken):
+            token = secrets.token_hex(32)
+
+        self.token: str = token
+        self._subprocess: subprocess.Popen[str]
+        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: int = 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()