diff --git a/autogen/agentchat/conversable_agent.py b/autogen/agentchat/conversable_agent.py index a49f8a5b43b..c7351d4287c 100644 --- a/autogen/agentchat/conversable_agent.py +++ b/autogen/agentchat/conversable_agent.py @@ -17,6 +17,7 @@ from .._pydantic import model_dump from ..cache.cache import AbstractCache from ..code_utils import ( + PYTHON_VARIANTS, UNKNOWN, check_can_use_docker_or_throw, content_str, @@ -2079,7 +2080,7 @@ def execute_code_blocks(self, code_blocks): ) if lang in ["bash", "shell", "sh"]: exitcode, logs, image = self.run_code(code, lang=lang, **self._code_execution_config) - elif lang in ["python", "Python"]: + elif lang in PYTHON_VARIANTS: if code.startswith("# filename: "): filename = code[11 : code.find("\n")].strip() else: diff --git a/autogen/code_utils.py b/autogen/code_utils.py index d4b2ae99cf0..aa75756e04a 100644 --- a/autogen/code_utils.py +++ b/autogen/code_utils.py @@ -35,6 +35,7 @@ DEFAULT_TIMEOUT = 600 WIN32 = sys.platform == "win32" PATH_SEPARATOR = WIN32 and "\\" or "/" +PYTHON_VARIANTS = ["python", "Python", "py"] logger = logging.getLogger(__name__) @@ -244,6 +245,8 @@ def get_powershell_command(): def _cmd(lang: str) -> str: + if lang in PYTHON_VARIANTS: + return "python" if lang.startswith("python") or lang in ["bash", "sh"]: return lang if lang in ["shell"]: diff --git a/autogen/coding/local_commandline_code_executor.py b/autogen/coding/local_commandline_code_executor.py index 44ba859265d..68ef76b7e7f 100644 --- a/autogen/coding/local_commandline_code_executor.py +++ b/autogen/coding/local_commandline_code_executor.py @@ -17,7 +17,7 @@ to_stub, ) -from ..code_utils import TIMEOUT_MSG, WIN32, _cmd +from ..code_utils import PYTHON_VARIANTS, TIMEOUT_MSG, WIN32, _cmd from .base import CodeBlock, CodeExecutor, CodeExtractor, CommandLineCodeResult from .markdown_code_extractor import MarkdownCodeExtractor from .utils import _get_file_name_from_content, silence_pip @@ -217,6 +217,9 @@ def _execute_code_dont_check_setup(self, code_blocks: List[CodeBlock]) -> Comman LocalCommandLineCodeExecutor.sanitize_command(lang, code) code = silence_pip(code, lang) + if lang in PYTHON_VARIANTS: + lang = "python" + if WIN32 and lang in ["sh", "shell"]: lang = "ps1" diff --git a/test/coding/test_commandline_code_executor.py b/test/coding/test_commandline_code_executor.py index 5c20e856d60..a83282dec78 100644 --- a/test/coding/test_commandline_code_executor.py +++ b/test/coding/test_commandline_code_executor.py @@ -23,6 +23,7 @@ UNIX_SHELLS = ["bash", "sh", "shell"] WINDOWS_SHELLS = ["ps1", "pwsh", "powershell"] +PYTHON_VARIANTS = ["python", "Python", "py"] @pytest.mark.parametrize("cls", classes_to_test) @@ -60,23 +61,26 @@ def test_commandline_executor_init(cls) -> None: executor = cls(timeout=111, work_dir="/invalid/directory") +@pytest.mark.parametrize("py_variant", PYTHON_VARIANTS) @pytest.mark.parametrize("cls", classes_to_test) -def test_commandline_executor_execute_code(cls) -> None: +def test_commandline_executor_execute_code(cls, py_variant) -> None: with tempfile.TemporaryDirectory() as temp_dir: executor = cls(work_dir=temp_dir) - _test_execute_code(executor=executor) + _test_execute_code(py_variant, executor=executor) -def _test_execute_code(executor: CodeExecutor) -> None: +@pytest.mark.parametrize("py_variant", PYTHON_VARIANTS) +def _test_execute_code(py_variant, executor: CodeExecutor) -> None: + # Test single code block. - code_blocks = [CodeBlock(code="import sys; print('hello world!')", language="python")] + code_blocks = [CodeBlock(code="import sys; print('hello world!')", language=py_variant)] code_result = executor.execute_code_blocks(code_blocks) assert code_result.exit_code == 0 and "hello world!" in code_result.output and code_result.code_file is not None # Test multiple code blocks. code_blocks = [ - CodeBlock(code="import sys; print('hello world!')", language="python"), - CodeBlock(code="a = 100 + 100; print(a)", language="python"), + CodeBlock(code="import sys; print('hello world!')", language=py_variant), + CodeBlock(code="a = 100 + 100; print(a)", language=py_variant), ] code_result = executor.execute_code_blocks(code_blocks) assert ( @@ -94,7 +98,7 @@ def _test_execute_code(executor: CodeExecutor) -> None: # Test running code. file_lines = ["import sys", "print('hello world!')", "a = 100 + 100", "print(a)"] - code_blocks = [CodeBlock(code="\n".join(file_lines), language="python")] + code_blocks = [CodeBlock(code="\n".join(file_lines), language=py_variant)] code_result = executor.execute_code_blocks(code_blocks) assert ( code_result.exit_code == 0