Feature: Add ability to use a separate python environment in local executor

autogen/code_utils.py

@@ -6,8 +6,10 @@
 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,11 +1,13 @@
 import logging
+import os
 import re
 import subprocess
 import sys
 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 @@ class LocalCommandLineCodeExecutor(CodeExecutor):
     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",
@@ -84,6 +87,7 @@ def __init__(
 
         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 +109,7 @@ def __init__(
 
         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 +201,11 @@ def _setup_functions(self) -> None:
         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 +278,18 @@ def _execute_code_dont_check_setup(self, code_blocks: List[CodeBlock]) -> Comman
 
             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 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 tempfile
 import unittest
 from io import StringIO
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import pytest
@@ -15,6 +16,7 @@
     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