From e6423bb8628833e20cb80bcaf5f3b2cc4b883f26 Mon Sep 17 00:00:00 2001 From: Victor Dibia Date: Thu, 13 Feb 2025 13:50:11 -0800 Subject: [PATCH] Make CodeExecutor Serializable/Declarative (#5527) ## Why are these changes needed? Make [PythonCodeExecutionTool](https://github.com/microsoft/autogen/blob/main/python/packages/autogen-ext/src/autogen_ext/tools/code_execution/_code_execution.py) declarative so it can be used in tools like AGS Summary of changes - Make CodeExecutor declarative (convert from Protocol to ABC, inherit from ComponentBase) - Make LocalCommandLineCodeExecutor, JupyterCodeExecutor and DockerCommandLineCodeExecutor declarative , best effort. Not all fields are serialized, warnings are shown where appropriate. - Make PythonCodeExecutionTool declarative. ## Related issue number Closes #5526 ## Checks - [ ] I've included any doc changes needed for https://microsoft.github.io/autogen/. See https://microsoft.github.io/autogen/docs/Contribute#documentation to build and test documentation locally. - [ ] I've added tests (if relevant) corresponding to the changes introduced in this PR. - [ ] I've made sure all auto checks have passed. --- .../src/autogen_core/code_executor/_base.py | 13 +++- .../docker/_docker_code_executor.py | 63 ++++++++++++++++++- .../jupyter/_jupyter_code_executor.py | 28 ++++++++- .../code_executors/local/__init__.py | 38 ++++++++++- .../tools/code_execution/_code_execution.py | 27 +++++++- .../test_commandline_code_executor.py | 8 +++ .../test_docker_commandline_code_executor.py | 9 +++ .../test_jupyter_code_executor.py | 9 +++ .../tools/test_python_code_executor_tool.py | 59 +++++++++++++++++ 9 files changed, 243 insertions(+), 11 deletions(-) create mode 100644 python/packages/autogen-ext/tests/tools/test_python_code_executor_tool.py diff --git a/python/packages/autogen-core/src/autogen_core/code_executor/_base.py b/python/packages/autogen-core/src/autogen_core/code_executor/_base.py index 16f023a2191f..661751908c66 100644 --- a/python/packages/autogen-core/src/autogen_core/code_executor/_base.py +++ b/python/packages/autogen-core/src/autogen_core/code_executor/_base.py @@ -3,10 +3,14 @@ from __future__ import annotations +from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import List, Protocol, runtime_checkable +from typing import List + +from pydantic import BaseModel from .._cancellation_token import CancellationToken +from .._component_config import ComponentBase @dataclass @@ -25,10 +29,12 @@ class CodeResult: output: str -@runtime_checkable -class CodeExecutor(Protocol): +class CodeExecutor(ABC, ComponentBase[BaseModel]): """Executes code blocks and returns the result.""" + component_type = "code_executor" + + @abstractmethod async def execute_code_blocks( self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken ) -> CodeResult: @@ -49,6 +55,7 @@ async def execute_code_blocks( """ ... + @abstractmethod async def restart(self) -> None: """Restart the code executor. diff --git a/python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py b/python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py index 1ac058a9680f..d52b08bb2505 100644 --- a/python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py +++ b/python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py @@ -14,13 +14,15 @@ from types import TracebackType from typing import Any, Callable, ClassVar, Dict, List, Optional, ParamSpec, Type, Union -from autogen_core import CancellationToken +from autogen_core import CancellationToken, Component from autogen_core.code_executor import ( CodeBlock, CodeExecutor, FunctionWithRequirements, FunctionWithRequirementsStr, ) +from pydantic import BaseModel +from typing_extensions import Self from .._common import ( CommandLineCodeResult, @@ -50,7 +52,23 @@ async def _wait_for_ready(container: Any, timeout: int = 60, stop_time: float = A = ParamSpec("A") -class DockerCommandLineCodeExecutor(CodeExecutor): +class DockerCommandLineCodeExecutorConfig(BaseModel): + """Configuration for DockerCommandLineCodeExecutor""" + + image: str = "python:3-slim" + container_name: Optional[str] = None + timeout: int = 60 + work_dir: str = "." # Stored as string, converted to Path + bind_dir: Optional[str] = None # Stored as string, converted to Path + 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 + + +class DockerCommandLineCodeExecutor(CodeExecutor, Component[DockerCommandLineCodeExecutorConfig]): """Executes code through a command line environment in a Docker container. .. note:: @@ -97,6 +115,9 @@ class DockerCommandLineCodeExecutor(CodeExecutor): Example: init_command="kubectl config use-context docker-hub" """ + component_config_schema = DockerCommandLineCodeExecutorConfig + component_provider_override = "autogen_ext.code_executors.docker.DockerCommandLineCodeExecutor" + SUPPORTED_LANGUAGES: ClassVar[List[str]] = [ "bash", "shell", @@ -412,3 +433,41 @@ async def __aexit__( ) -> Optional[bool]: await self.stop() return None + + 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), + 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, + ) + + @classmethod + def _from_config(cls, config: DockerCommandLineCodeExecutorConfig) -> Self: + """(Experimental) Create a component from a config object.""" + bind_dir = Path(config.bind_dir) if config.bind_dir else None + return cls( + image=config.image, + container_name=config.container_name, + timeout=config.timeout, + work_dir=Path(config.work_dir), + bind_dir=bind_dir, + 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, + ) diff --git a/python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_code_executor.py b/python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_code_executor.py index bdeb0a78db9b..6da6b18ff501 100644 --- a/python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_code_executor.py +++ b/python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_code_executor.py @@ -8,6 +8,9 @@ from pathlib import Path from types import TracebackType +from autogen_core import Component +from pydantic import BaseModel + if sys.version_info >= (3, 11): from typing import Self else: @@ -18,6 +21,7 @@ from nbclient import NotebookClient from nbformat import NotebookNode from nbformat import v4 as nbformat +from typing_extensions import Self from .._common import silence_pip @@ -29,7 +33,15 @@ class JupyterCodeResult(CodeResult): output_files: list[Path] -class JupyterCodeExecutor(CodeExecutor): +class JupyterCodeExecutorConfig(BaseModel): + """Configuration for JupyterCodeExecutor""" + + kernel_name: str = "python3" + timeout: int = 60 + output_dir: str = "." + + +class JupyterCodeExecutor(CodeExecutor, Component[JupyterCodeExecutorConfig]): """A code executor class that executes code statefully using [nbclient](https://github.com/jupyter/nbclient). .. danger:: @@ -113,6 +125,9 @@ async def main() -> None: output_dir (Path): The directory to save output files, by default ".". """ + component_config_schema = JupyterCodeExecutorConfig + component_provider_override = "autogen_ext.code_executors.jupyter.JupyterCodeExecutor" + def __init__( self, kernel_name: str = "python3", @@ -261,3 +276,14 @@ async def __aexit__( exc_tb: TracebackType | None, ) -> None: await self.stop() + + 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) + ) + + @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)) diff --git a/python/packages/autogen-ext/src/autogen_ext/code_executors/local/__init__.py b/python/packages/autogen-ext/src/autogen_ext/code_executors/local/__init__.py index b538b373eb61..1513f723e9c2 100644 --- a/python/packages/autogen-ext/src/autogen_ext/code_executors/local/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/code_executors/local/__init__.py @@ -12,9 +12,10 @@ from types import SimpleNamespace from typing import Any, Callable, ClassVar, List, Optional, Sequence, Union -from autogen_core import CancellationToken +from autogen_core import CancellationToken, Component from autogen_core.code_executor import CodeBlock, CodeExecutor, FunctionWithRequirements, FunctionWithRequirementsStr -from typing_extensions import ParamSpec +from pydantic import BaseModel +from typing_extensions import ParamSpec, Self from .._common import ( PYTHON_VARIANTS, @@ -31,7 +32,15 @@ A = ParamSpec("A") -class LocalCommandLineCodeExecutor(CodeExecutor): +class LocalCommandLineCodeExecutorConfig(BaseModel): + """Configuration for LocalCommandLineCodeExecutor""" + + timeout: int = 60 + work_dir: str = "." # Stored as string, converted to Path in _from_config + functions_module: str = "functions" + + +class LocalCommandLineCodeExecutor(CodeExecutor, Component[LocalCommandLineCodeExecutorConfig]): """A code executor class that executes code through a local command line environment. @@ -97,6 +106,9 @@ async def example(): """ + component_config_schema = LocalCommandLineCodeExecutorConfig + component_provider_override = "autogen_ext.code_executors.local.LocalCommandLineCodeExecutor" + SUPPORTED_LANGUAGES: ClassVar[List[str]] = [ "bash", "shell", @@ -357,3 +369,23 @@ async def restart(self) -> None: "Restarting local command line code executor is not supported. No action is taken.", stacklevel=2, ) + + 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, + ) + + @classmethod + def _from_config(cls, config: LocalCommandLineCodeExecutorConfig) -> Self: + return cls( + timeout=config.timeout, + work_dir=Path(config.work_dir), + functions_module=config.functions_module, + ) diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/code_execution/_code_execution.py b/python/packages/autogen-ext/src/autogen_ext/tools/code_execution/_code_execution.py index 3b72940f2445..14086b63b513 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/code_execution/_code_execution.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/code_execution/_code_execution.py @@ -1,7 +1,8 @@ -from autogen_core import CancellationToken +from autogen_core import CancellationToken, Component, ComponentModel from autogen_core.code_executor import CodeBlock, CodeExecutor from autogen_core.tools import BaseTool from pydantic import BaseModel, Field, model_serializer +from typing_extensions import Self class CodeExecutionInput(BaseModel): @@ -17,7 +18,16 @@ def ser_model(self) -> str: return self.output -class PythonCodeExecutionTool(BaseTool[CodeExecutionInput, CodeExecutionResult]): +class PythonCodeExecutionToolConfig(BaseModel): + """Configuration for PythonCodeExecutionTool""" + + executor: ComponentModel + description: str = "Execute Python code blocks." + + +class PythonCodeExecutionTool( + BaseTool[CodeExecutionInput, CodeExecutionResult], Component[PythonCodeExecutionToolConfig] +): """A tool that executes Python code in a code executor and returns output. Example executors: @@ -61,6 +71,9 @@ async def main() -> None: executor (CodeExecutor): The code executor that will be used to execute the code blocks. """ + component_config_schema = PythonCodeExecutionToolConfig + component_provider_override = "autogen_ext.tools.code_execution.PythonCodeExecutionTool" + def __init__(self, executor: CodeExecutor): super().__init__(CodeExecutionInput, CodeExecutionResult, "CodeExecutor", "Execute Python code blocks.") self._executor = executor @@ -72,3 +85,13 @@ async def run(self, args: CodeExecutionInput, cancellation_token: CancellationTo ) return CodeExecutionResult(success=result.exit_code == 0, output=result.output) + + def _to_config(self) -> PythonCodeExecutionToolConfig: + """Convert current instance to config object""" + return PythonCodeExecutionToolConfig(executor=self._executor.dump_component()) + + @classmethod + def _from_config(cls, config: PythonCodeExecutionToolConfig) -> Self: + """Create instance from config object""" + executor = CodeExecutor.load_component(config.executor) + return cls(executor=executor) diff --git a/python/packages/autogen-ext/tests/code_executors/test_commandline_code_executor.py b/python/packages/autogen-ext/tests/code_executors/test_commandline_code_executor.py index e5e4407f6bbc..9d27d39d96fb 100644 --- a/python/packages/autogen-ext/tests/code_executors/test_commandline_code_executor.py +++ b/python/packages/autogen-ext/tests/code_executors/test_commandline_code_executor.py @@ -195,3 +195,11 @@ async def test_local_executor_with_custom_venv_in_local_relative_path() -> None: finally: if os.path.isdir(relative_folder_path): shutil.rmtree(relative_folder_path) + + +def test_serialize_deserialize() -> None: + with tempfile.TemporaryDirectory() as temp_dir: + executor = LocalCommandLineCodeExecutor(work_dir=temp_dir) + executor_config = executor.dump_component() + loaded_executor = LocalCommandLineCodeExecutor.load_component(executor_config) + assert executor.work_dir == loaded_executor.work_dir diff --git a/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py b/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py index 6c65835d183d..c507be9ffbda 100644 --- a/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py +++ b/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py @@ -209,3 +209,12 @@ async def test_docker_commandline_code_executor_extra_args() -> None: code_result = await executor.execute_code_blocks(code_blocks, cancellation_token) assert code_result.exit_code == 0 assert "This is a test file." in code_result.output + + +@pytest.mark.asyncio +async def test_docker_commandline_code_executor_serialization() -> None: + with tempfile.TemporaryDirectory() as temp_dir: + executor = DockerCommandLineCodeExecutor(work_dir=temp_dir) + loaded_executor = DockerCommandLineCodeExecutor.load_component(executor.dump_component()) + assert executor.bind_dir == loaded_executor.bind_dir + assert executor.timeout == loaded_executor.timeout diff --git a/python/packages/autogen-ext/tests/code_executors/test_jupyter_code_executor.py b/python/packages/autogen-ext/tests/code_executors/test_jupyter_code_executor.py index e8ddace5ec57..2ff0f1cb25aa 100644 --- a/python/packages/autogen-ext/tests/code_executors/test_jupyter_code_executor.py +++ b/python/packages/autogen-ext/tests/code_executors/test_jupyter_code_executor.py @@ -167,3 +167,12 @@ async def test_execute_code_with_html_output(tmp_path: Path) -> None: output_files=code_result.output_files, ) assert code_result.output_files[0].parent == tmp_path + + +@pytest.mark.asyncio +async def test_jupyter_code_executor_serialization(tmp_path: Path) -> None: + executor = JupyterCodeExecutor(output_dir=tmp_path) + serialized = executor.dump_component() + loaded_executor = JupyterCodeExecutor.load_component(serialized) + + assert isinstance(loaded_executor, JupyterCodeExecutor) diff --git a/python/packages/autogen-ext/tests/tools/test_python_code_executor_tool.py b/python/packages/autogen-ext/tests/tools/test_python_code_executor_tool.py new file mode 100644 index 000000000000..2593a4475426 --- /dev/null +++ b/python/packages/autogen-ext/tests/tools/test_python_code_executor_tool.py @@ -0,0 +1,59 @@ +import tempfile + +import pytest +from autogen_core import CancellationToken +from autogen_ext.code_executors.local import LocalCommandLineCodeExecutor +from autogen_ext.tools.code_execution import CodeExecutionInput, PythonCodeExecutionTool + + +@pytest.mark.asyncio +async def test_python_code_execution_tool() -> None: + """Test basic functionality of PythonCodeExecutionTool.""" + # Create a temporary directory for the executor + with tempfile.TemporaryDirectory() as temp_dir: + # Initialize the executor and tool + executor = LocalCommandLineCodeExecutor(work_dir=temp_dir) + tool = PythonCodeExecutionTool(executor=executor) + + # Test simple code execution + code = "print('hello world!')" + result = await tool.run(args=CodeExecutionInput(code=code), cancellation_token=CancellationToken()) + + # Verify successful execution + assert result.success is True + assert "hello world!" in result.output + + # Test code with computation + code = """a = 100 + 200 \nprint(f'Result: {a}') + """ + result = await tool.run(args=CodeExecutionInput(code=code), cancellation_token=CancellationToken()) + + # Verify computation result + assert result.success is True + assert "Result: 300" in result.output + + # Test error handling + code = "print(undefined_variable)" + result = await tool.run(args=CodeExecutionInput(code=code), cancellation_token=CancellationToken()) + + # Verify error handling + assert result.success is False + assert "NameError" in result.output + + +def test_python_code_execution_tool_serialization() -> None: + """Test serialization and deserialization of PythonCodeExecutionTool.""" + with tempfile.TemporaryDirectory() as temp_dir: + # Create original tool + executor = LocalCommandLineCodeExecutor(work_dir=temp_dir) + original_tool = PythonCodeExecutionTool(executor=executor) + + # Serialize + config = original_tool.dump_component() + assert config.config.get("executor") is not None + + # Deserialize + loaded_tool = PythonCodeExecutionTool.load_component(config) + + # Verify the loaded tool has the same configuration + assert isinstance(loaded_tool, PythonCodeExecutionTool)