diff --git a/autogen/code_utils.py b/autogen/code_utils.py index e1bc951f09..98ed606706 100644 --- a/autogen/code_utils.py +++ b/autogen/code_utils.py @@ -6,8 +6,10 @@ import string import subprocess import sys import time +import venv from concurrent.futures import ThreadPoolExecutor, TimeoutError from hashlib import md5 +from types import SimpleNamespace from typing import Any, Callable, Dict, List, Optional, Tuple, Union import docker @@ -719,3 +721,19 @@ def implement( # cost += metrics["gen_cost"] # if metrics["succeed_assertions"] or i == len(configs) - 1: # return responses[metrics["index_selected"]], cost, i + + +def create_virtual_env(dir_path: str, **env_args) -> SimpleNamespace: + """Creates a python virtual environment and returns the context. + + Args: + dir_path (str): Directory path where the env will be created. + **env_args: Any extra args to pass to the `EnvBuilder` + + Returns: + SimpleNamespace: the virtual env context object.""" + if not env_args: + env_args = {"with_pip": True} + env_builder = venv.EnvBuilder(**env_args) + env_builder.create(dir_path) + return env_builder.ensure_directories(dir_path) diff --git a/autogen/coding/local_commandline_code_executor.py b/autogen/coding/local_commandline_code_executor.py index ed92cd527b..29172bbe92 100644 --- a/autogen/coding/local_commandline_code_executor.py +++ b/autogen/coding/local_commandline_code_executor.py @@ -1,4 +1,5 @@ import logging +import os import re import subprocess import sys @@ -6,6 +7,7 @@ import warnings from hashlib import md5 from pathlib import Path from string import Template +from types import SimpleNamespace from typing import Any, Callable, ClassVar, Dict, List, Optional, Union from typing_extensions import ParamSpec @@ -64,6 +66,7 @@ $functions""" def __init__( self, timeout: int = 60, + virtual_env_context: Optional[SimpleNamespace] = None, work_dir: Union[Path, str] = Path("."), functions: List[Union[FunctionWithRequirements[Any, A], Callable[..., Any], FunctionWithRequirementsStr]] = [], functions_module: str = "functions", @@ -82,8 +85,22 @@ $functions""" PowerShell (pwsh, powershell, ps1), HTML, CSS, and JavaScript. Execution policies determine whether each language's code blocks are executed or saved only. + ## Execution with a Python virtual environment + A python virtual env can be used to execute code and install dependencies. This has the added benefit of not polluting the + base environment with unwanted modules. + ```python + from autogen.code_utils import create_virtual_env + from autogen.coding import LocalCommandLineCodeExecutor + + venv_dir = ".venv" + venv_context = create_virtual_env(venv_dir) + + executor = LocalCommandLineCodeExecutor(virtual_env_context=venv_context) + ``` + Args: timeout (int): The timeout for code execution, default is 60 seconds. + virtual_env_context (Optional[SimpleNamespace]): The virtual environment context to use. work_dir (Union[Path, str]): The working directory for code execution, defaults to the current directory. functions (List[Union[FunctionWithRequirements[Any, A], Callable[..., Any], FunctionWithRequirementsStr]]): A list of callable functions available to the executor. functions_module (str): The module name under which functions are accessible. @@ -105,6 +122,7 @@ $functions""" self._timeout = timeout self._work_dir: Path = work_dir + self._virtual_env_context: Optional[SimpleNamespace] = virtual_env_context self._functions = functions # Setup could take some time so we intentionally wait for the first code block to do it. @@ -196,7 +214,11 @@ $functions""" required_packages = list(set(flattened_packages)) if len(required_packages) > 0: logging.info("Ensuring packages are installed in executor.") - cmd = [sys.executable, "-m", "pip", "install"] + required_packages + if self._virtual_env_context: + py_executable = self._virtual_env_context.env_exe + else: + py_executable = sys.executable + cmd = [py_executable, "-m", "pip", "install"] + required_packages try: result = subprocess.run( cmd, cwd=self._work_dir, capture_output=True, text=True, timeout=float(self._timeout) @@ -269,9 +291,18 @@ $functions""" program = _cmd(lang) cmd = [program, str(written_file.absolute())] + env = os.environ.copy() + + if self._virtual_env_context: + path_with_virtualenv = rf"{self._virtual_env_context.bin_path}{os.pathsep}{env['PATH']}" + env["PATH"] = path_with_virtualenv + if WIN32: + activation_script = os.path.join(self._virtual_env_context.bin_path, "activate.bat") + cmd = [activation_script, "&&", *cmd] + try: result = subprocess.run( - cmd, cwd=self._work_dir, capture_output=True, text=True, timeout=float(self._timeout) + cmd, cwd=self._work_dir, capture_output=True, text=True, timeout=float(self._timeout), env=env ) except subprocess.TimeoutExpired: logs_all += "\n" + TIMEOUT_MSG diff --git a/test/coding/test_commandline_code_executor.py b/test/coding/test_commandline_code_executor.py index 0a0ded71e6..4daf2e21bc 100644 --- a/test/coding/test_commandline_code_executor.py +++ b/test/coding/test_commandline_code_executor.py @@ -2,12 +2,13 @@ import os import sys import tempfile import uuid +import venv from pathlib import Path import pytest from autogen.agentchat.conversable_agent import ConversableAgent -from autogen.code_utils import decide_use_docker, is_docker_running +from autogen.code_utils import WIN32, decide_use_docker, is_docker_running from autogen.coding.base import CodeBlock, CodeExecutor from autogen.coding.docker_commandline_code_executor import DockerCommandLineCodeExecutor from autogen.coding.factory import CodeExecutorFactory @@ -393,3 +394,20 @@ def test_silent_pip_install(cls, lang: str) -> None: code_blocks = [CodeBlock(code=code, language=lang)] code_result = executor.execute_code_blocks(code_blocks) assert code_result.exit_code == error_exit_code and "ERROR: " in code_result.output + + +def test_local_executor_with_custom_python_env(): + 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"), + ] + execution = executor.execute_code_blocks(code_blocks) + + assert execution.exit_code == 0 + assert execution.output.strip() == "True" diff --git a/test/test_code_utils.py b/test/test_code_utils.py index d6084f9b02..ade3855ad8 100755 --- a/test/test_code_utils.py +++ b/test/test_code_utils.py @@ -5,6 +5,7 @@ import sys import tempfile import unittest from io import StringIO +from types import SimpleNamespace from unittest.mock import patch import pytest @@ -15,6 +16,7 @@ from autogen.code_utils import ( UNKNOWN, check_can_use_docker_or_throw, content_str, + create_virtual_env, decide_use_docker, execute_code, extract_code, @@ -500,6 +502,20 @@ def test_can_use_docker_or_throw(): check_can_use_docker_or_throw(True) +def test_create_virtual_env(): + with tempfile.TemporaryDirectory() as temp_dir: + venv_context = create_virtual_env(temp_dir) + assert isinstance(venv_context, SimpleNamespace) + assert venv_context.env_name == os.path.split(temp_dir)[1] + + +def test_create_virtual_env_with_extra_args(): + with tempfile.TemporaryDirectory() as temp_dir: + venv_context = create_virtual_env(temp_dir, with_pip=False) + assert isinstance(venv_context, SimpleNamespace) + assert venv_context.env_name == os.path.split(temp_dir)[1] + + def _test_improve(): try: import openai diff --git a/website/docs/topics/code-execution/cli-code-executor.ipynb b/website/docs/topics/code-execution/cli-code-executor.ipynb index 69df79754d..11649b15a5 100644 --- a/website/docs/topics/code-execution/cli-code-executor.ipynb +++ b/website/docs/topics/code-execution/cli-code-executor.ipynb @@ -126,6 +126,35 @@ ")" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using a Python virtual environment\n", + "\n", + "By default, the LocalCommandLineCodeExecutor executes code and installs dependencies within the same Python environment as the AutoGen code. You have the option to specify a Python virtual environment to prevent polluting the base Python environment.\n", + "\n", + "### Example" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from autogen.code_utils import create_virtual_env\n", + "from autogen.coding import CodeBlock, LocalCommandLineCodeExecutor\n", + "\n", + "venv_dir = \".venv\"\n", + "venv_context = create_virtual_env(venv_dir)\n", + "\n", + "executor = LocalCommandLineCodeExecutor(virtual_env_context=venv_context)\n", + "print(\n", + " executor.execute_code_blocks(code_blocks=[CodeBlock(language=\"python\", code=\"import sys; print(sys.executable)\")])\n", + ")" + ] + }, { "cell_type": "markdown", "metadata": {},