Silences Pip Install Messages in Code Executors (#2105)

* fix

* adds tests

* check if windows

* adds windows shells

* modifies exit code

* fix powershell

---------

Co-authored-by: Eric Zhu <[email protected]>

Author: WaelKarkoub

autogen/coding/docker_commandline_code_executor.py

@@ -10,7 +10,7 @@
 import docker
 from docker.errors import ImageNotFound
 
-from .utils import _get_file_name_from_content
+from .utils import _get_file_name_from_content, silence_pip
 from .base import CommandLineCodeResult
 
 from ..code_utils import TIMEOUT_MSG, _cmd
@@ -166,7 +166,7 @@ def execute_code_blocks(self, code_blocks: List[CodeBlock]) -> CommandLineCodeRe
         last_exit_code = 0
         for code_block in code_blocks:
             lang = code_block.language
-            code = code_block.code
+            code = silence_pip(code_block.code, lang)
 
             try:
                 # Check if there is a filename comment

autogen/coding/jupyter/jupyter_code_executor.py

@@ -8,6 +8,8 @@
 from typing import Any, ClassVar, List, Optional, Type, Union
 import sys
 
+from autogen.coding.utils import silence_pip
+
 if sys.version_info >= (3, 11):
     from typing import Self
 else:
@@ -91,7 +93,7 @@ def execute_code_blocks(self, code_blocks: List[CodeBlock]) -> IPythonCodeResult
         outputs = []
         output_files = []
         for code_block in code_blocks:
-            code = self._process_code(code_block.code)
+            code = silence_pip(code_block.code, code_block.language)
             result = self._jupyter_kernel_client.execute(code, timeout_seconds=self._timeout)
             if result.is_ok:
                 outputs.append(result.output)
@@ -140,18 +142,6 @@ def _save_html(self, html_data: str) -> str:
             f.write(html_data)
         return os.path.abspath(path)
 
-    def _process_code(self, code: str) -> str:
-        """Process code before execution."""
-        # Find lines that start with `! pip install` and make sure "-qqq" flag is added.
-        lines = code.split("\n")
-        for i, line in enumerate(lines):
-            # use regex to find lines that start with `! pip install` or `!pip install`.
-            match = re.search(r"^! ?pip install", line)
-            if match is not None:
-                if "-qqq" not in line:
-                    lines[i] = line.replace(match.group(0), match.group(0) + " -qqq")
-        return "\n".join(lines)
-
     def stop(self) -> None:
         """Stop the kernel."""
         self._jupyter_client.delete_kernel(self._kernel_id)

autogen/coding/local_commandline_code_executor.py

@@ -12,7 +12,7 @@
 from .base import CodeBlock, CodeExecutor, CodeExtractor, CommandLineCodeResult
 from .markdown_code_extractor import MarkdownCodeExtractor
 
-from .utils import _get_file_name_from_content
+from .utils import _get_file_name_from_content, silence_pip
 
 import subprocess
 
@@ -114,6 +114,7 @@ def execute_code_blocks(self, code_blocks: List[CodeBlock]) -> CommandLineCodeRe
             lang = lang.lower()
 
             LocalCommandLineCodeExecutor.sanitize_command(lang, code)
+            code = silence_pip(code, lang)
 
             if WIN32 and lang in ["sh", "shell"]:
                 lang = "ps1"

autogen/coding/utils.py

@@ -1,4 +1,5 @@
 # Will return the filename relative to the workspace path
+import re
 from pathlib import Path
 from typing import Optional
 
@@ -20,3 +21,23 @@ def _get_file_name_from_content(code: str, workspace_path: Path) -> Optional[str
         return str(relative)
 
     return None
+
+
+def silence_pip(code: str, lang: str) -> str:
+    """Apply -qqq flag to pip install commands."""
+    if lang == "python":
+        regex = r"^! ?pip install"
+    elif lang in ["bash", "shell", "sh", "pwsh", "powershell", "ps1"]:
+        regex = r"^pip install"
+    else:
+        return code
+
+    # Find lines that start with pip install and make sure "-qqq" flag is added.
+    lines = code.split("\n")
+    for i, line in enumerate(lines):
+        # use regex to find lines that start with pip install.
+        match = re.search(regex, line)
+        if match is not None:
+            if "-qqq" not in line:
+                lines[i] = line.replace(match.group(0), match.group(0) + " -qqq")
+    return "\n".join(lines)

test/coding/test_commandline_code_executor.py

@@ -1,22 +1,25 @@
 from pathlib import Path
 import sys
 import tempfile
+import uuid
 import pytest
 from autogen.agentchat.conversable_agent import ConversableAgent
 from autogen.code_utils import is_docker_running
 from autogen.coding.base import CodeBlock, CodeExecutor
 from autogen.coding.factory import CodeExecutorFactory
 from autogen.coding.docker_commandline_code_executor import DockerCommandLineCodeExecutor
 from autogen.coding.local_commandline_code_executor import LocalCommandLineCodeExecutor
-from autogen.oai.openai_utils import config_list_from_json
 
-from conftest import MOCK_OPEN_AI_API_KEY, skip_openai, skip_docker
+from conftest import MOCK_OPEN_AI_API_KEY, skip_docker
 
 if skip_docker or not is_docker_running():
     classes_to_test = [LocalCommandLineCodeExecutor]
 else:
     classes_to_test = [LocalCommandLineCodeExecutor, DockerCommandLineCodeExecutor]
 
+UNIX_SHELLS = ["bash", "sh", "shell"]
+WINDOWS_SHELLS = ["ps1", "pwsh", "powershell"]
+
 
 @pytest.mark.parametrize("cls", classes_to_test)
 def test_is_code_executor(cls) -> None:
@@ -218,3 +221,31 @@ def test_valid_relative_path(cls) -> None:
         assert "test.py" in result.code_file
         assert (temp_dir / "test.py").resolve() == Path(result.code_file).resolve()
         assert (temp_dir / "test.py").exists()
+
+
+@pytest.mark.parametrize("cls", classes_to_test)
+@pytest.mark.parametrize("lang", WINDOWS_SHELLS + UNIX_SHELLS)
+def test_silent_pip_install(cls, lang: str) -> None:
+    # Ensure that the shell is supported.
+    lang = "ps1" if lang in ["powershell", "pwsh"] else lang
+
+    if sys.platform in ["win32"] and lang in UNIX_SHELLS:
+        pytest.skip("Linux shells are not supported on Windows.")
+    elif sys.platform not in ["win32"] and lang in WINDOWS_SHELLS:
+        pytest.skip("Windows shells are not supported on Unix.")
+
+    error_exit_code = 0 if sys.platform in ["win32"] else 1
+
+    executor = cls(timeout=600)
+
+    code = "pip install matplotlib numpy"
+    code_blocks = [CodeBlock(code=code, language=lang)]
+    code_result = executor.execute_code_blocks(code_blocks)
+    assert code_result.exit_code == 0 and code_result.output.strip() == ""
+
+    none_existing_package = uuid.uuid4().hex
+
+    code = f"pip install matplotlib_{none_existing_package}"
+    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