diff --git a/autogen/code_utils.py b/autogen/code_utils.py index 96c3f179779f..3a7c67c136d6 100644 --- a/autogen/code_utils.py +++ b/autogen/code_utils.py @@ -213,13 +213,37 @@ def timeout_handler(signum, frame): raise TimeoutError("Timed out!") +def get_powershell_command(): + try: + result = subprocess.run(["powershell", "$PSVersionTable.PSVersion.Major"], capture_output=True, text=True) + if result.returncode == 0: + return "powershell" + + except FileNotFoundError: + # This means that 'powershell' command is not found so now we try looking for 'pwsh' + try: + result = subprocess.run( + ["pwsh", "-Command", "$PSVersionTable.PSVersion.Major"], capture_output=True, text=True + ) + if result.returncode == 0: + return "pwsh" + + except FileNotFoundError: + print("Neither powershell nor pwsh is installed.") + return None + + +powershell_command = get_powershell_command() + + def _cmd(lang): - if lang.startswith("python") or lang in ["bash", "sh", "powershell"]: + if lang.startswith("python") or lang in ["bash", "sh", powershell_command]: return lang if lang in ["shell"]: return "sh" - if lang in ["ps1"]: - return "powershell" + if lang in ["ps1", "pwsh", "powershell"]: + return powershell_command + raise NotImplementedError(f"{lang} not recognized in code execution") diff --git a/autogen/coding/local_commandline_code_executor.py b/autogen/coding/local_commandline_code_executor.py index fe3dbf8d3184..25f4e5ebe4f5 100644 --- a/autogen/coding/local_commandline_code_executor.py +++ b/autogen/coding/local_commandline_code_executor.py @@ -156,7 +156,7 @@ def execute_code_blocks(self, code_blocks: List[CodeBlock]) -> CommandlineCodeRe ) filename_uuid = uuid.uuid4().hex filename = None - if lang in ["bash", "shell", "sh"]: + if lang in ["bash", "shell", "sh", "pwsh", "powershell", "ps1"]: filename = f"{filename_uuid}.{lang}" exitcode, logs, _ = execute_code( code=code, diff --git a/test/test_code_utils.py b/test/test_code_utils.py index 6a5a662ae43b..74efd9ddc222 100644 --- a/test/test_code_utils.py +++ b/test/test_code_utils.py @@ -1,8 +1,10 @@ import os import tempfile import unittest - +from unittest.mock import patch +import sys import pytest +from io import StringIO import autogen from autogen.code_utils import ( @@ -11,6 +13,7 @@ execute_code, extract_code, improve_code, + get_powershell_command, improve_function, infer_lang, is_docker_running, @@ -553,6 +556,41 @@ def test_non_dict_in_list(self): content_str(content) +class TestGetPowerShellCommand(unittest.TestCase): + @patch("subprocess.run") + def test_get_powershell_command_powershell(self, mock_subprocess_run): + # Set up the mock to return a successful result for 'powershell' + mock_subprocess_run.return_value.returncode = 0 + mock_subprocess_run.return_value.stdout = StringIO("5") + + self.assertEqual(get_powershell_command(), "powershell") + + @patch("subprocess.run") + def test_get_powershell_command_pwsh(self, mock_subprocess_run): + # Set up the mock to return a successful result for 'pwsh' + mock_subprocess_run.side_effect = [FileNotFoundError, mock_subprocess_run.return_value] + mock_subprocess_run.return_value.returncode = 0 + mock_subprocess_run.return_value.stdout = StringIO("7") + + self.assertEqual(get_powershell_command(), "pwsh") + + @patch("subprocess.run") + def test_get_powershell_command_no_shell(self, mock_subprocess_run): + # Set up the mock to simulate 'powershell' and 'pwsh' not found + mock_subprocess_run.side_effect = [FileNotFoundError, FileNotFoundError] + + with patch("sys.stdout", new=StringIO()) as fake_out: + get_powershell_command() + self.assertEqual(fake_out.getvalue().strip(), "Neither powershell nor pwsh is installed.") + + @patch("subprocess.run") + def test_get_powershell_command_no_shell_no_output(self, mock_subprocess_run): + # Set up the mock to simulate 'powershell' and 'pwsh' not found without printing error message + mock_subprocess_run.side_effect = [FileNotFoundError, FileNotFoundError] + + self.assertIsNone(get_powershell_command()) + + if __name__ == "__main__": # test_infer_lang() test_extract_code()