diff --git a/autogen/coding/local_commandline_code_executor.py b/autogen/coding/local_commandline_code_executor.py index 16e37851c19..e012b322614 100644 --- a/autogen/coding/local_commandline_code_executor.py +++ b/autogen/coding/local_commandline_code_executor.py @@ -30,8 +30,6 @@ class LocalCommandLineCodeExecutor(CodeExecutor): SUPPORTED_LANGUAGES: ClassVar[List[str]] = ["bash", "shell", "sh", "pwsh", "powershell", "ps1", "python"] - FUNCTIONS_MODULE: ClassVar[str] = "functions" - FUNCTIONS_FILENAME: ClassVar[str] = "functions.py" FUNCTION_PROMPT_TEMPLATE: ClassVar[ str ] = """You have access to the following user defined functions. They can be accessed from the module called `$module_name` by their function names. @@ -45,6 +43,7 @@ def __init__( timeout: int = 60, work_dir: Union[Path, str] = Path("."), functions: List[Union[FunctionWithRequirements[Any, A], Callable[..., Any], FunctionWithRequirementsStr]] = [], + functions_module: str = "functions", ): """(Experimental) A code executor class that executes code through a local command line environment. @@ -76,6 +75,11 @@ def __init__( if isinstance(work_dir, str): work_dir = Path(work_dir) + if not functions_module.isidentifier(): + raise ValueError("Module name must be a valid Python identifier") + + self._functions_module = functions_module + work_dir.mkdir(exist_ok=True) self._timeout = timeout @@ -104,10 +108,15 @@ def format_functions_for_prompt(self, prompt_template: str = FUNCTION_PROMPT_TEM template = Template(prompt_template) return template.substitute( - module_name=self.FUNCTIONS_MODULE, + module_name=self._functions_module, functions="\n\n".join([to_stub(func) for func in self._functions]), ) + @property + def functions_module(self) -> str: + """(Experimental) The module name for the functions.""" + return self._functions_module + @property def functions( self, @@ -154,7 +163,7 @@ def sanitize_command(lang: str, code: str) -> None: def _setup_functions(self) -> None: func_file_content = _build_python_functions_file(self._functions) - func_file = self._work_dir / self.FUNCTIONS_FILENAME + func_file = self._work_dir / f"{self._functions_module}.py" func_file.write_text(func_file_content) # Collect requirements diff --git a/test/coding/test_user_defined_functions.py b/test/coding/test_user_defined_functions.py index 7a68acdc58c..53b3c5045dd 100644 --- a/test/coding/test_user_defined_functions.py +++ b/test/coding/test_user_defined_functions.py @@ -54,7 +54,7 @@ def function_missing_reqs() -> "pandas.DataFrame": def test_can_load_function_with_reqs(cls) -> None: with tempfile.TemporaryDirectory() as temp_dir: executor = cls(work_dir=temp_dir, functions=[load_data]) - code = f"""from {cls.FUNCTIONS_MODULE} import load_data + code = f"""from {executor.functions_module} import load_data import pandas # Get first row's name @@ -74,7 +74,7 @@ def test_can_load_function_with_reqs(cls) -> None: def test_can_load_function(cls) -> None: with tempfile.TemporaryDirectory() as temp_dir: executor = cls(work_dir=temp_dir, functions=[add_two_numbers]) - code = f"""from {cls.FUNCTIONS_MODULE} import add_two_numbers + code = f"""from {executor.functions_module} import add_two_numbers print(add_two_numbers(1, 2))""" result = executor.execute_code_blocks( @@ -93,7 +93,7 @@ def test_can_load_function(cls) -> None: # def test_fails_for_missing_reqs(cls) -> None: # with tempfile.TemporaryDirectory() as temp_dir: # executor = cls(work_dir=temp_dir, functions=[function_missing_reqs]) -# code = f"""from {cls.FUNCTIONS_MODULE} import function_missing_reqs +# code = f"""from {executor.functions_module} import function_missing_reqs # function_missing_reqs()""" # with pytest.raises(ValueError): @@ -109,7 +109,7 @@ def test_can_load_function(cls) -> None: def test_fails_for_function_incorrect_import(cls) -> None: with tempfile.TemporaryDirectory() as temp_dir: executor = cls(work_dir=temp_dir, functions=[function_incorrect_import]) - code = f"""from {cls.FUNCTIONS_MODULE} import function_incorrect_import + code = f"""from {executor.functions_module} import function_incorrect_import function_incorrect_import()""" with pytest.raises(ValueError): @@ -125,7 +125,7 @@ def test_fails_for_function_incorrect_import(cls) -> None: def test_fails_for_function_incorrect_dep(cls) -> None: with tempfile.TemporaryDirectory() as temp_dir: executor = cls(work_dir=temp_dir, functions=[function_incorrect_dep]) - code = f"""from {cls.FUNCTIONS_MODULE} import function_incorrect_dep + code = f"""from {executor.functions_module} import function_incorrect_dep function_incorrect_dep()""" with pytest.raises(ValueError): @@ -183,7 +183,7 @@ def add_two_numbers(a: int, b: int) -> int: ) executor = cls(work_dir=temp_dir, functions=[func]) - code = f"""from {cls.FUNCTIONS_MODULE} import add_two_numbers + code = f"""from {executor.functions_module} import add_two_numbers print(add_two_numbers(1, 2))""" result = executor.execute_code_blocks( @@ -219,10 +219,9 @@ def add_two_numbers(a: int, b: int) -> int: ''' ) - code = f"""from {cls.FUNCTIONS_MODULE} import add_two_numbers -print(add_two_numbers(object(), False))""" - executor = cls(work_dir=temp_dir, functions=[func]) + code = f"""from {executor.functions_module} import add_two_numbers +print(add_two_numbers(object(), False))""" result = executor.execute_code_blocks( code_blocks=[