first commit

agent_dhal/agentdhal_extensions/code_executors/_common.py

@@ -0,0 +1,199 @@
+import inspect
+import re
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+from textwrap import dedent, indent
+from typing import Any, Callable, Optional, Sequence, Set, TypeVar, Union
+
+from agentdhal_core.code_executor import Alias, CodeResult, FunctionWithRequirements, FunctionWithRequirementsStr, Import
+from typing_extensions import ParamSpec
+
+
+@dataclass
+class CommandLineCodeResult(CodeResult):
+    """A code result class for command line code executor."""
+
+    code_file: Optional[str]
+
+
+T = TypeVar("T")
+P = ParamSpec("P")
+
+
+def _to_code(func: Union[FunctionWithRequirements[T, P], Callable[P, T], FunctionWithRequirementsStr]) -> str:
+    if isinstance(func, FunctionWithRequirementsStr):
+        return func.func
+
+    code = inspect.getsource(func)
+    # Strip the decorator
+    if code.startswith("@"):
+        code = code[code.index("\n") + 1 :]
+    return code
+
+
+def _import_to_str(im: Import) -> str:
+    if isinstance(im, str):
+        return f"import {im}"
+    elif isinstance(im, Alias):
+        return f"import {im.name} as {im.alias}"
+    else:
+
+        def to_str(i: Union[str, Alias]) -> str:
+            if isinstance(i, str):
+                return i
+            else:
+                return f"{i.name} as {i.alias}"
+
+        imports = ", ".join(map(to_str, im.imports))
+        return f"from {im.module} import {imports}"
+
+
+def build_python_functions_file(
+    funcs: Sequence[Union[FunctionWithRequirements[Any, P], Callable[..., Any], FunctionWithRequirementsStr]],
+) -> str:
+    """:meta private:"""
+    # First collect all global imports
+    global_imports: Set[Import] = set()
+    for func in funcs:
+        if isinstance(func, (FunctionWithRequirements, FunctionWithRequirementsStr)):
+            global_imports.update(func.global_imports)
+
+    content = "\n".join(map(_import_to_str, global_imports)) + "\n\n"
+
+    for func in funcs:
+        content += _to_code(func) + "\n\n"
+
+    return content
+
+
+def to_stub(func: Union[Callable[..., Any], FunctionWithRequirementsStr]) -> str:
+    """Generate a stub for a function as a string
+
+    Args:
+        func (Callable[..., Any]): The function to generate a stub for
+
+    Returns:
+        str: The stub for the function
+    """
+    if isinstance(func, FunctionWithRequirementsStr):
+        return to_stub(func.compiled_func)
+
+    content = f"def {func.__name__}{inspect.signature(func)}:\n"
+    docstring = func.__doc__
+
+    if docstring:
+        docstring = dedent(docstring)
+        docstring = '"""' + docstring + '"""'
+        docstring = indent(docstring, "    ")
+        content += docstring + "\n"
+
+    content += "    ..."
+    return content
+
+
+# Raises ValueError if the file is not in the workspace
+def get_file_name_from_content(code: str, workspace_path: Path) -> Optional[str]:
+    first_line = code.split("\n")[0]
+    # TODO - support other languages
+    if first_line.startswith("# filename:"):
+        filename = first_line.split(":")[1].strip()
+
+        # Handle relative paths in the filename
+        path = Path(filename)
+        if not path.is_absolute():
+            path = workspace_path / path
+        path = path.resolve()
+        # Throws an error if the file is not in the workspace
+        relative = path.relative_to(workspace_path.resolve())
+        return str(relative)
+
+    return None
+
+
+def silence_pip(code: str, lang: str) -> str:
+    """Apply -qqq flag to pip install commands."""
+    if lang == "python":
+        regex = r"^! ?pip install"
+    elif lang in ["bash", "shell", "sh", "pwsh", "powershell", "ps1"]:
+        regex = r"^pip install"
+    else:
+        return code
+
+    # 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.
+        match = re.search(regex, 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)
+
+
+def get_required_packages(code: str, lang: str) -> set[str]:
+    ret: set[str] = set()
+    if lang == "python":
+        regex = r"^! ?pip install(.*)$"
+    else:
+        return ret
+
+    # Find lines that start with pip install and make sure "-qqq" flag is added.
+    lines = code.split("\n")
+    for _, line in enumerate(lines):
+        # use regex to find lines that start with pip install.
+        match = re.search(regex, line)
+        if match is not None:
+            reqs = match.group(1).split(",")
+            ret = {req.strip(" ") for req in reqs}
+    return ret
+
+
+PYTHON_VARIANTS = ["python", "Python", "py"]
+
+
+def lang_to_cmd(lang: str) -> str:
+    if lang in PYTHON_VARIANTS:
+        return "python"
+    if lang.startswith("python") or lang in ["bash", "sh"]:
+        return lang
+    if lang in ["shell"]:
+        return "sh"
+    if lang in ["pwsh", "powershell", "ps1"]:
+        # Check if pwsh is available, otherwise fall back to powershell
+        if shutil.which("pwsh") is not None:
+            return "pwsh"
+        elif shutil.which("powershell") is not None:
+            return "powershell"
+        else:
+            raise ValueError("Powershell or pwsh is not installed. Please install one of them.")
+    else:
+        raise ValueError(f"Unsupported language: {lang}")
+
+
+# Regular expression for finding a code block
+# ```[ \t]*(\w+)?[ \t]*\r?\n(.*?)[ \t]*\r?\n``` Matches multi-line code blocks.
+#   The [ \t]* matches the potential spaces before language name.
+#   The (\w+)? matches the language, where the ? indicates it is optional.
+#   The [ \t]* matches the potential spaces (not newlines) after language name.
+#   The \r?\n makes sure there is a linebreak after ```.
+#   The (.*?) matches the code itself (non-greedy).
+#   The \r?\n makes sure there is a linebreak before ```.
+#   The [ \t]* matches the potential spaces before closing ``` (the spec allows indentation).
+CODE_BLOCK_PATTERN = r"```[ \t]*(\w+)?[ \t]*\r?\n(.*?)\r?\n[ \t]*```"
+
+
+def infer_lang(code: str) -> str:
+    """infer the language for the code.
+    TODO: make it robust.
+    """
+    if code.startswith("python ") or code.startswith("pip") or code.startswith("python3 "):
+        return "sh"
+
+    # check if code is a valid python code
+    try:
+        compile(code, "test", "exec")
+        return "python"
+    except SyntaxError:
+        # not a valid python code
+        return "unknown"

agent_dhal/agentdhal_extensions/code_executors/azure/__init__.py

@@ -0,0 +1,3 @@
+from ._azure_container_code_executor import ACADynamicSessionsCodeExecutor, TokenProvider
+
+__all__ = ["TokenProvider", "ACADynamicSessionsCodeExecutor"]

agent_dhal/agentdhal_extensions/code_executors/azure/_azure_container_code_executor.py

