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

autogen/coding/base.py

@@ -0,0 +1,119 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Portions derived from https://github.com/microsoft/autogen are under the MIT License.
+# SPDX-License-Identifier: MIT
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any, Literal, Optional, Protocol, TypedDict, Union, runtime_checkable
+
+from pydantic import BaseModel, Field
+
+from ..doc_utils import export_module
+from ..types import UserMessageImageContentPart, UserMessageTextContentPart
+
+__all__ = ("CodeBlock", "CodeExecutionConfig", "CodeExecutor", "CodeExtractor", "CodeResult")
+
+
+@export_module("autogen.coding")
+class CodeBlock(BaseModel):
+    """(Experimental) A class that represents a code block."""
+
+    code: str = Field(description="The code to execute.")
+
+    language: str = Field(description="The language of the code.")
+
+
+@export_module("autogen.coding")
+class CodeResult(BaseModel):
+    """(Experimental) A class that represents the result of a code execution."""
+
+    exit_code: int = Field(description="The exit code of the code execution.")
+
+    output: str = Field(description="The output of the code execution.")
+
+
+@export_module("autogen.coding")
+class CodeExtractor(Protocol):
+    """(Experimental) A code extractor class that extracts code blocks from a message."""
+
+    def extract_code_blocks(
+        self, message: Optional[Union[str, list[Union[UserMessageTextContentPart, UserMessageImageContentPart]]]]
+    ) -> list[CodeBlock]:
+        """(Experimental) Extract code blocks from a message.
+
+        Args:
+            message (str): The message to extract code blocks from.
+
+        Returns:
+            List[CodeBlock]: The extracted code blocks.
+        """
+        ...  # pragma: no cover
+
+
+@runtime_checkable
+@export_module("autogen.coding")
+class CodeExecutor(Protocol):
+    """(Experimental) A code executor class that executes code blocks and returns the result."""
+
+    @property
+    def code_extractor(self) -> CodeExtractor:
+        """(Experimental) The code extractor used by this code executor."""
+        ...  # pragma: no cover
+
+    def execute_code_blocks(self, code_blocks: list[CodeBlock]) -> CodeResult:
+        """(Experimental) Execute code blocks and return the result.
+
+        This method should be implemented by the code executor.
+
+        Args:
+            code_blocks (List[CodeBlock]): The code blocks to execute.
+
+        Returns:
+            CodeResult: The result of the code execution.
+        """
+        ...  # pragma: no cover
+
+    def restart(self) -> None:
+        """(Experimental) Restart the code executor.
+
+        This method should be implemented by the code executor.
+
+        This method is called when the agent is reset.
+        """
+        ...  # pragma: no cover
+
+
+class IPythonCodeResult(CodeResult):
+    """(Experimental) A code result class for IPython code executor."""
+
+    output_files: list[str] = Field(
+        default_factory=list,
+        description="The list of files that the executed code blocks generated.",
+    )
+
+
+CodeExecutionConfig = TypedDict(
+    "CodeExecutionConfig",
+    {
+        "executor": Union[Literal["ipython-embedded", "commandline-local"], CodeExecutor],
+        "last_n_messages": Union[int, Literal["auto"]],
+        "timeout": int,
+        "use_docker": Union[bool, str, list[str]],
+        "work_dir": str,
+        "ipython-embedded": Mapping[str, Any],
+        "commandline-local": Mapping[str, Any],
+    },
+    total=False,
+)
+
+
+class CommandLineCodeResult(CodeResult):
+    """(Experimental) A code result class for command line code executor."""
+
+    code_file: Optional[str] = Field(
+        default=None,
+        description="The file that the executed code block was saved to.",
+    )

autogen/coding/docker_commandline_code_executor.py

