diff --git a/autogen/code_utils.py b/autogen/code_utils.py index 299a39fba9d..f468c3d1bfd 100644 --- a/autogen/code_utils.py +++ b/autogen/code_utils.py @@ -2,6 +2,7 @@ import os import pathlib import re +import string import subprocess import sys import time @@ -224,6 +225,29 @@ def _cmd(lang): raise NotImplementedError(f"{lang} not recognized in code execution") +def _sanitize_filename_for_docker_tag(filename: str) -> str: + """Convert a filename to a valid docker tag. + See https://docs.docker.com/engine/reference/commandline/tag/ for valid tag + format. + + Args: + filename (str): The filename to be converted. + + Returns: + str: The sanitized Docker tag. + """ + # Replace any character not allowed with an underscore + allowed_chars = set(string.ascii_letters + string.digits + "_.-") + sanitized = "".join(char if char in allowed_chars else "_" for char in filename) + + # Ensure it does not start with a period or a dash + if sanitized.startswith(".") or sanitized.startswith("-"): + sanitized = "_" + sanitized[1:] + + # Truncate if longer than 128 characters + return sanitized[:128] + + def execute_code( code: Optional[str] = None, timeout: Optional[int] = None, @@ -379,7 +403,7 @@ def execute_code( cmd = [ "sh", "-c", - f"{_cmd(lang)} {filename}; exit_code=$?; echo -n {exit_code_str}; echo -n $exit_code; echo {exit_code_str}", + f'{_cmd(lang)} "{filename}"; exit_code=$?; echo -n {exit_code_str}; echo -n $exit_code; echo {exit_code_str}', ] # create a docker container container = client.containers.run( @@ -403,7 +427,7 @@ def execute_code( # get the container logs logs = container.logs().decode("utf-8").rstrip() # commit the image - tag = filename.replace("/", "") + tag = _sanitize_filename_for_docker_tag(filename) container.commit(repository="python", tag=tag) # remove the container container.remove() diff --git a/test/test_code.py b/test/test_code.py index ec8ada2c906..e7306125de2 100644 --- a/test/test_code.py +++ b/test/test_code.py @@ -1,3 +1,4 @@ +import importlib.metadata import os import sys import unittest @@ -22,6 +23,21 @@ here = os.path.abspath(os.path.dirname(__file__)) +def is_package_installed(package_name): + """Check if a package is installed. This is a preferred way to check if a + package is installed or not than doing a try-catch around an import + because it avoids name conflict with local modules and + code execution in the imported module.""" + try: + importlib.metadata.version(package_name) + return True + except importlib.metadata.PackageNotFoundError: + return False + + +docker_package_installed = is_package_installed("docker") + + # def test_find_code(): # try: # import openai @@ -293,10 +309,6 @@ def scrape(url): assert len(codeblocks) == 1 and codeblocks[0] == ("", "source setup.sh") -@pytest.mark.skipif( - sys.platform in ["darwin"], - reason="do not run on MacOS", -) def test_execute_code(use_docker=None): try: import docker @@ -338,15 +350,27 @@ def test_execute_code(use_docker=None): assert isinstance(image, str) or docker is None or os.path.exists("/.dockerenv") or use_docker is False +@pytest.mark.skipif(docker_package_installed is False, reason="docker package not installed") +def test_execute_code_with_custom_filename_on_docker(): + exit_code, msg, image = execute_code("print('hello world')", filename="tmp/codetest.py", use_docker=True) + assert exit_code == 0 and msg == "hello world\n", msg + assert image == "python:tmp_codetest.py" + + +@pytest.mark.skipif(docker_package_installed is False, reason="docker package not installed") +def test_execute_code_with_misformed_filename_on_docker(): + exit_code, msg, image = execute_code( + "print('hello world')", filename="tmp/codetest.py (some extra information)", use_docker=True + ) + assert exit_code == 0 and msg == "hello world\n", msg + assert image == "python:tmp_codetest.py__some_extra_information_" + + def test_execute_code_raises_when_code_and_filename_are_both_none(): with pytest.raises(AssertionError): execute_code(code=None, filename=None) -@pytest.mark.skipif( - sys.platform in ["darwin"], - reason="do not run on MacOS", -) def test_execute_code_nodocker(): test_execute_code(use_docker=False)