@@ -0,0 +1,522 @@
+# Credit to original authors
+
+from __future__ import annotations
+
+import asyncio
+import os
+import tempfile
+import warnings
+from pathlib import Path
+from string import Template
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, List, Optional, Protocol, Sequence, Union
+from uuid import uuid4
+
+import aiohttp
+
+# async functions shouldn't use open()
+from anyio import open_file
+from agentdhal_core import CancellationToken
+from agentdhal_core.code_executor import (
+    CodeBlock,
+    CodeExecutor,
+    CodeResult,
+    FunctionWithRequirements,
+    FunctionWithRequirementsStr,
+)
+from typing_extensions import ParamSpec
+
+from .._common import build_python_functions_file, get_required_packages, to_stub
+
+if TYPE_CHECKING:
+    from azure.core.credentials import AccessToken
+
+PYTHON_VARIANTS = ["python", "Python", "py"]
+
+__all__ = ("ACADynamicSessionsCodeExecutor", "TokenProvider")
+
+A = ParamSpec("A")
+
+
+class TokenProvider(Protocol):
+    def get_token(
+        self, *scopes: str, claims: Optional[str] = None, tenant_id: Optional[str] = None, **kwargs: Any
+    ) -> AccessToken: ...
+
+
+class ACADynamicSessionsCodeExecutor(CodeExecutor):
+    """(Experimental) A code executor class that executes code through a an Azure
+    Container Apps Dynamic Sessions instance.
+
+    .. note::
+
+        This class requires the :code:`azure` extra for the :code:`autogen-ext` package:
+
+        .. code-block:: bash
+
+            pip install "agentdhal-ext[azure]"
+
+    .. caution::
+
+        **This will execute LLM generated code on an Azure dynamic code container.**
+
+    The execution environment is similar to that of a jupyter notebook which allows for incremental code execution. The parameter functions are executed in order once at the beginning of each session. Each code block is then executed serially and in the order they are received. Each environment has a statically defined set of available packages which cannot be changed.
+    Currently, attempting to use packages beyond what is available on the environment will result in an error. To get the list of supported packages, call the `get_available_packages` function.
+    Currently the only supported language is Python.
+    For Python code, use the language "python" for the code block.
+
+    Args:
+        pool_management_endpoint (str): The azure container apps dynamic sessions endpoint.
+        credential (TokenProvider): An object that implements the get_token function.
+        timeout (int): The timeout for the execution of any single code block. Default is 60.
+        work_dir (str): The working directory for the code execution. If None,
+            a default working directory will be used. The default working
+            directory is a temporal directory.
+        functions (List[Union[FunctionWithRequirements[Any, A], Callable[..., Any]]]): A list of functions that are available to the code executor. Default is an empty list.
+        suppress_result_output bool: By default the executor will attach any result info in the execution response to the result outpu. Set this to True to prevent this.
+        session_id (str): The session id for the code execution (passed to Dynamic Sessions). If None, a new session id will be generated. Default is None. Note this value will be reset when calling `restart`
+
+    .. note::
+        Using the current directory (".") as working directory is deprecated. Using it will raise a deprecation warning.
+    """
+
+    SUPPORTED_LANGUAGES: ClassVar[List[str]] = [
+        "python",
+    ]
+    FUNCTION_PROMPT_TEMPLATE: ClassVar[str] = """You have access to the following user defined functions.
+
+$functions"""
+
+    _AZURE_API_VER = "2024-02-02-preview"
+
+    def __init__(
+        self,
+        pool_management_endpoint: str,
+        credential: TokenProvider,
+        timeout: int = 60,
+        work_dir: Union[Path, str, None] = None,
+        functions: Sequence[
+            Union[
+                FunctionWithRequirements[Any, A],
+                Callable[..., Any],
+                FunctionWithRequirementsStr,
+            ]
+        ] = [],
+        functions_module: str = "functions",
+        suppress_result_output: bool = False,
+        session_id: Optional[str] = None,
+    ):
+        if timeout < 1:
+            raise ValueError("Timeout must be greater than or equal to 1.")
+
+        self._work_dir: Optional[Path] = None
+        self._temp_dir: Optional[tempfile.TemporaryDirectory[str]] = None
+
+        # If a user specifies a working directory, use that
+        if work_dir is not None:
+            if isinstance(work_dir, str):
+                self._work_dir = Path(work_dir)
+            else:
+                self._work_dir = work_dir
+            # Create the directory if it doesn't exist
+            self._work_dir.mkdir(exist_ok=True, parents=True)
+        # If a user does not specify a working directory, use the default directory (tempfile.TemporaryDirectory)
+        else:
+            self._temp_dir = tempfile.TemporaryDirectory()
+            temp_dir_path = Path(self._temp_dir.name)
+            temp_dir_path.mkdir(exist_ok=True, parents=True)
+
+        self._started = False
+
+        # Rest of initialization remains the same
+        self._functions_module = functions_module
+        self._timeout = timeout
+        self._functions = functions
+        self._func_code: Optional[str] = None
+
+        # Setup could take some time so we intentionally wait for the first code block to do it.
+        if len(functions) > 0:
+            self._setup_functions_complete = False
+        else:
+            self._setup_functions_complete = True
+
+        self._suppress_result_output = suppress_result_output
+
+        self._pool_management_endpoint = pool_management_endpoint
+        self._access_token: str | None = None
+        self._session_id: str = session_id or str(uuid4())
+        self._available_packages: set[str] | None = None
+        self._credential: TokenProvider = credential
+        # cwd needs to be set to /mnt/data to properly read uploaded files and download written files
+        self._setup_cwd_complete = False
+
+    # TODO: expiration?
+    def _ensure_access_token(self) -> None:
+        if not self._access_token:
+            scope = "https://dynamicsessions.io/.default"
+            self._access_token = self._credential.get_token(scope).token
+
+    def format_functions_for_prompt(self, prompt_template: str = FUNCTION_PROMPT_TEMPLATE) -> str:
+        """(Experimental) Format the functions for a prompt.
+
+        The template includes one variable:
+        - `$functions`: The functions formatted as stubs with two newlines between each function.
+
+        Args:
+            prompt_template (str): The prompt template. Default is the class default.
+
+        Returns:
+            str: The formatted prompt.
+        """
+
+        template = Template(prompt_template)
+        return template.substitute(
+            functions="\n\n".join([to_stub(func) for func in self._functions]),
+        )
+
+    @property
+    def functions_module(self) -> str:
+        """(Experimental) The module name for the functions."""
+        return self._functions_module
+
+    @property
+    def functions(self) -> List[str]:
+        raise NotImplementedError
+
+    @property
+    def timeout(self) -> int:
+        """(Experimental) The timeout for code execution."""
+        return self._timeout
+
+    @property
+    def work_dir(self) -> Path:
+        # If a user specifies a working directory, use that
+        if self._work_dir is not None:
+            # If a user specifies the current directory, warn them that this is deprecated
+            if self._work_dir == Path("."):
+                warnings.warn(
+                    "Using the current directory as work_dir is deprecated",
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
+            return self._work_dir
+        # If a user does not specify a working directory, use the default directory (tempfile.TemporaryDirectory)
+        elif self._temp_dir is not None:
+            return Path(self._temp_dir.name)
+        else:
+            raise RuntimeError("Working directory not properly initialized")
+
+    def _construct_url(self, path: str) -> str:
+        endpoint = self._pool_management_endpoint
+        if not endpoint.endswith("/"):
+            endpoint += "/"
+        url = endpoint + f"{path}?api-version={self._AZURE_API_VER}&identifier={self._session_id}"
+        return url
+
+    async def get_available_packages(self, cancellation_token: CancellationToken) -> set[str]:
+        if self._available_packages is not None:
+            return self._available_packages
+        avail_pkgs = """
+import pkg_resources\n[d.project_name for d in pkg_resources.working_set]
+"""
+        ret = await self._execute_code_dont_check_setup(
+            [CodeBlock(code=avail_pkgs, language="python")], cancellation_token
+        )
+        if ret.exit_code != 0:
+            raise ValueError(f"Failed to get list of available packages: {ret.output.strip()}")
+        pkgs = ret.output.strip("[]")
+        pkglist = pkgs.split(",\n")
+        return {pkg.strip(" '") for pkg in pkglist}
+
+    async def _populate_available_packages(self, cancellation_token: CancellationToken) -> None:
+        self._available_packages = await self.get_available_packages(cancellation_token)
+
+    async def _setup_functions(self, cancellation_token: CancellationToken) -> None:
+        if not self._func_code:
+            self._func_code = build_python_functions_file(self._functions)
+
+            # Check required function imports and packages
+            lists_of_packages = [x.python_packages for x in self._functions if isinstance(x, FunctionWithRequirements)]
+            # Should we also be checking the imports?
+
+            flattened_packages = [item for sublist in lists_of_packages for item in sublist]
+            required_packages = set(flattened_packages)
+
+            if self._available_packages is None:
+                await self._populate_available_packages(cancellation_token)
+
+            if self._available_packages is not None:
+                missing_pkgs = set(required_packages - self._available_packages)
+                if len(missing_pkgs) > 0:
+                    raise ValueError(f"Packages unavailable in environment: {missing_pkgs}")
+
+        func_file = self.work_dir / f"{self._functions_module}.py"
+        func_file.write_text(self._func_code)
+
+        # Attempt to load the function file to check for syntax errors, imports etc.
+        exec_result = await self._execute_code_dont_check_setup(
+            [CodeBlock(code=self._func_code, language="python")], cancellation_token
+        )
+
+        if exec_result.exit_code != 0:
+            raise ValueError(f"Functions failed to load: {exec_result.output.strip()}")
+
+        self._setup_functions_complete = True
+
+    async def _setup_cwd(self, cancellation_token: CancellationToken) -> None:
+        # Change the cwd to /mnt/data to properly have access to uploaded files
+        exec_result = await self._execute_code_dont_check_setup(
+            [CodeBlock(code="import os; os.chdir('/mnt/data')", language="python")], cancellation_token
+        )
+
+        if exec_result.exit_code != 0:
+            raise ValueError("Failed to set up Azure container working directory")
+        self._setup_cwd_complete = True
+
+    async def get_file_list(self, cancellation_token: CancellationToken) -> List[str]:
+        self._ensure_access_token()
+        timeout = aiohttp.ClientTimeout(total=float(self._timeout))
+        headers = {
+            "Authorization": f"Bearer {self._access_token}",
+        }
+        url = self._construct_url("files")
+        async with aiohttp.ClientSession(timeout=timeout) as client:
+            task = asyncio.create_task(
+                client.get(
+                    url,
+                    headers=headers,
+                )
+            )
+            cancellation_token.link_future(task)
+            try:
+                resp = await task
+                resp.raise_for_status()
+                data = await resp.json()
+            except asyncio.TimeoutError as e:
+                # e.add_note is only in py 3.11+
+                raise asyncio.TimeoutError("Timeout getting file list") from e
+            except asyncio.CancelledError as e:
+                # e.add_note is only in py 3.11+
+                raise asyncio.CancelledError("File list retrieval cancelled") from e
+            except aiohttp.ClientResponseError as e:
+                raise ConnectionError("Error while getting file list") from e
+
+        values = data["value"]
+        file_info_list: List[str] = []
+        for value in values:
+            file = value["properties"]
+            file_info_list.append(file["filename"])
+        return file_info_list
+
+    async def upload_files(self, files: List[Union[Path, str]], cancellation_token: CancellationToken) -> None:
+        self._ensure_access_token()
+        # TODO: Better to use the client auth system rather than headers
+        headers = {"Authorization": f"Bearer {self._access_token}"}
+        url = self._construct_url("files/upload")
+        timeout = aiohttp.ClientTimeout(total=float(self._timeout))
+        async with aiohttp.ClientSession(timeout=timeout) as client:
+            for file in files:
+                file_path = self.work_dir / file
+                if not file_path.is_file():
+                    # TODO: what to do here?
+                    raise FileNotFoundError(f"{file} does not exist")
+
+                data = aiohttp.FormData()
+                async with await open_file(file_path, "rb") as f:
+                    data.add_field(
+                        "file",
+                        f,
+                        filename=os.path.basename(file_path),
+                        content_type="application/octet-stream",
+                    )
+
+                    task = asyncio.create_task(
+                        client.post(
+                            url,
+                            headers=headers,
+                            data=data,
+                        )
+                    )
+
+                    cancellation_token.link_future(task)
+                    try:
+                        resp = await task
+                        resp.raise_for_status()
+
+                    except asyncio.TimeoutError as e:
+                        # e.add_note is only in py 3.11+
+                        raise asyncio.TimeoutError("Timeout uploading files") from e
+                    except asyncio.CancelledError as e:
+                        # e.add_note is only in py 3.11+
+                        raise asyncio.CancelledError("Uploading files cancelled") from e
+                    except aiohttp.ClientResponseError as e:
+                        raise ConnectionError("Error while uploading files") from e
+
+    async def download_files(self, files: List[Union[Path, str]], cancellation_token: CancellationToken) -> List[str]:
+        self._ensure_access_token()
+        available_files = await self.get_file_list(cancellation_token)
+        # TODO: Better to use the client auth system rather than headers
+        headers = {"Authorization": f"Bearer {self._access_token}"}
+        timeout = aiohttp.ClientTimeout(total=float(self._timeout))
+        local_paths: List[str] = []
+        async with aiohttp.ClientSession(timeout=timeout) as client:
+            for file in files:
+                if file not in available_files:
+                    # TODO: what's the right thing to do here?
+                    raise FileNotFoundError(f"{file} does not exist")
+
+                url = self._construct_url(f"files/content/{file}")
+
+                task = asyncio.create_task(
+                    client.get(
+                        url,
+                        headers=headers,
+                    )
+                )
+                cancellation_token.link_future(task)
+                try:
+                    resp = await task
+                    resp.raise_for_status()
+                    local_path = self.work_dir / file
+                    local_paths.append(str(local_path))
+                    async with await open_file(local_path, "wb") as f:
+                        await f.write(await resp.read())
+                except asyncio.TimeoutError as e:
+                    # e.add_note is only in py 3.11+
+                    raise asyncio.TimeoutError("Timeout downloading files") from e
+                except asyncio.CancelledError as e:
+                    # e.add_note is only in py 3.11+
+                    raise asyncio.CancelledError("Downloading files cancelled") from e
+                except aiohttp.ClientResponseError as e:
+                    raise ConnectionError("Error while downloading files") from e
+        return local_paths
+
+    async def execute_code_blocks(
+        self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken
+    ) -> CodeResult:
+        """(Experimental) Execute the code blocks and return the result.
+
+        Args:
+            code_blocks (List[CodeBlock]): The code blocks to execute.
+            cancellation_token (CancellationToken): a token to cancel the operation
+            input_files (Optional[Union[Path, str]]): Any files the code blocks will need to access
+
+        Returns:
+            CodeResult: The result of the code execution."""
+
+        self._ensure_access_token()
+        if self._available_packages is None:
+            await self._populate_available_packages(cancellation_token)
+        if not self._setup_functions_complete:
+            await self._setup_functions(cancellation_token)
+        if not self._setup_cwd_complete:
+            await self._setup_cwd(cancellation_token)
+
+        return await self._execute_code_dont_check_setup(code_blocks, cancellation_token)
+
+    # The http call here should be replaced by an actual Azure client call once its available
+    async def _execute_code_dont_check_setup(
+        self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken
+    ) -> CodeResult:
+        logs_all = ""
+        exitcode = 0
+
+        # TODO: Better to use the client auth system rather than headers
+        assert self._access_token is not None
+        headers = {
+            "Authorization": f"Bearer {self._access_token}",
+            "Content-Type": "application/json",
+        }
+        properties = {
+            "codeInputType": "inline",
+            "executionType": "synchronous",
+            "code": "",  # Filled in later
+        }
+        url = self._construct_url("code/execute")
+        timeout = aiohttp.ClientTimeout(total=float(self._timeout))
+        async with aiohttp.ClientSession(timeout=timeout) as client:
+            for code_block in code_blocks:
+                lang, code = code_block.language, code_block.code
+                lang = lang.lower()
+
+                if lang in PYTHON_VARIANTS:
+                    lang = "python"
+
+                if lang not in self.SUPPORTED_LANGUAGES:
+                    # In case the language is not supported, we return an error message.
+                    exitcode = 1
+                    logs_all += "\n" + f"unknown language {lang}"
+                    break
+
+                if self._available_packages is not None:
+                    req_pkgs = get_required_packages(code, lang)
+                    missing_pkgs = set(req_pkgs - self._available_packages)
+                    if len(missing_pkgs) > 0:
+                        # In case the code requires packages that are not available in the environment
+                        exitcode = 1
+                        logs_all += "\n" + f"Python packages unavailable in environment: {missing_pkgs}"
+                        break
+
+                properties["code"] = code_block.code
+
+                task = asyncio.create_task(
+                    client.post(
+                        url,
+                        headers=headers,
+                        json={"properties": properties},
+                    )
+                )
+
+                cancellation_token.link_future(task)
+                try:
+                    response = await task
+                    response.raise_for_status()
+                    data = await response.json()
+                    data = data["properties"]
+                    logs_all += data.get("stderr", "") + data.get("stdout", "")
+                    if "Success" in data["status"]:
+                        if not self._suppress_result_output:
+                            logs_all += str(data["result"])
+                    elif "Failure" in data["status"]:
+                        exitcode = 1
+
+                except asyncio.TimeoutError as e:
+                    logs_all += "\n Timeout"
+                    # e.add_note is only in py 3.11+
+                    raise asyncio.TimeoutError(logs_all) from e
+                except asyncio.CancelledError as e:
+                    logs_all += "\n Cancelled"
+                    # e.add_note is only in py 3.11+
+                    raise asyncio.CancelledError(logs_all) from e
+                except aiohttp.ClientResponseError as e:
+                    logs_all += "\nError while sending code block to endpoint"
+                    raise ConnectionError(logs_all) from e
+
+        return CodeResult(exit_code=exitcode, output=logs_all)
+
+    async def restart(self) -> None:
+        """(Experimental) Restart the code executor.
+
+        Resets the internal state of the executor by generating a new session ID and resetting the setup variables.
+        This causes the next code execution to reinitialize the environment and re-run any setup code.
+        """
+        self._session_id = str(uuid4())
+        self._setup_functions_complete = False
+        self._access_token = None
+        self._available_packages = None
+        self._setup_cwd_complete = False
+
+    async def start(self) -> None:
+        """(Experimental) Start the code executor.
+
+        Marks the code executor as started."""
+        # No setup needed for this executor
+        self._started = True
+
+    async def stop(self) -> None:
+        """(Experimental) Stop the code executor.
+
+        Stops the code executor after cleaning up the temporary working directory (if it was created)."""
+        if self._temp_dir is not None:
+            self._temp_dir.cleanup()
+            self._temp_dir = None
+        self._started = False

agent_dhal/agentdhal_extensions/code_executors/docker/__init__.py

@@ -0,0 +1,3 @@
+from ._docker_code_executor import DockerCommandLineCodeExecutor
+
+__all__ = ["DockerCommandLineCodeExecutor"]

agent_dhal/agentdhal_extensions/code_executors/docker/_docker_code_executor.py

@@ -0,0 +1,613 @@
+# File based from: https://github.com/microsoft/autogen/blob/main/autogen/coding/docker_commandline_code_executor.py
+# Credit to original authors
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import shlex
+import sys
+import tempfile
+import uuid
+import warnings
+from collections.abc import Sequence
+from concurrent.futures import Future as ConcurrentFuture
+from hashlib import sha256
+from pathlib import Path
+from typing import Any, Callable, ClassVar, Dict, List, Optional, ParamSpec, Tuple, Union
+
+from agentdhal_core import CancellationToken, Component
+from agentdhal_core.code_executor import (
+    CodeBlock,
+    CodeExecutor,
+    FunctionWithRequirements,
+    FunctionWithRequirementsStr,
+)
+from pydantic import BaseModel
+from typing_extensions import Self
+
+from docker.types import DeviceRequest
+
+from .._common import (
+    CommandLineCodeResult,
+    build_python_functions_file,
+    get_file_name_from_content,
+    lang_to_cmd,
+    silence_pip,
+)
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+try:
+    import asyncio_atexit
+
+    import docker
+    from docker.errors import DockerException, ImageNotFound, NotFound
+    from docker.models.containers import Container
+except ImportError as e:
+    raise RuntimeError(
+        "Missing dependecies for DockerCommandLineCodeExecutor. Please ensure the autogen-ext package was installed with the 'docker' extra."
+    ) from e
+
+
+async def _wait_for_ready(container: Any, timeout: int = 60, stop_time: float = 0.1) -> None:
+    elapsed_time = 0.0
+    while container.status != "running" and elapsed_time < timeout:
+        await asyncio.sleep(stop_time)
+        elapsed_time += stop_time
+        await asyncio.to_thread(container.reload)
+        continue
+    if container.status != "running":
+        raise ValueError("Container failed to start")
+
+
+A = ParamSpec("A")
+
+
+class DockerCommandLineCodeExecutorConfig(BaseModel):
+    """Configuration for DockerCommandLineCodeExecutor"""
+
+    image: str = "python:3-slim"
+    container_name: Optional[str] = None
+    timeout: int = 60
+    work_dir: Optional[str] = None
+    bind_dir: Optional[str] = None
+    auto_remove: bool = True
+    stop_container: bool = True
+    functions_module: str = "functions"
+    extra_volumes: Dict[str, Dict[str, str]] = {}
+    extra_hosts: Dict[str, str] = {}
+    init_command: Optional[str] = None
+    delete_tmp_files: bool = False
+
+
+class DockerCommandLineCodeExecutor(CodeExecutor, Component[DockerCommandLineCodeExecutorConfig]):
+    """Executes code through a command line environment in a Docker container.
+
+    .. note::
+
+        This class requires the :code:`docker` extra for the :code:`autogen-ext` package:
+
+        .. code-block:: bash
+
+            pip install "agentdhal-ext[docker]"
+
+
+    The executor first saves each code block in a file in the working
+    directory, and then executes the code file in the container.
+    The executor executes the code blocks in the order they are received.
+    Currently, the executor only supports Python and shell scripts.
+    For Python code, use the language "python" for the code block.
+    For shell scripts, use the language "bash", "shell", "sh", "pwsh", "powershell", or "ps1" for the code block.
+
+    Args:
+        image (_type_, optional): Docker image to use for code execution.
+            Defaults to "python:3-slim".
+        container_name (Optional[str], optional): Name of the Docker container
+            which is created. If None, will autogenerate a name. Defaults to None.
+        timeout (int, optional): The timeout for code execution. Defaults to 60.
+        work_dir (Union[Path, str], optional): The working directory for the code
+            execution. Defaults to temporary directory.
+        bind_dir (Union[Path, str], optional): The directory that will be bound
+        to the code executor container. Useful for cases where you want to spawn
+        the container from within a container. Defaults to work_dir.
+        auto_remove (bool, optional): If true, will automatically remove the Docker
+            container when it is stopped. Defaults to True.
+        stop_container (bool, optional): If true, will automatically stop the
+            container when stop is called, when the context manager exits or when
+            the Python process exits with atext. Defaults to True.
+        device_requests (Optional[List[DeviceRequest]], optional): A list of device request instances to add to the container for exposing GPUs (e.g., [docker.types.DeviceRequest(count=-1, capabilities=[['gpu']])]). Defaults to None.
+        functions (List[Union[FunctionWithRequirements[Any, A], Callable[..., Any]]]): A list of functions that are available to the code executor. Default is an empty list.
+        functions_module (str, optional): The name of the module that will be created to store the functions. Defaults to "functions".
+        extra_volumes (Optional[Dict[str, Dict[str, str]]], optional): A dictionary of extra volumes (beyond the work_dir) to mount to the container;
+            key is host source path and value 'bind' is the container path. See  Defaults to None.
+            Example: extra_volumes = {'/home/user1/': {'bind': '/mnt/vol2', 'mode': 'rw'}, '/var/www': {'bind': '/mnt/vol1', 'mode': 'ro'}}
+        extra_hosts (Optional[Dict[str, str]], optional): A dictionary of host mappings to add to the container. (See Docker docs on extra_hosts) Defaults to None.
+            Example: extra_hosts = {"kubernetes.docker.internal": "host-gateway"}
+        init_command (Optional[str], optional): A shell command to run before each shell operation execution. Defaults to None.
+            Example: init_command="kubectl config use-context docker-hub"
+        delete_tmp_files (bool, optional): If true, will delete temporary files after execution. Defaults to False.
+
+    .. note::
+        Using the current directory (".") as working directory is deprecated. Using it will raise a deprecation warning.
+
+    """
+
+    component_config_schema = DockerCommandLineCodeExecutorConfig
+    component_provider_override = "agentdhal_extensions.code_executors.docker.DockerCommandLineCodeExecutor"
+
+    SUPPORTED_LANGUAGES: ClassVar[List[str]] = [
+        "bash",
+        "shell",
+        "sh",
+        "pwsh",
+        "powershell",
+        "ps1",
+        "python",
+    ]
+
+    FUNCTION_PROMPT_TEMPLATE: ClassVar[
+        str
+    ] = """You have access to the following user defined functions. They can be accessed from the module called `$module_name` by their function names.
+
+For example, if there was a function called `foo` you could import it by writing `from $module_name import foo`
+
+$functions"""
+
+    def __init__(
+        self,
+        image: str = "python:3-slim",
+        container_name: Optional[str] = None,
+        *,
+        timeout: int = 60,
+        work_dir: Union[Path, str, None] = None,
+        bind_dir: Optional[Union[Path, str]] = None,
+        auto_remove: bool = True,
+        stop_container: bool = True,
+        device_requests: Optional[List[DeviceRequest]] = None,
+        functions: Sequence[
+            Union[
+                FunctionWithRequirements[Any, A],
+                Callable[..., Any],
+                FunctionWithRequirementsStr,
+            ]
+        ] = [],
+        functions_module: str = "functions",
+        extra_volumes: Optional[Dict[str, Dict[str, str]]] = None,
+        extra_hosts: Optional[Dict[str, str]] = None,
+        init_command: Optional[str] = None,
+        delete_tmp_files: bool = False,
+    ):
+        if timeout < 1:
+            raise ValueError("Timeout must be greater than or equal to 1.")
+
+        # Handle working directory logic
+        if work_dir is None:
+            self._work_dir = None
+        else:
+            if isinstance(work_dir, str):
+                work_dir = Path(work_dir)
+            # Emit a deprecation warning if the user is using the current directory as working directory
+            if work_dir.resolve() == Path.cwd().resolve():
+                warnings.warn(
+                    "Using the current directory as work_dir is deprecated.",
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
+            self._work_dir = work_dir
+            # Create the working directory if it doesn't exist
+            self._work_dir.mkdir(exist_ok=True, parents=True)
+
+        if container_name is None:
+            self.container_name = f"agentdhal-code-exec-{uuid.uuid4()}"
+        else:
+            self.container_name = container_name
+
+        self._timeout = timeout
+
+        # Handle bind_dir
+        self._bind_dir: Optional[Path] = None
+        if bind_dir is not None:
+            self._bind_dir = Path(bind_dir) if isinstance(bind_dir, str) else bind_dir
+        else:
+            self._bind_dir = self._work_dir  # Default to work_dir if not provided
+
+        # Track temporary directory
+        self._temp_dir: Optional[tempfile.TemporaryDirectory[str]] = None
+        self._temp_dir_path: Optional[Path] = None
+
+        self._started = False
+
+        self._auto_remove = auto_remove
+        self._stop_container = stop_container
+        self._image = image
+
+        if not functions_module.isidentifier():
+            raise ValueError("Module name must be a valid Python identifier")
+
+        self._functions_module = functions_module
+        self._functions = functions
+        self._extra_volumes = extra_volumes if extra_volumes is not None else {}
+        self._extra_hosts = extra_hosts if extra_hosts is not None else {}
+        self._init_command = init_command
+        self._delete_tmp_files = delete_tmp_files
+        self._device_requests = device_requests
+
+        # Setup could take some time so we intentionally wait for the first code block to do it.
+        if len(functions) > 0:
+            self._setup_functions_complete = False
+        else:
+            self._setup_functions_complete = True
+
+        self._container: Container | None = None
+        self._running = False
+
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+        self._cancellation_futures: List[ConcurrentFuture[None]] = []
+
+    @property
+    def timeout(self) -> int:
+        """(Experimental) The timeout for code execution."""
+        return self._timeout
+
+    async def _setup_functions(self, cancellation_token: CancellationToken) -> None:
+        func_file_content = build_python_functions_file(self._functions)
+        func_file = self.work_dir / f"{self._functions_module}.py"
+        func_file.write_text(func_file_content)
+
+        # Collect requirements
+        lists_of_packages = [x.python_packages for x in self._functions if isinstance(x, FunctionWithRequirements)]
+        flattened_packages = [item for sublist in lists_of_packages for item in sublist]
+        required_packages = list(set(flattened_packages))
+        if len(required_packages) > 0:
+            logging.info("Ensuring packages are installed in executor.")
+
+            packages = shlex.join(required_packages)
+
+            result = await self._execute_code_dont_check_setup(
+                [CodeBlock(code=f"python -m pip install {packages}", language="sh")], cancellation_token
+            )
+
+            if result.exit_code != 0:
+                stdout = result.output
+                stderr = result.output
+                raise ValueError(f"Pip install failed. {stdout}, {stderr}")
+
+        # Attempt to load the function file to check for syntax errors, imports etc.
+        exec_result = await self._execute_code_dont_check_setup(
+            [CodeBlock(code=func_file_content, language="python")], cancellation_token
+        )
+
+        if exec_result.exit_code != 0:
+            raise ValueError(f"Functions failed to load: {exec_result.output}")
+
+        self._setup_functions_complete = True
+
+    async def _kill_running_command(self, command: List[str]) -> None:
+        if self._container is None or not self._running:
+            return
+        await asyncio.to_thread(self._container.exec_run, ["pkill", "-f", " ".join(command)])
+
+    async def _execute_command(self, command: List[str], cancellation_token: CancellationToken) -> Tuple[str, int]:
+        if self._container is None or not self._running:
+            raise ValueError("Container is not running. Must first be started with either start or a context manager.")
+
+        exec_task = asyncio.create_task(asyncio.to_thread(self._container.exec_run, command))
+        cancellation_token.link_future(exec_task)
+
+        # Wait for the exec task to finish.
+        try:
+            result = await exec_task
+            exit_code = result.exit_code
+            output = result.output.decode("utf-8")
+            if exit_code == 124:
+                output += "\n Timeout"
+            return output, exit_code
+        except asyncio.CancelledError:
+            # Schedule a task to kill the running command in the background.
+            if self._loop and not self._loop.is_closed():
+                try:
+                    logging.debug(f"Scheduling kill command via run_coroutine_threadsafe on loop {self._loop!r}")
+                    future: ConcurrentFuture[None] = asyncio.run_coroutine_threadsafe(
+                        self._kill_running_command(command), self._loop
+                    )
+                    self._cancellation_futures.append(future)
+                    logging.debug(f"Kill command scheduled, future: {future!r}")
+                except RuntimeError as e:
+                    logging.error(f"Failed to schedule kill command on loop {self._loop!r}: {e}")
+                except Exception as e:
+                    logging.exception(f"Unexpected error scheduling kill command: {e}")
+            else:
+                logging.warning(
+                    f"Cannot schedule kill command: Executor loop is not available or closed (loop: {self._loop!r})."
+                )
+            return "Code execution was cancelled.", 1
+
+    async def _execute_code_dont_check_setup(
+        self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken
+    ) -> CommandLineCodeResult:
+        if self._container is None or not self._running:
+            raise ValueError("Container is not running. Must first be started with either start or a context manager.")
+
+        if len(code_blocks) == 0:
+            raise ValueError("No code blocks to execute.")
+
+        outputs: List[str] = []
+        files: List[Path] = []
+        last_exit_code = 0
+        try:
+            for code_block in code_blocks:
+                lang = code_block.language.lower()
+                code = silence_pip(code_block.code, lang)
+
+                # Check if there is a filename comment
+                try:
+                    filename = get_file_name_from_content(code, self.work_dir)
+                except ValueError:
+                    outputs.append("Filename is not in the workspace")
+                    last_exit_code = 1
+                    break
+
+                if not filename:
+                    filename = f"tmp_code_{sha256(code.encode()).hexdigest()}.{lang}"
+
+                code_path = self.work_dir / filename
+                with code_path.open("w", encoding="utf-8") as fout:
+                    fout.write(code)
+                files.append(code_path)
+
+                command = ["timeout", str(self._timeout), lang_to_cmd(lang), filename]
+
+                output, exit_code = await self._execute_command(command, cancellation_token)
+                outputs.append(output)
+                last_exit_code = exit_code
+                if exit_code != 0:
+                    break
+        finally:
+            if self._delete_tmp_files:
+                for file in files:
+                    try:
+                        file.unlink()
+                    except (OSError, FileNotFoundError):
+                        pass
+
+        code_file = str(files[0]) if files else None
+        return CommandLineCodeResult(exit_code=last_exit_code, output="".join(outputs), code_file=code_file)
+
+    @property
+    def work_dir(self) -> Path:
+        # If a user specifies a working directory, use that
+        if self._work_dir is not None:
+            # If a user specifies the current directory, warn them that this is deprecated
+            if self._work_dir == Path("."):
+                warnings.warn(
+                    "Using the current directory as work_dir is deprecated.",
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
+            return self._work_dir
+        # If a user does not specify a working directory, use the default directory (tempfile.TemporaryDirectory)
+        elif self._temp_dir is not None:
+            return Path(self._temp_dir.name)
+        else:
+            raise RuntimeError("Working directory not properly initialized")
+
+    @property
+    def bind_dir(self) -> Path:
+        # If the user specified a bind directory, return it
+        if self._bind_dir is not None:
+            return self._bind_dir
+        # Otherwise bind_dir is set to the current work_dir as default
+        else:
+            return self.work_dir
+
+    async def execute_code_blocks(
+        self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken
+    ) -> CommandLineCodeResult:
+        """(Experimental) Execute the code blocks and return the result.
+
+        Args:
+            code_blocks (List[CodeBlock]): The code blocks to execute.
+
+        Returns:
+            CommandlineCodeResult: The result of the code execution."""
+
+        if not self._setup_functions_complete:
+            await self._setup_functions(cancellation_token)
+
+        return await self._execute_code_dont_check_setup(code_blocks, cancellation_token)
+
+    async def restart(self) -> None:
+        """(Experimental) Restart the Docker container code executor."""
+        if self._container is None or not self._running:
+            raise ValueError("Container is not running. Must first be started with either start or a context manager.")
+
+        await asyncio.to_thread(self._container.restart)  # type: ignore
+        if self._container.status != "running":
+            self._running = False
+            logs_str = self._container.logs().decode("utf-8")
+            raise ValueError(f"Failed to restart container. Logs: {logs_str}")
+
+    async def stop(self) -> None:
+        """(Experimental) Stop the code executor.
+
+        Stops the Docker container and cleans up any temporary files (if they were created), along with the temporary directory.
+        The method first waits for all cancellation tasks to finish before stopping the container. Finally it marks the executor as not running.
+        If the container is not running, the method does nothing.
+        """
+        if not self._running:
+            return
+
+        if self._temp_dir is not None:
+            self._temp_dir.cleanup()
+            self._temp_dir = None
+
+        client = docker.from_env()
+        try:
+            try:
+                container = await asyncio.to_thread(client.containers.get, self.container_name)
+            except NotFound:
+                logging.debug(f"Container {self.container_name} not found during stop...")
+                self._running = False
+                self._cancellation_futures.clear()
+                return
+
+            if self._cancellation_futures:
+                if not self._loop or self._loop.is_closed():
+                    logging.warning(
+                        f"Executor loop ({self._loop!r}) is closed or unavailable. Cannot reliably wait for "
+                        f"{len(self._cancellation_futures)} cancellation futures."
+                    )
+                    self._cancellation_futures.clear()
+                else:
+                    # concurrent.futures.Future -> asyncio.Future
+                    asyncio_futures = [asyncio.wrap_future(f, loop=self._loop) for f in self._cancellation_futures]
+
+                    if asyncio_futures:
+                        logging.debug(
+                            f"Waiting for {len(asyncio_futures)} cancellation futures to complete on loop {self._loop!r}..."
+                        )
+                        results = await asyncio.gather(*asyncio_futures, return_exceptions=True)
+                        for i, result in enumerate(results):
+                            original_future = self._cancellation_futures[i]
+                            if isinstance(result, Exception):
+                                logging.warning(f"Cancellation future {original_future!r} failed: {result}")
+                            else:
+                                logging.debug(f"Cancellation future {original_future!r} completed successfully.")
+                    else:
+                        logging.debug("No valid cancellation futures to await.")
+
+                    self._cancellation_futures.clear()
+
+            logging.debug(f"Stopping container {self.container_name}...")
+            await asyncio.to_thread(container.stop)
+            logging.debug(f"Container {self.container_name} stopped.")
+
+        except DockerException as e:
+            logging.error(f"Docker error while stopping container {self.container_name}: {e}")
+        except Exception as e:
+            logging.exception(f"Unexpected error during stop operation for container {self.container_name}: {e}")
+        finally:
+            self._running = False
+            self._cancellation_futures.clear()
+
+    async def start(self) -> None:
+        """(Experimental) Start the code executor.
+
+        This method sets the working environment variables, connects to Docker and starts the code executor.
+        If no working directory was provided to the code executor, it creates a temporary directory and sets it as the code executor working directory.
+        """
+
+        if self._work_dir is None and self._temp_dir is None:
+            self._temp_dir = tempfile.TemporaryDirectory()
+            self._temp_dir_path = Path(self._temp_dir.name)
+            self._temp_dir_path.mkdir(exist_ok=True)
+
+        # Start a container from the image, read to exec commands later
+        try:
+            client = docker.from_env()
+        except DockerException as e:
+            if "FileNotFoundError" in str(e):
+                raise RuntimeError("Failed to connect to Docker. Please ensure Docker is installed and running.") from e
+            raise
+        except Exception as e:
+            raise RuntimeError(f"Unexpected error while connecting to Docker: {str(e)}") from e
+
+        # Check if the image exists
+        try:
+            await asyncio.to_thread(client.images.get, self._image)
+        except ImageNotFound:
+            # TODO logger
+            logging.info(f"Pulling image {self._image}...")
+            # Let the docker exception escape if this fails.
+            await asyncio.to_thread(client.images.pull, self._image)
+
+        # Prepare the command (if needed)
+        shell_command = "/bin/sh"
+        command = ["-c", f"{(self._init_command)};exec {shell_command}"] if self._init_command else None
+
+        # Check if a container with the same name already exists and remove it
+        try:
+            existing_container = await asyncio.to_thread(client.containers.get, self.container_name)
+            await asyncio.to_thread(existing_container.remove, force=True)
+        except NotFound:
+            pass
+
+        self._container = await asyncio.to_thread(
+            client.containers.create,
+            self._image,
+            name=self.container_name,
+            entrypoint=shell_command,
+            command=command,
+            tty=True,
+            detach=True,
+            auto_remove=self._auto_remove,
+            volumes={str(self.bind_dir.resolve()): {"bind": "/workspace", "mode": "rw"}, **self._extra_volumes},
+            working_dir="/workspace",
+            extra_hosts=self._extra_hosts,
+            device_requests=self._device_requests,
+        )
+        await asyncio.to_thread(self._container.start)
+
+        await _wait_for_ready(self._container)
+
+        async def cleanup() -> None:
+            await self.stop()
+            asyncio_atexit.unregister(cleanup)  # type: ignore
+
+        if self._stop_container:
+            asyncio_atexit.register(cleanup)  # type: ignore
+
+        # Check if the container is running
+        if self._container.status != "running":
+            logs_str = self._container.logs().decode("utf-8")
+            raise ValueError(f"Failed to start container from image {self._image}. Logs: {logs_str}")
+
+        self._loop = asyncio.get_running_loop()
+        self._cancellation_futures = []
+        logging.debug(f"Executor started, associated with event loop: {self._loop!r}")
+
+        self._running = True
+
+    def _to_config(self) -> DockerCommandLineCodeExecutorConfig:
+        """(Experimental) Convert the component to a config object."""
+        if self._functions:
+            logging.info("Functions will not be included in serialized configuration")
+
+        return DockerCommandLineCodeExecutorConfig(
+            image=self._image,
+            container_name=self.container_name,
+            timeout=self._timeout,
+            work_dir=str(self._work_dir) if self._work_dir else None,
+            bind_dir=str(self._bind_dir) if self._bind_dir else None,
+            auto_remove=self._auto_remove,
+            stop_container=self._stop_container,
+            functions_module=self._functions_module,
+            extra_volumes=self._extra_volumes,
+            extra_hosts=self._extra_hosts,
+            init_command=self._init_command,
+            delete_tmp_files=self._delete_tmp_files,
+        )
+
+    @classmethod
+    def _from_config(cls, config: DockerCommandLineCodeExecutorConfig) -> Self:
+        """(Experimental) Create a component from a config object."""
+
+        return cls(
+            image=config.image,
+            container_name=config.container_name,
+            timeout=config.timeout,
+            work_dir=Path(config.work_dir) if config.work_dir else None,
+            bind_dir=Path(config.bind_dir) if config.bind_dir else None,
+            auto_remove=config.auto_remove,
+            stop_container=config.stop_container,
+            functions=[],  # Functions not restored from config
+            functions_module=config.functions_module,
+            extra_volumes=config.extra_volumes,
+            extra_hosts=config.extra_hosts,
+            init_command=config.init_command,
+            delete_tmp_files=config.delete_tmp_files,
+        )

agent_dhal/agentdhal_extensions/code_executors/docker_jupyter/__init__.py

@@ -0,0 +1,10 @@
+from ._docker_jupyter import DockerJupyterCodeExecutor, DockerJupyterCodeResult
+from ._jupyter_server import DockerJupyterServer, JupyterClient, JupyterKernelClient
+
+__all__ = [
+    "DockerJupyterCodeExecutor",
+    "DockerJupyterServer",
+    "JupyterClient",
+    "JupyterKernelClient",
+    "DockerJupyterCodeResult",
+]

agent_dhal/agentdhal_extensions/code_executors/docker_jupyter/_docker_jupyter.py

@@ -0,0 +1,300 @@
+import asyncio
+import base64
+import json
+import os
+import tempfile
+import uuid
+from dataclasses import dataclass
+from pathlib import Path
+from types import TracebackType
+from typing import List, Optional, Union
+
+from agentdhal_core import CancellationToken, Component
+from agentdhal_core.code_executor import CodeBlock, CodeExecutor, CodeResult
+from agentdhal_extensions.code_executors._common import silence_pip
+from pydantic import BaseModel
+from typing_extensions import Self
+
+from ._jupyter_server import JupyterClient, JupyterConnectable, JupyterConnectionInfo, JupyterKernelClient
+
+
+@dataclass
+class DockerJupyterCodeResult(CodeResult):
+    """(Experimental) A code result class for IPython code executor."""
+
+    output_files: list[Path]
+
+
+class DockerJupyterCodeExecutorConfig(BaseModel):
+    """Configuration for JupyterCodeExecutor"""
+
+    jupyter_server: Union[JupyterConnectable, JupyterConnectionInfo]
+    kernel_name: str = "python3"
+    timeout: int = 60
+    output_dir: Optional[Union[Path, str]] = None
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class DockerJupyterCodeExecutor(CodeExecutor, Component[DockerJupyterCodeExecutorConfig]):
+    """(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.
+
+    To use this, you need to install the following dependencies:
+
+    .. code-block:: shell
+
+        pip install "agentdhal-ext[docker-jupyter-executor]"
+
+    Args:
+        jupyter_server (Union[JupyterConnectable, JupyterConnectionInfo]): The Jupyter server to use.
+        kernel_name (str): The kernel name to use. Make sure it is installed.
+            By default, it is "python3".
+        timeout (int): The timeout for code execution, by default 60.
+        output_dir (str): The directory to save output files, by default None.
+
+    Example of using it directly:
+
+    .. code-block:: python
+
+        import asyncio
+        from agentdhal_core import CancellationToken
+        from agentdhal_core.code_executor import CodeBlock
+        from agentdhal_extensions.code_executors.docker_jupyter import DockerJupyterCodeExecutor, DockerJupyterServer
+
+
+        async def main() -> None:
+            async with DockerJupyterServer() as jupyter_server:
+                async with DockerJupyterCodeExecutor(jupyter_server=jupyter_server) as executor:
+                    code_blocks = [CodeBlock(code="print('hello world!')", language="python")]
+                    code_result = await executor.execute_code_blocks(code_blocks, cancellation_token=CancellationToken())
+                    print(code_result)
+
+
+        asyncio.run(main())
+
+    Example of using it with your own jupyter image:
+
+    .. code-block:: python
+
+        import asyncio
+        from agentdhal_core import CancellationToken
+        from agentdhal_core.code_executor import CodeBlock
+        from agentdhal_extensions.code_executors.docker_jupyter import DockerJupyterCodeExecutor, DockerJupyterServer
+
+
+        async def main() -> None:
+            async with DockerJupyterServer(custom_image_name="your_custom_images_name", expose_port=8888) as jupyter_server:
+                async with DockerJupyterCodeExecutor(jupyter_server=jupyter_server) as executor:
+                    code_blocks = [CodeBlock(code="print('hello world!')", language="python")]
+                    code_result = await executor.execute_code_blocks(code_blocks, cancellation_token=CancellationToken())
+                    print(code_result)
+
+
+        asyncio.run(main())
+
+    Example of using it with :class:`~agentdhal_extensions.tools.code_execution.PythonCodeExecutionTool`:
+
+    .. code-block:: python
+
+        import asyncio
+        from agentdhal_agentchat.agents import AssistantAgent
+        from agentdhal_extensions.code_executors.docker_jupyter import DockerJupyterCodeExecutor, DockerJupyterServer
+        from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
+        from agentdhal_extensions.tools.code_execution import PythonCodeExecutionTool
+
+
+        async def main() -> None:
+            async with DockerJupyterServer() as jupyter_server:
+                async with DockerJupyterCodeExecutor(jupyter_server=jupyter_server) as executor:
+                    tool = PythonCodeExecutionTool(executor)
+                    model_client = OpenAIChatCompletionClient(model="gpt-4o")
+                    agent = AssistantAgent("assistant", model_client=model_client, tools=[tool])
+                    result = await agent.run(task="What is the 10th Fibonacci number? Use Python to calculate it.")
+                    print(result)
+
+
+        asyncio.run(main())
+
+    Example of using it inside a :class:`~agentdhal_agentchat.agents._code_executor_agent.CodeExecutorAgent`:
+
+    .. code-block:: python
+
+        import asyncio
+        from agentdhal_agentchat.agents import CodeExecutorAgent
+        from agentdhal_agentchat.messages import TextMessage
+        from agentdhal_extensions.code_executors.docker_jupyter import DockerJupyterCodeExecutor, DockerJupyterServer
+        from agentdhal_core import CancellationToken
+
+
+        async def main() -> None:
+            async with DockerJupyterServer() as jupyter_server:
+                async with DockerJupyterCodeExecutor(jupyter_server=jupyter_server) as executor:
+                    code_executor_agent = CodeExecutorAgent("code_executor", code_executor=executor)
+                    task = TextMessage(
+                        content='''Here is some code
+                ```python
+                print('Hello world')
+                ```
+                ''',
+                        source="user",
+                    )
+                    response = await code_executor_agent.on_messages([task], CancellationToken())
+                    print(response.chat_message)
+
+
+        asyncio.run(main())
+
+    """
+
+    component_config_schema = DockerJupyterCodeExecutorConfig
+    component_provider_override = "agentdhal_extensions.code_executors.docker_jupyter.DockerJupyterCodeExecutor"
+
+    def __init__(
+        self,
+        jupyter_server: Union[JupyterConnectable, JupyterConnectionInfo],
+        kernel_name: str = "python3",
+        timeout: int = 60,
+        output_dir: Path | None = None,
+    ):
+        if timeout < 1:
+            raise ValueError("Timeout must be greater than or equal to 1.")
+
+        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._output_dir = output_dir or getattr(jupyter_server, "_bind_dir", None)
+        if not self._output_dir:
+            with tempfile.TemporaryDirectory() as temp_dir:
+                self._output_dir = Path(temp_dir)
+                self._output_dir.mkdir(exist_ok=True)
+
+        self._jupyter_client = JupyterClient(self._connection_info)
+
+        self._kernel_name = kernel_name
+        self._timeout = timeout
+        self._async_jupyter_kernel_client: Optional[JupyterKernelClient] = None
+        self._kernel_id: Optional[str] = None
+
+    async def _ensure_async_kernel_client(self) -> JupyterKernelClient:
+        """Ensure that an async kernel client exists and return it."""
+        if self._kernel_id is None:
+            await self.start()
+            assert self._kernel_id is not None
+        if self._async_jupyter_kernel_client is None:
+            self._async_jupyter_kernel_client = await self._jupyter_client.get_kernel_client(self._kernel_id)
+        return self._async_jupyter_kernel_client
+
+    async def execute_code_blocks(
+        self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken
+    ) -> DockerJupyterCodeResult:
+        """(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:
+            DockerJupyterCodeResult: The result of the code execution.
+        """
+        kernel_client = await self._ensure_async_kernel_client()
+        # Wait for kernel to be ready using async client
+        is_ready = await kernel_client.wait_for_ready(timeout_seconds=self._timeout)
+        if not is_ready:
+            return DockerJupyterCodeResult(exit_code=1, output="ERROR: Kernel not ready", output_files=[])
+
+        outputs: List[str] = []
+        output_files: List[Path] = []
+        for code_block in code_blocks:
+            code = silence_pip(code_block.code, code_block.language)
+            # Execute code using async client
+            exec_task = asyncio.create_task(kernel_client.execute(code, timeout_seconds=self._timeout))
+            cancellation_token.link_future(exec_task)
+            result = await exec_task
+            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(path)
+                        output_files.append(Path(path))
+                    elif data.mime_type == "text/html":
+                        path = self._save_html(data.data)
+                        outputs.append(path)
+                        output_files.append(Path(path))
+                    else:
+                        outputs.append(json.dumps(data.data))
+            else:
+                existing_output = "\n".join([str(output) for output in outputs])
+                return DockerJupyterCodeResult(
+                    exit_code=1, output=existing_output + "\nERROR: " + result.output, output_files=output_files
+                )
+        return DockerJupyterCodeResult(
+            exit_code=0, output="\n".join([str(output) for output in outputs]), output_files=output_files
+        )
+
+    async def restart(self) -> None:
+        """(Experimental) Restart a new session."""
+        # Use async client to restart kernel
+        if self._kernel_id is not None:
+            await self._jupyter_client.restart_kernel(self._kernel_id)
+        # Reset the clients to force recreation
+        if self._async_jupyter_kernel_client is not None:
+            await self._async_jupyter_kernel_client.stop()
+            self._async_jupyter_kernel_client = None
+
+    async def start(self) -> None:
+        """(Experimental) Start a new session."""
+        available_kernels = await self._jupyter_client.list_kernel_specs()
+        if self._kernel_name not in available_kernels["kernelspecs"]:
+            raise ValueError(f"Kernel {self._kernel_name} is not installed.")
+        self._kernel_id = await self._jupyter_client.start_kernel(self._kernel_name)
+
+    def _save_image(self, image_data_base64: str) -> str:
+        """Save image data to a file."""
+        image_data = base64.b64decode(image_data_base64)
+        filename = f"{uuid.uuid4().hex}.png"
+        path = os.path.join(str(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."""
+        filename = f"{uuid.uuid4().hex}.html"
+        path = os.path.join(str(self._output_dir), filename)
+        with open(path, "w") as f:
+            f.write(html_data)
+        return os.path.abspath(path)
+
+    async def stop(self) -> None:
+        """Stop the kernel."""
+        if self._kernel_id is not None:
+            await self._jupyter_client.delete_kernel(self._kernel_id)
+        if self._async_jupyter_kernel_client is not None:
+            await self._async_jupyter_kernel_client.stop()
+            self._async_jupyter_kernel_client = None
+        await self._jupyter_client.close()
+
+    async def __aenter__(self) -> Self:
+        await self.start()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        await self.stop()

agent_dhal/agentdhal_extensions/code_executors/docker_jupyter/_jupyter_server.py

@@ -0,0 +1,430 @@
+import asyncio
+import atexit
+import datetime
+import io
+import json
+import logging
+import os
+import secrets
+import uuid
+from dataclasses import dataclass
+from pathlib import Path
+from time import sleep
+from types import TracebackType
+from typing import Any, Dict, List, Optional, Protocol, Type, Union, cast, runtime_checkable
+
+import aiohttp
+import docker
+import docker.errors
+import requests
+import websockets
+from requests.adapters import HTTPAdapter, Retry
+from typing_extensions import Self
+
+
+@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."""
+        ...
+
+
+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))
+        # Create aiohttp session for async requests
+        self._async_session: aiohttp.ClientSession | None = None
+
+    async def _ensure_async_session(self) -> aiohttp.ClientSession:
+        if self._async_session is None:
+            self._async_session = aiohttp.ClientSession()
+        return self._async_session
+
+    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}"
+
+    async 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())
+
+    async 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())
+
+    async def start_kernel(self, kernel_spec_name: str) -> str:
+        """Start a new kernel asynchronously.
+
+        Args:
+            kernel_spec_name (str): Name of the kernel spec to start
+
+        Returns:
+            str: ID of the started kernel
+        """
+        session = await self._ensure_async_session()
+        async with session.post(
+            f"{self._get_api_base_url()}/api/kernels",
+            headers=self._get_headers(),
+            json={"name": kernel_spec_name},
+        ) as response:
+            data = await response.json()
+            return cast(str, data["id"])
+
+    async def delete_kernel(self, kernel_id: str) -> None:
+        session = await self._ensure_async_session()
+        async with session.delete(
+            f"{self._get_api_base_url()}/api/kernels/{kernel_id}", headers=self._get_headers()
+        ) as response:
+            response.raise_for_status()
+
+    async def restart_kernel(self, kernel_id: str) -> None:
+        session = await self._ensure_async_session()
+        async with session.post(
+            f"{self._get_api_base_url()}/api/kernels/{kernel_id}/restart", headers=self._get_headers()
+        ) as response:
+            response.raise_for_status()
+
+    async def get_kernel_client(self, kernel_id: str) -> "JupyterKernelClient":
+        ws_url = f"{self._get_ws_base_url()}/api/kernels/{kernel_id}/channels"
+        # Using websockets library for async websocket connections
+        ws = await websockets.connect(ws_url, additional_headers=self._get_headers())
+        return JupyterKernelClient(ws)
+
+    async def close(self) -> None:
+        """Close the async session"""
+        if self._async_session is not None:
+            await self._async_session.close()
+            self._async_session = None
+        self._session.close()
+
+
+@dataclass
+class DataItem:
+    mime_type: str
+    data: str
+
+
+@dataclass
+class ExecutionResult:
+    is_ok: bool
+    output: str
+    data_items: List[DataItem]
+
+
+class JupyterKernelClient:
+    """An asynchronous client for communicating with a Jupyter kernel."""
+
+    def __init__(self, websocket: websockets.ClientConnection) -> None:
+        self._session_id = uuid.uuid4().hex
+        self._websocket = websocket
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(
+        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+    ) -> None:
+        await self.stop()
+
+    async def stop(self) -> None:
+        await self._websocket.close()
+
+    async 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": "agentdhal",
+                "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": {},
+        }
+        await self._websocket.send(json.dumps(message))
+        return message_id
+
+    async def _receive_message(self, timeout_seconds: Optional[float]) -> Optional[Dict[str, Any]]:
+        try:
+            if timeout_seconds is not None:
+                data = await asyncio.wait_for(self._websocket.recv(), timeout=timeout_seconds)
+            else:
+                data = await self._websocket.recv()
+            if isinstance(data, bytes):
+                return cast(Dict[str, Any], json.loads(data.decode("utf-8")))
+            return cast(Dict[str, Any], json.loads(data))
+        except asyncio.TimeoutError:
+            return None
+
+    async def wait_for_ready(self, timeout_seconds: Optional[float] = None) -> bool:
+        message_id = await self._send_message(content={}, channel="shell", message_type="kernel_info_request")
+        while True:
+            message = await 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
+
+    async def execute(self, code: str, timeout_seconds: Optional[float] = None) -> ExecutionResult:
+        message_id = await 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: List[str] = []
+        data_output: List[DataItem] = []
+        while True:
+            message = await self._receive_message(timeout_seconds)
+            if message is None:
+                return 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(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 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 ExecutionResult(
+            is_ok=True, output="\n".join([str(output) for output in text_output]), data_items=data_output
+        )
+
+
+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: Optional[Dict[str, str]] = None,
+        expose_port: int = 8888,
+        token: Optional[Union[str, GenerateToken]] = None,
+        work_dir: Union[Path, str] = "/workspace",
+        bind_dir: Optional[Union[Path, str]] = None,
+    ):
+        """Start a Jupyter kernel gateway server in a Docker container.
+
+        Args:
+            custom_image_name: Custom Docker image to use. If None, builds and uses bundled image.
+            container_name: Name for the Docker container. Auto-generated if None.
+            auto_remove: If True, container will be deleted when stopped.
+            stop_container: If True, container stops on program exit or when context manager exits.
+            docker_env: Additional environment variables for the container.
+            expose_port: Port to expose for Jupyter connection.
+            token: Authentication token. If GenerateToken, creates random token. Empty for no auth.
+            work_dir: Working directory inside the container.
+            bind_dir: Local directory to bind to container's work_dir.
+        """
+        # Generate container name if not provided
+        container_name = container_name or f"agentdhal-jupyterkernelgateway-{uuid.uuid4()}"
+
+        # Initialize Docker client
+        client = docker.from_env()
+        # Set up bind directory if specified
+        self._bind_dir: Optional[Path] = None
+        if bind_dir:
+            self._bind_dir = Path(bind_dir) if isinstance(bind_dir, str) else bind_dir
+            self._bind_dir.mkdir(exist_ok=True)
+            os.chmod(bind_dir, 0o777)
+
+        # Determine and prepare Docker image
+        image_name = custom_image_name or "agentdhal-jupyterkernelgateway"
+        if not custom_image_name:
+            try:
+                client.images.get(image_name)
+            except docker.errors.ImageNotFound:
+                # Build default image if not found
+                here = Path(__file__).parent
+                dockerfile = io.BytesIO(self.DEFAULT_DOCKERFILE.encode("utf-8"))
+                logging.info(f"Building image {image_name}...")
+                client.images.build(path=str(here), fileobj=dockerfile, tag=image_name)
+                logging.info(f"Image {image_name} built successfully")
+        else:
+            # Verify custom image exists
+            try:
+                client.images.get(image_name)
+            except docker.errors.ImageNotFound as err:
+                raise ValueError(f"Custom image {image_name} does not exist") from err
+        if docker_env is None:
+            docker_env = {}
+        if token is None:
+            token = DockerJupyterServer.GenerateToken()
+        # Set up authentication token
+        self._token = secrets.token_hex(32) if isinstance(token, DockerJupyterServer.GenerateToken) else token
+
+        # Prepare environment variables
+        env = {"TOKEN": self._token}
+        env.update(docker_env)
+
+        # Define volume configuration if bind directory is specified
+        volumes = {str(self._bind_dir): {"bind": str(work_dir), "mode": "rw"}} if self._bind_dir else None
+
+        # Start the container
+        container = client.containers.run(
+            image_name,
+            detach=True,
+            auto_remove=auto_remove,
+            environment=env,
+            publish_all_ports=True,
+            name=container_name,
+            volumes=volumes,
+            working_dir=str(work_dir),
+        )
+
+        # Wait for container to be ready
+        self._wait_for_ready(container)
+
+        # Store container information
+        self._container = container
+        self._port = int(container.ports[f"{expose_port}/tcp"][0]["HostPort"])
+        self._container_id = container.id
+        self._expose_port = expose_port
+
+        if self._container_id is None:
+            raise ValueError("Failed to obtain container id.")
+
+        # Define cleanup function
+        def cleanup() -> None:
+            try:
+                assert self._container_id is not None
+                inner_container = client.containers.get(self._container_id)
+                inner_container.stop()
+            except docker.errors.NotFound:
+                pass
+            atexit.unregister(cleanup)
+
+        # Register cleanup if container should be stopped automatically
+        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 _wait_for_ready(self, container: Any, timeout: int = 60, stop_time: float = 0.1) -> None:
+        elapsed_time = 0.0
+        while container.status != "running" and elapsed_time < timeout:
+            sleep(stop_time)
+            elapsed_time += stop_time
+            container.reload()
+            continue
+        if container.status != "running":
+            raise ValueError("Container failed to start")
+
+    async def stop(self) -> None:
+        loop = asyncio.get_event_loop()
+        await loop.run_in_executor(None, self._cleanup_func)
+
+    async def get_client(self) -> JupyterClient:
+        return JupyterClient(self.connection_info)
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(
+        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+    ) -> None:
+        await self.stop()

