Feature: Add ability to use a separate python environment in local executor (#2615)

* Add ability to use virtual environments in local executor

* Copy environment variables from parent environment

* Fix mypy errors and formatting

* Account for venv on Windows

* Use a virtual environment context object instead of path

* Add utility method to create a virtual environment

* Remove assertion using `_venv_path`

* Add tests for `create_virtual_env`

* Modify test code and add output assertion

* Modify test code and assertion

* Execute activation script before actual command on windows

* Add docs for using a virtual env

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)

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

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"

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

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": {},