Make CodeExecutor Serializable/Declarative (#5527)

<!-- Thank you for your contribution! Please review
https://microsoft.github.io/autogen/docs/Contribute before opening a
pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR.
If you don't have the access to it, we will shortly find a reviewer and
assign them to your PR. -->

## 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.

<!-- Please give a short summary of the change and the problem this
solves. -->

## Related issue number

Closes #5526 
<!-- For example: "Closes #1234" -->

## 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.

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.
 

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,
+        )

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))

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,
+        )

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)

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

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

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)