agent_dhal/agentdhal_extensions/code_executors/jupyter/__init__.py

@@ -0,0 +1,6 @@
+from ._jupyter_code_executor import JupyterCodeExecutor, JupyterCodeResult
+
+__all__ = [
+    "JupyterCodeExecutor",
+    "JupyterCodeResult",
+]

agent_dhal/agentdhal_extensions/code_executors/jupyter/_jupyter_code_executor.py

@@ -0,0 +1,335 @@
+import asyncio
+import base64
+import json
+import re
+import sys
+import tempfile
+import uuid
+import warnings
+from dataclasses import dataclass
+from pathlib import Path
+
+from agentdhal_core import Component
+from pydantic import BaseModel
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+from contextlib import AbstractAsyncContextManager
+from typing import Optional, Union
+
+from agentdhal_core import CancellationToken
+from agentdhal_core.code_executor import CodeBlock, CodeExecutor, CodeResult
+from nbclient import NotebookClient
+from nbformat import NotebookNode
+from nbformat import v4 as nbformat
+from typing_extensions import Self
+
+from .._common import silence_pip
+
+
+@dataclass
+class JupyterCodeResult(CodeResult):
+    """A code result class for Jupyter code executor."""
+
+    output_files: list[Path]
+
+
+class JupyterCodeExecutorConfig(BaseModel):
+    """Configuration for JupyterCodeExecutor"""
+
+    kernel_name: str = "python3"
+    timeout: int = 60
+    output_dir: Optional[str] = None
+
+
+class JupyterCodeExecutor(CodeExecutor, Component[JupyterCodeExecutorConfig]):
+    """A code executor class that executes code statefully using [nbclient](https://github.com/jupyter/nbclient).
+
+    .. danger::
+
+        This will execute code on the local machine. If being used with LLM generated code, caution should be used.
+
+    Example of using it directly:
+
+    .. code-block:: python
+
+        import asyncio
+        from agentdhal_core import CancellationToken
+        from agentdhal_core.code_executor import CodeBlock
+        from agentdhal_extensions.code_executors.jupyter import JupyterCodeExecutor
+
+
+        async def main() -> None:
+            async with JupyterCodeExecutor() as executor:
+                cancel_token = CancellationToken()
+                code_blocks = [CodeBlock(code="print('hello world!')", language="python")]
+                code_result = await executor.execute_code_blocks(code_blocks, cancel_token)
+                print(code_result)
+
+
+        asyncio.run(main())
+
+    Example of using it with :class:`~agentdhal_extensions.tools.code_execution.PythonCodeExecutionTool`:
+
+    .. code-block:: python
+
+        import asyncio
+        from agentdhal_agentchat.agents import AssistantAgent
+        from agentdhal_extensions.code_executors.jupyter import JupyterCodeExecutor
+        from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
+        from agentdhal_extensions.tools.code_execution import PythonCodeExecutionTool
+
+
+        async def main() -> None:
+            async with JupyterCodeExecutor() as executor:
+                tool = PythonCodeExecutionTool(executor)
+                model_client = OpenAIChatCompletionClient(model="gpt-4o")
+                agent = AssistantAgent("assistant", model_client=model_client, tools=[tool])
+                result = await agent.run(task="What is the 10th Fibonacci number? Use Python to calculate it.")
+                print(result)
+
+
+        asyncio.run(main())
+
+    Example of using it inside a :class:`~agentdhal_agentchat.agents._code_executor_agent.CodeExecutorAgent`:
+
+    .. code-block:: python
+
+        import asyncio
+        from agentdhal_agentchat.agents import CodeExecutorAgent
+        from agentdhal_agentchat.messages import TextMessage
+        from agentdhal_extensions.code_executors.jupyter import JupyterCodeExecutor
+        from agentdhal_core import CancellationToken
+
+
+        async def main() -> None:
+            async with JupyterCodeExecutor() as executor:
+                code_executor_agent = CodeExecutorAgent("code_executor", code_executor=executor)
+                task = TextMessage(
+                    content='''Here is some code
+            ```python
+            print('Hello world')
+            ```
+            ''',
+                    source="user",
+                )
+                response = await code_executor_agent.on_messages([task], CancellationToken())
+                print(response.chat_message)
+
+
+        asyncio.run(main())
+
+
+    Args:
+        kernel_name (str): The kernel name to use. By default, "python3".
+        timeout (int): The timeout for code execution, by default 60.
+        output_dir (Path): The directory to save output files, by default a temporary directory.
+
+
+    .. note::
+        Using the current directory (".") as output directory is deprecated. Using it will raise a deprecation warning.
+    """
+
+    component_config_schema = JupyterCodeExecutorConfig
+    component_provider_override = "agentdhal_extensions.code_executors.jupyter.JupyterCodeExecutor"
+
+    def __init__(
+        self,
+        kernel_name: str = "python3",
+        timeout: int = 60,
+        output_dir: Optional[Union[Path, str]] = None,
+    ):
+        if timeout < 1:
+            raise ValueError("Timeout must be greater than or equal to 1.")
+
+        self._output_dir: Path = Path(tempfile.mkdtemp()) if output_dir is None else Path(output_dir)
+        self._output_dir.mkdir(exist_ok=True, parents=True)
+
+        self._temp_dir: Optional[tempfile.TemporaryDirectory[str]] = None
+        self._temp_dir_path: Optional[Path] = None
+
+        self._started = False
+
+        self._kernel_name = kernel_name
+        self._timeout = timeout
+
+        self._client: Optional[NotebookClient] = None
+        self.kernel_context: Optional[AbstractAsyncContextManager[None]] = None
+
+    async def execute_code_blocks(
+        self, code_blocks: list[CodeBlock], cancellation_token: CancellationToken
+    ) -> JupyterCodeResult:
+        """Execute code blocks and return the result.
+
+        Args:
+            code_blocks (list[CodeBlock]): The code blocks to execute.
+
+        Returns:
+            JupyterCodeResult: The result of the code execution.
+        """
+        outputs: list[str] = []
+        output_files: list[Path] = []
+        exit_code = 0
+
+        for code_block in code_blocks:
+            result = await self._execute_code_block(code_block, cancellation_token)
+            exit_code = result.exit_code
+            outputs.append(result.output)
+            output_files.extend(result.output_files)
+
+            # Stop execution if one code block fails
+            if exit_code != 0:
+                break
+
+        return JupyterCodeResult(exit_code=exit_code, output="\n".join(outputs), output_files=output_files)
+
+    async def _execute_code_block(
+        self, code_block: CodeBlock, cancellation_token: CancellationToken
+    ) -> JupyterCodeResult:
+        """Execute single code block and return the result.
+
+        Args:
+            code_block (CodeBlock): The code block to execute.
+
+        Returns:
+            JupyterCodeResult: The result of the code execution.
+        """
+        execute_task = asyncio.create_task(
+            self._execute_cell(
+                nbformat.new_code_cell(silence_pip(code_block.code, code_block.language))  # type: ignore
+            )
+        )
+
+        cancellation_token.link_future(execute_task)
+        output_cell = await asyncio.wait_for(asyncio.shield(execute_task), timeout=self._timeout)
+
+        outputs: list[str] = []
+        output_files: list[Path] = []
+        exit_code = 0
+
+        for output in output_cell.get("outputs", []):
+            match output.get("output_type"):
+                case "stream":
+                    outputs.append(output.get("text", ""))
+                case "error":
+                    traceback = re.sub(r"\x1b\[[0-9;]*[A-Za-z]", "", "\n".join(output["traceback"]))
+                    outputs.append(traceback)
+                    exit_code = 1
+                case "execute_result" | "display_data":
+                    data = output.get("data", {})
+                    for mime, content in data.items():
+                        match mime:
+                            case "text/plain":
+                                outputs.append(content)
+                            case "image/png":
+                                path = self._save_image(content)
+                                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(content)
+                                output_files.append(path)
+                            case _:
+                                outputs.append(json.dumps(content))
+                case _:
+                    pass
+
+        return JupyterCodeResult(exit_code=exit_code, output="\n".join(outputs), output_files=output_files)
+
+    async def _execute_cell(self, cell: NotebookNode) -> NotebookNode:
+        # Temporary push cell to nb as async_execute_cell expects it. But then we want to remove it again as cells can take up significant amount of memory (especially with images)
+        if not self._client:
+            raise RuntimeError("Executor must be started before executing cells")
+        self._client.nb.cells.append(cell)
+        output = await self._client.async_execute_cell(
+            cell,
+            cell_index=0,
+        )
+        self._client.nb.cells.pop()
+        return output
+
+    def _save_image(self, image_data_base64: str) -> Path:
+        """Save image data to a file."""
+        image_data = base64.b64decode(image_data_base64)
+        path = self._output_dir / f"{uuid.uuid4().hex}.png"
+        path.write_bytes(image_data)
+        return path.absolute()
+
+    def _save_html(self, html_data: str) -> Path:
+        """Save HTML data to a file."""
+        path = self._output_dir / f"{uuid.uuid4().hex}.html"
+        path.write_text(html_data)
+        return path.absolute()
+
+    async def restart(self) -> None:
+        """Restart the code executor."""
+        await self.stop()
+        await self.start()
+
+    async def start(self) -> None:
+        """(Experimental) Start the code executor.
+
+        Initializes the Jupyter Notebook execution environment by creating a new notebook and setting it up with the specified Jupyter Kernel.
+        Marks the executor as started, allowing for code execution.
+        This method should be called before executing any code blocks.
+        """
+        if self._started:
+            return
+
+        notebook: NotebookNode = nbformat.new_notebook()  # type: ignore
+
+        self._client = NotebookClient(
+            nb=notebook,
+            kernel_name=self._kernel_name,
+            timeout=self._timeout,
+            allow_errors=True,
+        )
+
+        self.kernel_context = self._client.async_setup_kernel()
+        await self.kernel_context.__aenter__()
+
+        self._started = True
+
+    async def stop(self) -> None:
+        """(Experimental) Stop the code executor.
+
+        Terminates the Jupyter Notebook execution by exiting the kernel context and cleaning up the associated resources."""
+        if not self._started:
+            return
+
+        if self.kernel_context is not None:
+            await self.kernel_context.__aexit__(None, None, None)
+            self.kernel_context = None
+
+        self._client = None
+        self._started = False
+
+    def _to_config(self) -> JupyterCodeExecutorConfig:
+        """Convert current instance to config object"""
+        return JupyterCodeExecutorConfig(
+            kernel_name=self._kernel_name, timeout=self._timeout, output_dir=str(self.output_dir)
+        )
+
+    @property
+    def output_dir(self) -> Path:
+        # If a user specifies the current directory, warn them that this is deprecated
+        if self._output_dir == Path("."):
+            warnings.warn(
+                "Using the current directory as output_dir is deprecated",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        return self._output_dir
+
+    @classmethod
+    def _from_config(cls, config: JupyterCodeExecutorConfig) -> Self:
+        """Create instance from config object"""
+        return cls(
+            kernel_name=config.kernel_name,
+            timeout=config.timeout,
+            output_dir=Path(config.output_dir) if config.output_dir else None,
+        )

agent_dhal/agentdhal_extensions/code_executors/local/__init__.py

@@ -0,0 +1,517 @@
+# File based from: https://github.com/microsoft/autogen/blob/main/autogen/coding/local_commandline_code_executor.py
+# Credit to original authors
+
+import asyncio
+import logging
+import os
+import sys
+import tempfile
+import warnings
+from hashlib import sha256
+from pathlib import Path
+from string import Template
+from types import SimpleNamespace
+from typing import Any, Callable, ClassVar, List, Optional, Sequence, Union
+
+from agentdhal_core import CancellationToken, Component
+from agentdhal_core.code_executor import CodeBlock, CodeExecutor, FunctionWithRequirements, FunctionWithRequirementsStr
+from pydantic import BaseModel
+from typing_extensions import ParamSpec, Self
+
+from .._common import (
+    PYTHON_VARIANTS,
+    CommandLineCodeResult,
+    build_python_functions_file,
+    get_file_name_from_content,
+    lang_to_cmd,
+    silence_pip,
+    to_stub,
+)
+
+__all__ = ("LocalCommandLineCodeExecutor",)
+
+A = ParamSpec("A")
+
+
+class LocalCommandLineCodeExecutorConfig(BaseModel):
+    """Configuration for LocalCommandLineCodeExecutor"""
+
+    timeout: int = 60
+    work_dir: Optional[str] = None
+    functions_module: str = "functions"
+    cleanup_temp_files: bool = True
+
+
+class LocalCommandLineCodeExecutor(CodeExecutor, Component[LocalCommandLineCodeExecutorConfig]):
+    """A code executor class that executes code through a local command line
+    environment.
+
+    .. danger::
+
+        This will execute code on the local machine. If being used with LLM generated code, caution should be used.
+
+    Each code block is saved as a file and executed in a separate process in
+    the working directory, and a unique file is generated and saved in the
+    working directory for each code block.
+    The code blocks are executed in the order they are received.
+    Command line code is sanitized using regular expression match against a list of dangerous commands in order to prevent self-destructive
+    commands from being executed which may potentially affect the users environment.
+    Currently the only supported languages is Python and shell scripts.
+    For Python code, use the language "python" for the code block.
+    For shell scripts, use the language "bash", "shell", "sh", "pwsh", "powershell", or "ps1" for the code
+    block.
+
+    .. note::
+
+        On Windows, the event loop policy must be set to `WindowsProactorEventLoopPolicy` to avoid issues with subprocesses.
+
+        .. code-block:: python
+
+            import sys
+            import asyncio
+
+            if sys.platform == "win32":
+                asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
+
+    Args:
+        timeout (int): The timeout for the execution of any single code block. Default is 60.
+        work_dir (str): The working directory for the code execution. If None,
+            a default working directory will be used. The default working directory is a temporary directory.
+        functions (List[Union[FunctionWithRequirements[Any, A], Callable[..., Any]]]): A list of functions that are available to the code executor. Default is an empty list.
+        functions_module (str, optional): The name of the module that will be created to store the functions. Defaults to "functions".
+        cleanup_temp_files (bool, optional): Whether to automatically clean up temporary files after execution. Defaults to True.
+        virtual_env_context (Optional[SimpleNamespace], optional): The virtual environment context. Defaults to None.
+
+    .. note::
+        Using the current directory (".") as working directory is deprecated. Using it will raise a deprecation warning.
+
+
+    Example:
+
+    How to use `LocalCommandLineCodeExecutor` with a virtual environment different from the one used to run the autogen application:
+    Set up a virtual environment using the `venv` module, and pass its context to the initializer of `LocalCommandLineCodeExecutor`. This way, the executor will run code within the new environment.
+
+        .. code-block:: python
+
+            import venv
+            from pathlib import Path
+            import asyncio
+
+            from agentdhal_core import CancellationToken
+            from agentdhal_core.code_executor import CodeBlock
+            from agentdhal_extensions.code_executors.local import LocalCommandLineCodeExecutor
+
+
+            async def example():
+                work_dir = Path("coding")
+                work_dir.mkdir(exist_ok=True)
+
+                venv_dir = work_dir / ".venv"
+                venv_builder = venv.EnvBuilder(with_pip=True)
+                venv_builder.create(venv_dir)
+                venv_context = venv_builder.ensure_directories(venv_dir)
+
+                local_executor = LocalCommandLineCodeExecutor(work_dir=work_dir, virtual_env_context=venv_context)
+                await local_executor.execute_code_blocks(
+                    code_blocks=[
+                        CodeBlock(language="bash", code="pip install matplotlib"),
+                    ],
+                    cancellation_token=CancellationToken(),
+                )
+
+
+            asyncio.run(example())
+
+    """
+
+    component_config_schema = LocalCommandLineCodeExecutorConfig
+    component_provider_override = "agentdhal_extensions.code_executors.local.LocalCommandLineCodeExecutor"
+
+    SUPPORTED_LANGUAGES: ClassVar[List[str]] = [
+        "bash",
+        "shell",
+        "sh",
+        "pwsh",
+        "powershell",
+        "ps1",
+        "python",
+    ]
+    FUNCTION_PROMPT_TEMPLATE: ClassVar[
+        str
+    ] = """You have access to the following user defined functions. They can be accessed from the module called `$module_name` by their function names.
+
+For example, if there was a function called `foo` you could import it by writing `from $module_name import foo`
+
+$functions"""
+
+    def __init__(
+        self,
+        timeout: int = 60,
+        work_dir: Optional[Union[Path, str]] = None,
+        functions: Sequence[
+            Union[
+                FunctionWithRequirements[Any, A],
+                Callable[..., Any],
+                FunctionWithRequirementsStr,
+            ]
+        ] = [],
+        functions_module: str = "functions",
+        cleanup_temp_files: bool = True,
+        virtual_env_context: Optional[SimpleNamespace] = None,
+    ):
+        if timeout < 1:
+            raise ValueError("Timeout must be greater than or equal to 1.")
+        self._timeout = timeout
+
+        self._work_dir: Optional[Path] = None
+        if work_dir is not None:
+            # Check if user provided work_dir is the current directory and warn if so.
+            if Path(work_dir).resolve() == Path.cwd().resolve():
+                warnings.warn(
+                    "Using the current directory as work_dir is deprecated.",
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
+            if isinstance(work_dir, str):
+                self._work_dir = Path(work_dir)
+            else:
+                self._work_dir = work_dir
+            self._work_dir.mkdir(exist_ok=True)
+
+        self._functions = functions
+        # Setup could take some time so we intentionally wait for the first code block to do it.
+        if len(functions) > 0:
+            self._setup_functions_complete = False
+        else:
+            self._setup_functions_complete = True
+
+        if not functions_module.isidentifier():
+            raise ValueError("Module name must be a valid Python identifier")
+        self._functions_module = functions_module
+
+        self._cleanup_temp_files = cleanup_temp_files
+        self._virtual_env_context: Optional[SimpleNamespace] = virtual_env_context
+
+        self._temp_dir: Optional[tempfile.TemporaryDirectory[str]] = None
+        self._started = False
+
+        # Check the current event loop policy if on windows.
+        if sys.platform == "win32":
+            current_policy = asyncio.get_event_loop_policy()
+            if hasattr(asyncio, "WindowsProactorEventLoopPolicy") and not isinstance(
+                current_policy, asyncio.WindowsProactorEventLoopPolicy
+            ):
+                warnings.warn(
+                    "The current event loop policy is not WindowsProactorEventLoopPolicy. "
+                    "This may cause issues with subprocesses. "
+                    "Try setting the event loop policy to WindowsProactorEventLoopPolicy. "
+                    "For example: `asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())`. "
+                    "See https://docs.python.org/3/library/asyncio-eventloop.html#asyncio.ProactorEventLoop.",
+                    stacklevel=2,
+                )
+
+    def format_functions_for_prompt(self, prompt_template: str = FUNCTION_PROMPT_TEMPLATE) -> str:
+        """(Experimental) Format the functions for a prompt.
+
+        The template includes two variables:
+        - `$module_name`: The module name.
+        - `$functions`: The functions formatted as stubs with two newlines between each function.
+
+        Args:
+            prompt_template (str): The prompt template. Default is the class default.
+
+        Returns:
+            str: The formatted prompt.
+        """
+
+        template = Template(prompt_template)
+        return template.substitute(
+            module_name=self._functions_module,
+            functions="\n\n".join([to_stub(func) for func in self._functions]),
+        )
+
+    @property
+    def timeout(self) -> int:
+        """(Experimental) The timeout for code execution."""
+        return self._timeout
+
+    @property
+    def work_dir(self) -> Path:
+        """(Experimental) The working directory for the code execution."""
+        if self._work_dir is not None:
+            return self._work_dir
+        else:
+            # Automatically create temp directory if not exists
+            if self._temp_dir is None:
+                self._temp_dir = tempfile.TemporaryDirectory()
+                self._started = True
+            return Path(self._temp_dir.name)
+
+    @property
+    def functions(self) -> List[str]:
+        raise NotImplementedError
+
+    @property
+    def functions_module(self) -> str:
+        """(Experimental) The module name for the functions."""
+        return self._functions_module
+
+    @property
+    def cleanup_temp_files(self) -> bool:
+        """(Experimental) Whether to automatically clean up temporary files after execution."""
+        return self._cleanup_temp_files
+
+    async def _setup_functions(self, cancellation_token: CancellationToken) -> None:
+        func_file_content = build_python_functions_file(self._functions)
+        func_file = self.work_dir / f"{self._functions_module}.py"
+        func_file.write_text(func_file_content)
+
+        # Collect requirements
+        lists_of_packages = [x.python_packages for x in self._functions if isinstance(x, FunctionWithRequirements)]
+        flattened_packages = [item for sublist in lists_of_packages for item in sublist]
+        required_packages = list(set(flattened_packages))
+        if len(required_packages) > 0:
+            logging.info("Ensuring packages are installed in executor.")
+
+            cmd_args = ["-m", "pip", "install"]
+            cmd_args.extend(required_packages)
+
+            if self._virtual_env_context:
+                py_executable = self._virtual_env_context.env_exe
+            else:
+                py_executable = sys.executable
+
+            task = asyncio.create_task(
+                asyncio.create_subprocess_exec(
+                    py_executable,
+                    *cmd_args,
+                    cwd=self.work_dir,
+                    stdout=asyncio.subprocess.PIPE,
+                    stderr=asyncio.subprocess.PIPE,
+                )
+            )
+            cancellation_token.link_future(task)
+            try:
+                proc = await task
+                stdout, stderr = await asyncio.wait_for(proc.communicate(), self._timeout)
+            except asyncio.TimeoutError as e:
+                raise ValueError("Pip install timed out") from e
+            except asyncio.CancelledError as e:
+                raise ValueError("Pip install was cancelled") from e
+
+            if proc.returncode is not None and proc.returncode != 0:
+                raise ValueError(f"Pip install failed. {stdout.decode()}, {stderr.decode()}")
+
+        # Attempt to load the function file to check for syntax errors, imports etc.
+        exec_result = await self._execute_code_dont_check_setup(
+            [CodeBlock(code=func_file_content, language="python")], cancellation_token
+        )
+
+        if exec_result.exit_code != 0:
+            raise ValueError(f"Functions failed to load: {exec_result.output}")
+
+        self._setup_functions_complete = True
+
+    async def execute_code_blocks(
+        self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken
+    ) -> CommandLineCodeResult:
+        """(Experimental) Execute the code blocks and return the result.
+
+        Args:
+            code_blocks (List[CodeBlock]): The code blocks to execute.
+            cancellation_token (CancellationToken): a token to cancel the operation
+
+        Returns:
+            CommandLineCodeResult: The result of the code execution."""
+
+        if not self._setup_functions_complete:
+            await self._setup_functions(cancellation_token)
+
+        return await self._execute_code_dont_check_setup(code_blocks, cancellation_token)
+
+    async def _execute_code_dont_check_setup(
+        self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken
+    ) -> CommandLineCodeResult:
+        """
+        Execute the provided code blocks in the local command line without re-checking setup.
+        Returns a CommandLineCodeResult indicating success or failure.
+        """
+        logs_all: str = ""
+        file_names: List[Path] = []
+        exitcode = 0
+
+        for code_block in code_blocks:
+            lang, code = code_block.language, code_block.code
+            lang = lang.lower()
+
+            # Remove pip output where possible
+            code = silence_pip(code, lang)
+
+            # Normalize python variants to "python"
+            if lang in PYTHON_VARIANTS:
+                lang = "python"
+
+            # Abort if not supported
+            if lang not in self.SUPPORTED_LANGUAGES:
+                exitcode = 1
+                logs_all += "\n" + f"unknown language {lang}"
+                break
+
+            # Try extracting a filename (if present)
+            try:
+                filename = get_file_name_from_content(code, self.work_dir)
+            except ValueError:
+                return CommandLineCodeResult(
+                    exit_code=1,
+                    output="Filename is not in the workspace",
+                    code_file=None,
+                )
+
+            # If no filename is found, create one
+            if filename is None:
+                code_hash = sha256(code.encode()).hexdigest()
+                if lang.startswith("python"):
+                    ext = "py"
+                elif lang in ["pwsh", "powershell", "ps1"]:
+                    ext = "ps1"
+                else:
+                    ext = lang
+
+                filename = f"tmp_code_{code_hash}.{ext}"
+
+            written_file = (self.work_dir / filename).resolve()
+            with written_file.open("w", encoding="utf-8") as f:
+                f.write(code)
+            file_names.append(written_file)
+
+            # Build environment
+            env = os.environ.copy()
+            if self._virtual_env_context:
+                virtual_env_bin_abs_path = os.path.abspath(self._virtual_env_context.bin_path)
+                env["PATH"] = f"{virtual_env_bin_abs_path}{os.pathsep}{env['PATH']}"
+
+            # Decide how to invoke the script
+            if lang == "python":
+                program = (
+                    os.path.abspath(self._virtual_env_context.env_exe) if self._virtual_env_context else sys.executable
+                )
+                extra_args = [str(written_file.absolute())]
+            else:
+                # Get the appropriate command for the language
+                program = lang_to_cmd(lang)
+
+                # Special handling for PowerShell
+                if program == "pwsh":
+                    extra_args = [
+                        "-NoProfile",
+                        "-ExecutionPolicy",
+                        "Bypass",
+                        "-File",
+                        str(written_file.absolute()),
+                    ]
+                else:
+                    # Shell commands (bash, sh, etc.)
+                    extra_args = [str(written_file.absolute())]
+
+            # Create a subprocess and run
+            task = asyncio.create_task(
+                asyncio.create_subprocess_exec(
+                    program,
+                    *extra_args,
+                    cwd=self.work_dir,
+                    stdout=asyncio.subprocess.PIPE,
+                    stderr=asyncio.subprocess.PIPE,
+                    env=env,
+                )
+            )
+            cancellation_token.link_future(task)
+
+            proc = None  # Track the process
+            try:
+                proc = await task
+                stdout, stderr = await asyncio.wait_for(proc.communicate(), self._timeout)
+                exitcode = proc.returncode or 0
+            except asyncio.TimeoutError:
+                logs_all += "\nTimeout"
+                exitcode = 124
+                if proc:
+                    proc.terminate()
+                    await proc.wait()  # Ensure process is fully dead
+                break
+            except asyncio.CancelledError:
+                logs_all += "\nCancelled"
+                exitcode = 125
+                if proc:
+                    proc.terminate()
+                    await proc.wait()
+                break
+
+            logs_all += stderr.decode()
+            logs_all += stdout.decode()
+
+            if exitcode != 0:
+                break
+
+        code_file = str(file_names[0]) if file_names else None
+        code_result = CommandLineCodeResult(exit_code=exitcode, output=logs_all, code_file=code_file)
+
+        if self._cleanup_temp_files:
+            for file in file_names:
+                try:
+                    file.unlink(missing_ok=True)
+                except OSError as error:
+                    logging.error(f"Failed to delete temporary file {file}: {error}")
+
+        return code_result
+
+    async def restart(self) -> None:
+        """(Experimental) Restart the code executor."""
+        warnings.warn(
+            "Restarting local command line code executor is not supported. No action is taken.",
+            stacklevel=2,
+        )
+
+    async def start(self) -> None:
+        """(Experimental) Start the code executor.
+
+        Initializes the local code executor and should be called before executing any code blocks.
+        It marks the executor internal state as started.
+        If no working directory is provided, the method creates a temporary directory for the executor to use.
+        """
+        if self._work_dir is None and self._temp_dir is None:
+            self._temp_dir = tempfile.TemporaryDirectory()
+        self._started = True
+
+    async def stop(self) -> None:
+        """(Experimental) Stop the code executor.
+
+        Stops the local code executor and performs the cleanup of the temporary working directory (if it was created).
+        The executor's internal state is markes as no longer started.
+        """
+        if self._temp_dir is not None:
+            self._temp_dir.cleanup()
+            self._temp_dir = None
+        self._started = False
+        pass
+
+    def _to_config(self) -> LocalCommandLineCodeExecutorConfig:
+        if self._functions:
+            logging.info("Functions will not be included in serialized configuration")
+        if self._virtual_env_context:
+            logging.info("Virtual environment context will not be included in serialized configuration")
+
+        return LocalCommandLineCodeExecutorConfig(
+            timeout=self._timeout,
+            work_dir=str(self.work_dir),
+            functions_module=self._functions_module,
+            cleanup_temp_files=self._cleanup_temp_files,
+        )
+
+    @classmethod
+    def _from_config(cls, config: LocalCommandLineCodeExecutorConfig) -> Self:
+        return cls(
+            timeout=config.timeout,
+            work_dir=Path(config.work_dir) if config.work_dir is not None else None,
+            functions_module=config.functions_module,
+            cleanup_temp_files=config.cleanup_temp_files,
+        )