Run LocalCommandLineCodeExecutor within venv (#3977)

* Run LocalCommandLineCodeExecutor within venv

* Remove create_virtual_env func and add docstring

* Add explanation for LocalCommandLineExecutor docstring example

* Enhance docstring example explanation

---------

Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com>

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,

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)

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)