diff --git a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/command-line-code-executors.ipynb b/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/command-line-code-executors.ipynb index a408cd09d..62de20d95 100644 --- a/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/command-line-code-executors.ipynb +++ b/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/command-line-code-executors.ipynb @@ -136,6 +136,76 @@ " )\n", ")" ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Local within a Virtual Environment\n", + "\n", + "If you want the code to run within a virtual environment created as part of the application’s setup, you can specify a directory for the newly created environment and pass its context to {py:class}`~autogen_core.components.code_executor.LocalCommandLineCodeExecutor`. This setup allows the executor to use the specified virtual environment consistently throughout the application's lifetime, ensuring isolated dependencies and a controlled runtime environment." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "CommandLineCodeResult(exit_code=0, output='', code_file='/Users/gziz/Dev/autogen/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/coding/tmp_code_d2a7db48799db3cc785156a11a38822a45c19f3956f02ec69b92e4169ecbf2ca.bash')" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "import venv\n", + "from pathlib import Path\n", + "\n", + "from autogen_core.base import CancellationToken\n", + "from autogen_core.components.code_executor import CodeBlock, LocalCommandLineCodeExecutor\n", + "\n", + "work_dir = Path(\"coding\")\n", + "work_dir.mkdir(exist_ok=True)\n", + "\n", + "venv_dir = work_dir / \".venv\"\n", + "venv_builder = venv.EnvBuilder(with_pip=True)\n", + "venv_builder.create(venv_dir)\n", + "venv_context = venv_builder.ensure_directories(venv_dir)\n", + "\n", + "local_executor = LocalCommandLineCodeExecutor(work_dir=work_dir, virtual_env_context=venv_context)\n", + "await local_executor.execute_code_blocks(\n", + " code_blocks=[\n", + " CodeBlock(language=\"bash\", code=\"pip install matplotlib\"),\n", + " ],\n", + " cancellation_token=CancellationToken(),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "As we can see, the code has executed successfully, and the installation has been isolated to the newly created virtual environment, without affecting our global environment." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { @@ -154,7 +224,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.9" + "version": "3.12.4" } }, "nbformat": 4, diff --git a/python/packages/autogen-core/src/autogen_core/components/code_executor/_impl/local_commandline_code_executor.py b/python/packages/autogen-core/src/autogen_core/components/code_executor/_impl/local_commandline_code_executor.py index f74111ef1..deca8355f 100644 --- a/python/packages/autogen-core/src/autogen_core/components/code_executor/_impl/local_commandline_code_executor.py +++ b/python/packages/autogen-core/src/autogen_core/components/code_executor/_impl/local_commandline_code_executor.py @@ -3,12 +3,14 @@ import asyncio import logging +import os import sys import warnings from hashlib import sha256 from pathlib import Path from string import Template -from typing import Any, Callable, ClassVar, List, Sequence, Union +from types import SimpleNamespace +from typing import Any, Callable, ClassVar, List, Optional, Sequence, Union from typing_extensions import ParamSpec @@ -54,6 +56,36 @@ class LocalCommandLineCodeExecutor(CodeExecutor): directory is the current 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". + virtual_env_context (Optional[SimpleNamespace], optional): The virtual environment context. Defaults to None. + + 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 + + from autogen_core.base import CancellationToken + from autogen_core.components.code_executor import CodeBlock, LocalCommandLineCodeExecutor + + 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(), + ) """ @@ -86,6 +118,7 @@ $functions""" ] ] = [], functions_module: str = "functions", + virtual_env_context: Optional[SimpleNamespace] = None, ): if timeout < 1: raise ValueError("Timeout must be greater than or equal to 1.") @@ -110,6 +143,8 @@ $functions""" else: self._setup_functions_complete = True + self._virtual_env_context: Optional[SimpleNamespace] = virtual_env_context + def format_functions_for_prompt(self, prompt_template: str = FUNCTION_PROMPT_TEMPLATE) -> str: """(Experimental) Format the functions for a prompt. @@ -164,9 +199,14 @@ $functions""" 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( - sys.executable, + py_executable, *cmd_args, cwd=self._work_dir, stdout=asyncio.subprocess.PIPE, @@ -253,7 +293,17 @@ $functions""" f.write(code) file_names.append(written_file) - program = sys.executable if lang.startswith("python") else lang_to_cmd(lang) + env = os.environ.copy() + + if self._virtual_env_context: + virtual_env_exe_abs_path = os.path.abspath(self._virtual_env_context.env_exe) + 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']}" + + program = virtual_env_exe_abs_path if lang.startswith("python") else lang_to_cmd(lang) + else: + program = sys.executable if lang.startswith("python") else lang_to_cmd(lang) + # Wrap in a task to make it cancellable task = asyncio.create_task( asyncio.create_subprocess_exec( @@ -262,6 +312,7 @@ $functions""" cwd=self._work_dir, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, + env=env, ) ) cancellation_token.link_future(task) diff --git a/python/packages/autogen-core/tests/execution/test_commandline_code_executor.py b/python/packages/autogen-core/tests/execution/test_commandline_code_executor.py index bb3ff2830..aff36b216 100644 --- a/python/packages/autogen-core/tests/execution/test_commandline_code_executor.py +++ b/python/packages/autogen-core/tests/execution/test_commandline_code_executor.py @@ -2,8 +2,11 @@ # Credit to original authors import asyncio +import os +import shutil import sys import tempfile +import venv from pathlib import Path from typing import AsyncGenerator, TypeAlias @@ -143,3 +146,51 @@ print("hello world") assert "test.py" in result.code_file assert (temp_dir / Path("test.py")).resolve() == Path(result.code_file).resolve() assert (temp_dir / Path("test.py")).exists() + + +@pytest.mark.asyncio +async def test_local_executor_with_custom_venv() -> None: + with tempfile.TemporaryDirectory() as temp_dir: + env_builder = venv.EnvBuilder(with_pip=True) + env_builder.create(temp_dir) + env_builder_context = env_builder.ensure_directories(temp_dir) + + executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, virtual_env_context=env_builder_context) + code_blocks = [ + # https://stackoverflow.com/questions/1871549/how-to-determine-if-python-is-running-inside-a-virtualenv + CodeBlock(code="import sys; print(sys.prefix != sys.base_prefix)", language="python"), + ] + cancellation_token = CancellationToken() + result = await executor.execute_code_blocks(code_blocks, cancellation_token=cancellation_token) + + assert result.exit_code == 0 + assert result.output.strip() == "True" + + +@pytest.mark.asyncio +async def test_local_executor_with_custom_venv_in_local_relative_path() -> None: + relative_folder_path = "tmp_dir" + try: + if not os.path.isdir(relative_folder_path): + os.mkdir(relative_folder_path) + + env_path = os.path.join(relative_folder_path, ".venv") + env_builder = venv.EnvBuilder(with_pip=True) + env_builder.create(env_path) + env_builder_context = env_builder.ensure_directories(env_path) + + executor = LocalCommandLineCodeExecutor(work_dir=relative_folder_path, virtual_env_context=env_builder_context) + code_blocks = [ + CodeBlock(code="import sys; print(sys.executable)", language="python"), + ] + cancellation_token = CancellationToken() + result = await executor.execute_code_blocks(code_blocks, cancellation_token=cancellation_token) + + assert result.exit_code == 0 + + # Check if the expected venv has been used + bin_path = os.path.abspath(env_builder_context.bin_path) + assert Path(result.output.strip()).parent.samefile(bin_path) + finally: + if os.path.isdir(relative_folder_path): + shutil.rmtree(relative_folder_path)