@@ -0,0 +1,268 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Portions derived from https://github.com/microsoft/autogen are under the MIT License.
+# SPDX-License-Identifier: MIT
+from __future__ import annotations
+
+import atexit
+import logging
+import sys
+import uuid
+from hashlib import md5
+from pathlib import Path
+from time import sleep
+from types import TracebackType
+from typing import Any, ClassVar, Optional, Union
+
+import docker
+from docker.errors import ImageNotFound
+
+from ..code_utils import TIMEOUT_MSG, _cmd
+from ..doc_utils import export_module
+from .base import CodeBlock, CodeExecutor, CodeExtractor, CommandLineCodeResult
+from .markdown_code_extractor import MarkdownCodeExtractor
+from .utils import _get_file_name_from_content, silence_pip
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+
+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:
+        sleep(stop_time)
+        elapsed_time += stop_time
+        container.reload()
+        continue
+    if container.status != "running":
+        raise ValueError("Container failed to start")
+
+
+__all__ = ("DockerCommandLineCodeExecutor",)
+
+
+@export_module("autogen.coding")
+class DockerCommandLineCodeExecutor(CodeExecutor):
+    DEFAULT_EXECUTION_POLICY: ClassVar[dict[str, bool]] = {
+        "bash": True,
+        "shell": True,
+        "sh": True,
+        "pwsh": True,
+        "powershell": True,
+        "ps1": True,
+        "python": True,
+        "javascript": False,
+        "html": False,
+        "css": False,
+    }
+    LANGUAGE_ALIASES: ClassVar[dict[str, str]] = {"py": "python", "js": "javascript"}
+
+    def __init__(
+        self,
+        image: str = "python:3-slim",
+        container_name: Optional[str] = None,
+        timeout: int = 60,
+        work_dir: Optional[Union[Path, str]] = None,
+        bind_dir: Optional[Union[Path, str]] = None,
+        auto_remove: bool = True,
+        stop_container: bool = True,
+        execution_policies: Optional[dict[str, bool]] = None,
+    ):
+        """(Experimental) A code executor class that executes code through
+        a command line environment in a Docker container.
+
+        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", or "sh" for the code
+        block.
+
+        Args:
+            image: Docker image to use for code execution. Defaults to "python:3-slim".
+            container_name: Name of the Docker container which is created. If None, will autogenerate a name. Defaults to None.
+            timeout: The timeout for code execution. Defaults to 60.
+            work_dir: The working directory for the code execution. Defaults to Path(".").
+            bind_dir: 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: If true, will automatically remove the Docker container when it is stopped. Defaults to True.
+            stop_container: 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.
+            execution_policies: A dictionary mapping language names to boolean values that determine
+                whether code in that language should be executed. True means code in that language
+                will be executed, False means it will only be saved to a file. This overrides the
+                default execution policies. Defaults to None.
+
+        Raises:
+            ValueError: On argument error, or if the container fails to start.
+        """
+        work_dir = work_dir if work_dir is not None else Path()
+
+        if timeout < 1:
+            raise ValueError("Timeout must be greater than or equal to 1.")
+
+        if isinstance(work_dir, str):
+            work_dir = Path(work_dir)
+        work_dir.mkdir(exist_ok=True)
+
+        if bind_dir is None:
+            bind_dir = work_dir
+        elif isinstance(bind_dir, str):
+            bind_dir = Path(bind_dir)
+
+        client = docker.from_env()
+        # Check if the image exists
+        try:
+            client.images.get(image)
+        except ImageNotFound:
+            logging.info(f"Pulling image {image}...")
+            # Let the docker exception escape if this fails.
+            client.images.pull(image)
+
+        if container_name is None:
+            container_name = f"autogen-code-exec-{uuid.uuid4()}"
+
+        # Start a container from the image, read to exec commands later
+        self._container = client.containers.create(
+            image,
+            name=container_name,
+            entrypoint="/bin/sh",
+            tty=True,
+            auto_remove=auto_remove,
+            volumes={str(bind_dir.resolve()): {"bind": "/workspace", "mode": "rw"}},
+            working_dir="/workspace",
+        )
+        self._container.start()
+
+        _wait_for_ready(self._container)
+
+        def cleanup() -> None:
+            try:
+                container = client.containers.get(container_name)
+                container.stop()
+            except docker.errors.NotFound:
+                pass
+            atexit.unregister(cleanup)
+
+        if stop_container:
+            atexit.register(cleanup)
+
+        self._cleanup = cleanup
+
+        # Check if the container is running
+        if self._container.status != "running":
+            raise ValueError(f"Failed to start container from image {image}. Logs: {self._container.logs()}")
+
+        self._timeout = timeout
+        self._work_dir: Path = work_dir
+        self._bind_dir: Path = bind_dir
+        self.execution_policies = self.DEFAULT_EXECUTION_POLICY.copy()
+        if execution_policies is not None:
+            self.execution_policies.update(execution_policies)
+
+    @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."""
+        return self._work_dir
+
+    @property
+    def bind_dir(self) -> Path:
+        """(Experimental) The binding directory for the code execution container."""
+        return self._bind_dir
+
+    @property
+    def code_extractor(self) -> CodeExtractor:
+        """(Experimental) Export a code extractor that can be used by an agent."""
+        return MarkdownCodeExtractor()
+
+    def execute_code_blocks(self, code_blocks: list[CodeBlock]) -> 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 len(code_blocks) == 0:
+            raise ValueError("No code blocks to execute.")
+
+        outputs = []
+        files = []
+        last_exit_code = 0
+        for code_block in code_blocks:
+            lang = self.LANGUAGE_ALIASES.get(code_block.language.lower(), code_block.language.lower())
+            if lang not in self.DEFAULT_EXECUTION_POLICY:
+                outputs.append(f"Unsupported language {lang}\n")
+                last_exit_code = 1
+                break
+
+            execute_code = self.execution_policies.get(lang, False)
+            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_{md5(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)
+
+            if not execute_code:
+                outputs.append(f"Code saved to {code_path!s}\n")
+                continue
+
+            command = ["timeout", str(self._timeout), _cmd(lang), filename]
+            result = self._container.exec_run(command)
+            exit_code = result.exit_code
+            output = result.output.decode("utf-8")
+            if exit_code == 124:
+                output += "\n" + TIMEOUT_MSG
+            outputs.append(output)
+
+            last_exit_code = exit_code
+            if exit_code != 0:
+                break
+
+        code_file = str(files[0]) if files else None
+        return CommandLineCodeResult(exit_code=last_exit_code, output="".join(outputs), code_file=code_file)
+
+    def restart(self) -> None:
+        """(Experimental) Restart the code executor."""
+        self._container.restart()
+        if self._container.status != "running":
+            raise ValueError(f"Failed to restart container. Logs: {self._container.logs()}")
+
+    def stop(self) -> None:
+        """(Experimental) Stop the code executor."""
+        self._cleanup()
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        self.stop()

autogen/coding/factory.py

@@ -0,0 +1,47 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Portions derived from https://github.com/microsoft/autogen are under the MIT License.
+# SPDX-License-Identifier: MIT
+from ..doc_utils import export_module
+from .base import CodeExecutionConfig, CodeExecutor
+
+__all__ = ("CodeExecutorFactory",)
+
+
+@export_module("autogen.coding")
+class CodeExecutorFactory:
+    """(Experimental) A factory class for creating code executors."""
+
+    @staticmethod
+    def create(code_execution_config: CodeExecutionConfig) -> CodeExecutor:
+        """(Experimental) Get a code executor based on the code execution config.
+
+        Args:
+            code_execution_config (Dict): The code execution config,
+                which is a dictionary that must contain the key "executor".
+                The value of the key "executor" can be either a string
+                or an instance of CodeExecutor, in which case the code
+                executor is returned directly.
+
+        Returns:
+            CodeExecutor: The code executor.
+
+        Raises:
+            ValueError: If the code executor is unknown or not specified.
+        """
+        executor = code_execution_config.get("executor")
+        if isinstance(executor, CodeExecutor):
+            # If the executor is already an instance of CodeExecutor, return it.
+            return executor
+        if executor == "ipython-embedded":
+            from .jupyter.embedded_ipython_code_executor import EmbeddedIPythonCodeExecutor
+
+            return EmbeddedIPythonCodeExecutor(**code_execution_config.get("ipython-embedded", {}))
+        elif executor == "commandline-local":
+            from .local_commandline_code_executor import LocalCommandLineCodeExecutor
+
+            return LocalCommandLineCodeExecutor(**code_execution_config.get("commandline-local", {}))
+        else:
+            raise ValueError(f"Unknown code executor {executor}")

autogen/coding/func_with_reqs.py

@@ -0,0 +1,202 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Portions derived from https://github.com/microsoft/autogen are under the MIT License.
+# SPDX-License-Identifier: MIT
+
+import functools
+import importlib
+import inspect
+from dataclasses import dataclass, field
+from importlib.abc import SourceLoader
+from textwrap import dedent, indent
+from typing import Any, Callable, Generic, TypeVar, Union
+
+from typing_extensions import ParamSpec
+
+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
+
+
+@dataclass
+class Alias:
+    name: str
+    alias: str
+
+
+@dataclass
+class ImportFromModule:
+    module: str
+    imports: list[Union[str, Alias]]
+
+
+Import = Union[str, ImportFromModule, Alias]
+
+
+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}"
+
+
+class _StringLoader(SourceLoader):
+    def __init__(self, data: str):
+        self.data = data
+
+    def get_source(self, fullname: str) -> str:
+        return self.data
+
+    def get_data(self, path: str) -> bytes:
+        return self.data.encode("utf-8")
+
+    def get_filename(self, fullname: str) -> str:
+        return "<not a real path>/" + fullname + ".py"
+
+
+@dataclass
+class FunctionWithRequirementsStr:
+    func: str
+    _compiled_func: Callable[..., Any]
+    _func_name: str
+    python_packages: list[str] = field(default_factory=list)
+    global_imports: list[Import] = field(default_factory=list)
+
+    def __init__(self, func: str, python_packages: list[str] = [], global_imports: list[Import] = []):
+        self.func = func
+        self.python_packages = python_packages
+        self.global_imports = global_imports
+
+        module_name = "func_module"
+        loader = _StringLoader(func)
+        spec = importlib.util.spec_from_loader(module_name, loader)
+        if spec is None:
+            raise ValueError("Could not create spec")
+        module = importlib.util.module_from_spec(spec)
+        if spec.loader is None:
+            raise ValueError("Could not create loader")
+
+        try:
+            spec.loader.exec_module(module)
+        except Exception as e:
+            raise ValueError(f"Could not compile function: {e}") from e
+
+        functions = inspect.getmembers(module, inspect.isfunction)
+        if len(functions) != 1:
+            raise ValueError("The string must contain exactly one function")
+
+        self._func_name, self._compiled_func = functions[0]
+
+    def __call__(self, *args: Any, **kwargs: Any) -> None:
+        raise NotImplementedError("String based function with requirement objects are not directly callable")
+
+
+@dataclass
+class FunctionWithRequirements(Generic[T, P]):
+    func: Callable[P, T]
+    python_packages: list[str] = field(default_factory=list)
+    global_imports: list[Import] = field(default_factory=list)
+
+    @classmethod
+    def from_callable(
+        cls, func: Callable[P, T], python_packages: list[str] = [], global_imports: list[Import] = []
+    ) -> "FunctionWithRequirements[T, P]":
+        return cls(python_packages=python_packages, global_imports=global_imports, func=func)
+
+    @staticmethod
+    def from_str(
+        func: str, python_packages: list[str] = [], global_imports: list[Import] = []
+    ) -> FunctionWithRequirementsStr:
+        return FunctionWithRequirementsStr(func=func, python_packages=python_packages, global_imports=global_imports)
+
+    # Type this based on F
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T:
+        return self.func(*args, **kwargs)
+
+
+def with_requirements(
+    python_packages: list[str] = [], global_imports: list[Import] = []
+) -> Callable[[Callable[P, T]], FunctionWithRequirements[T, P]]:
+    """Decorate a function with package and import requirements
+
+    Args:
+        python_packages (List[str], optional): Packages required to function. Can include version info.. Defaults to [].
+        global_imports (List[Import], optional): Required imports. Defaults to [].
+
+    Returns:
+        Callable[[Callable[P, T]], FunctionWithRequirements[T, P]]: The decorated function
+    """
+
+    def wrapper(func: Callable[P, T]) -> FunctionWithRequirements[T, P]:
+        func_with_reqs = FunctionWithRequirements(
+            python_packages=python_packages, global_imports=global_imports, func=func
+        )
+
+        functools.update_wrapper(func_with_reqs, func)
+        return func_with_reqs
+
+    return wrapper
+
+
+def _build_python_functions_file(
+    funcs: list[Union[FunctionWithRequirements[Any, P], Callable[..., Any], FunctionWithRequirementsStr]],
+) -> str:
+    # First collect all global imports
+    global_imports: set[str] = set()
+    for func in funcs:
+        if isinstance(func, (FunctionWithRequirements, FunctionWithRequirementsStr)):
+            global_imports.update(map(_import_to_str, func.global_imports))
+
+    content = "\n".join(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

autogen/coding/jupyter/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Original portions of this file are derived from https://github.com/microsoft/autogen under the MIT License.
+# SPDX-License-Identifier: MIT
+
+from .base import JupyterConnectable, JupyterConnectionInfo
+from .docker_jupyter_server import DockerJupyterServer
+from .embedded_ipython_code_executor import EmbeddedIPythonCodeExecutor
+from .jupyter_client import JupyterClient
+from .jupyter_code_executor import JupyterCodeExecutor
+from .local_jupyter_server import LocalJupyterServer
+
+__all__ = [
+    "DockerJupyterServer",
+    "EmbeddedIPythonCodeExecutor",
+    "JupyterClient",
+    "JupyterCodeExecutor",
+    "JupyterConnectable",
+    "JupyterConnectionInfo",
+    "LocalJupyterServer",
+]

autogen/coding/jupyter/base.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Portions derived from https://github.com/microsoft/autogen are under the MIT License.
+# SPDX-License-Identifier: MIT
+from dataclasses import dataclass
+from typing import Optional, Protocol, runtime_checkable
+
+from ...doc_utils import export_module
+
+
+@dataclass
+@export_module("autogen.coding.jupyter")
+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
+@export_module("autogen.coding.jupyter")
+class JupyterConnectable(Protocol):
+    """(Experimental)"""
+
+    @property
+    def connection_info(self) -> JupyterConnectionInfo:
+        """Return the connection information for this connectable."""
+        pass