Implement User Defined Functions for Local CLI Executor (microsoft#2102)

* Implement user defined functions feature for local cli exec, add docs

* add tests, update docs

* fixes

* fix test

* add pandas test dep

* install test

* provide template as func

* formatting

* undo change

* address comments

* add test deps

* formatting

* test only in 1 env

* formatting

* remove test for local only

---------

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

.github/workflows/build.yml

@@ -43,10 +43,10 @@ jobs:
           python -c "import autogen"
           pip install pytest mock
       - name: Install optional dependencies for code executors
-        # code executors auto skip without deps, so only run for python 3.11
+        # code executors and udfs auto skip without deps, so only run for python 3.11
         if: matrix.python-version == '3.11'
         run: |
-          pip install -e ".[jupyter-executor]"
+          pip install -e ".[jupyter-executor,test]"
           python -m ipykernel install --user --name python3
       - name: Set AUTOGEN_USE_DOCKER based on OS
         shell: bash

autogen/coding/func_with_reqs.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+import inspect
+import functools
+from typing import Any, Callable, List, TypeVar, Generic, Union
+from typing_extensions import ParamSpec
+from textwrap import indent, dedent
+from dataclasses import dataclass, field
+
+T = TypeVar("T")
+P = ParamSpec("P")
+
+
+def _to_code(func: Union[FunctionWithRequirements[T, P], Callable[P, T]]) -> str:
+    code = inspect.getsource(func)
+    # Strip the decorator
+    if code.startswith("@"):
+        code = code[code.index("\n") + 1 :]
+    return code
+
+
+@dataclass
+class Alias:
+    name: str
+    alias: str
+
+
+@dataclass
+class ImportFromModule:
+    module: str
+    imports: List[Union[str, Alias]]
+
+
+Import = Union[str, ImportFromModule, Alias]
+
+
+def _import_to_str(im: Import) -> str:
+    if isinstance(im, str):
+        return f"import {im}"
+    elif isinstance(im, Alias):
+        return f"import {im.name} as {im.alias}"
+    else:
+
+        def to_str(i: Union[str, Alias]) -> str:
+            if isinstance(i, str):
+                return i
+            else:
+                return f"{i.name} as {i.alias}"
+
+        imports = ", ".join(map(to_str, im.imports))
+        return f"from {im.module} import {imports}"
+
+
+@dataclass
+class FunctionWithRequirements(Generic[T, P]):
+    func: Callable[P, T]
+    python_packages: List[str] = field(default_factory=list)
+    global_imports: List[Import] = field(default_factory=list)
+
+    @classmethod
+    def from_callable(
+        cls, func: Callable[P, T], python_packages: List[str] = [], global_imports: List[Import] = []
+    ) -> FunctionWithRequirements[T, P]:
+        return cls(python_packages=python_packages, global_imports=global_imports, func=func)
+
+    # Type this based on F
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T:
+        return self.func(*args, **kwargs)
+
+
+def with_requirements(
+    python_packages: List[str] = [], global_imports: List[Import] = []
+) -> Callable[[Callable[P, T]], FunctionWithRequirements[T, P]]:
+    """Decorate a function with package and import requirements
+
+    Args:
+        python_packages (List[str], optional): Packages required to function. Can include version info.. Defaults to [].
+        global_imports (List[Import], optional): Required imports. Defaults to [].
+
+    Returns:
+        Callable[[Callable[P, T]], FunctionWithRequirements[T, P]]: The decorated function
+    """
+
+    def wrapper(func: Callable[P, T]) -> FunctionWithRequirements[T, P]:
+        func_with_reqs = FunctionWithRequirements(
+            python_packages=python_packages, global_imports=global_imports, func=func
+        )
+
+        functools.update_wrapper(func_with_reqs, func)
+        return func_with_reqs
+
+    return wrapper
+
+
+def _build_python_functions_file(funcs: List[Union[FunctionWithRequirements[Any, P], Callable[..., Any]]]) -> str:
+    # First collect all global imports
+    global_imports = set()
+    for func in funcs:
+        if isinstance(func, FunctionWithRequirements):
+            global_imports.update(func.global_imports)
+
+    content = "\n".join(map(_import_to_str, global_imports)) + "\n\n"
+
+    for func in funcs:
+        content += _to_code(func) + "\n\n"
+
+    return content
+
+
+def to_stub(func: Callable[..., Any]) -> str:
+    """Generate a stub for a function as a string
+
+    Args:
+        func (Callable[..., Any]): The function to generate a stub for
+
+    Returns:
+        str: The stub for the function
+    """
+    content = f"def {func.__name__}{inspect.signature(func)}:\n"
+    docstring = func.__doc__
+
+    if docstring:
+        docstring = dedent(docstring)
+        docstring = '"""' + docstring + '"""'
+        docstring = indent(docstring, "    ")
+        content += docstring + "\n"
+
+    content += "    ..."
+    return content

autogen/coding/local_commandline_code_executor.py

@@ -1,31 +1,45 @@
 from hashlib import md5
-import os
 from pathlib import Path
 import re
+from string import Template
 import sys
-import uuid
 import warnings
-from typing import ClassVar, List, Union
+from typing import Any, Callable, ClassVar, List, TypeVar, Union, cast
+from typing_extensions import ParamSpec
+from autogen.coding.func_with_reqs import FunctionWithRequirements, _build_python_functions_file, to_stub
 
-from ..agentchat.agent import LLMAgent
-from ..code_utils import TIMEOUT_MSG, WIN32, _cmd, execute_code
+from ..code_utils import 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
 
 import subprocess
 
+import logging
+
 __all__ = ("LocalCommandLineCodeExecutor",)
 
+A = ParamSpec("A")
+
 
 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.
+
+For example, if there was a function called `foo` you could import it by writing `from $module_name import foo`
+
+$functions"""
 
     def __init__(
         self,
         timeout: int = 60,
         work_dir: Union[Path, str] = Path("."),
+        functions: List[Union[FunctionWithRequirements[Any, A], Callable[..., Any]]] = [],
     ):
         """(Experimental) A code executor class that executes code through a local command line
         environment.
@@ -48,6 +62,7 @@ def __init__(
             work_dir (str): The working directory for the code execution. If None,
                 a default working directory will be used. The default working
                 directory is the current directory ".".
+            functions (List[Union[FunctionWithRequirements[Any, A], Callable[..., Any]]]): A list of functions that are available to the code executor. Default is an empty list.
         """
 
         if timeout < 1:
@@ -62,6 +77,38 @@ def __init__(
         self._timeout = timeout
         self._work_dir: Path = work_dir
 
+        self._functions = functions
+        # Setup could take some time so we intentionally wait for the first code block to do it.
+        if len(functions) > 0:
+            self._setup_functions_complete = False
+        else:
+            self._setup_functions_complete = True
+
+    def format_functions_for_prompt(self, prompt_template: str = FUNCTION_PROMPT_TEMPLATE) -> str:
+        """(Experimental) Format the functions for a prompt.
+
+        The template includes two variables:
+        - `$module_name`: The module name.
+        - `$functions`: The functions formatted as stubs with two newlines between each function.
+
+        Args:
+            prompt_template (str): The prompt template. Default is the class default.
+
+        Returns:
+            str: The formatted prompt.
+        """
+
+        template = Template(prompt_template)
+        return template.substitute(
+            module_name=self.FUNCTIONS_MODULE,
+            functions="\n\n".join([to_stub(func) for func in self._functions]),
+        )
+
+    @property
+    def functions(self) -> List[Union[FunctionWithRequirements[Any, A], Callable[..., Any]]]:
+        """(Experimental) The functions that are available to the code executor."""
+        return self._functions
+
     @property
     def timeout(self) -> int:
         """(Experimental) The timeout for code execution."""
@@ -99,6 +146,39 @@ def sanitize_command(lang: str, code: str) -> None:
                 if re.search(pattern, code):
                     raise ValueError(f"Potentially dangerous command detected: {message}")
 
+    def _setup_functions(self) -> None:
+        func_file_content = _build_python_functions_file(self._functions)
+        func_file = self._work_dir / self.FUNCTIONS_FILENAME
+        func_file.write_text(func_file_content)
+
+        # Collect requirements
+        lists_of_packages = [x.python_packages for x in self._functions if isinstance(x, FunctionWithRequirements)]
+        flattened_packages = [item for sublist in lists_of_packages for item in sublist]
+        required_packages = list(set(flattened_packages))
+        if len(required_packages) > 0:
+            logging.info("Ensuring packages are installed in executor.")
+
+            cmd = [sys.executable, "-m", "pip", "install"]
+            cmd.extend(required_packages)
+
+            try:
+                result = subprocess.run(
+                    cmd, cwd=self._work_dir, capture_output=True, text=True, timeout=float(self._timeout)
+                )
+            except subprocess.TimeoutExpired as e:
+                raise ValueError("Pip install timed out") from e
+
+            if result.returncode != 0:
+                raise ValueError(f"Pip install failed. {result.stdout}, {result.stderr}")
+
+        # Attempt to load the function file to check for syntax errors, imports etc.
+        exec_result = self._execute_code_dont_check_setup([CodeBlock(code=func_file_content, language="python")])
+
+        if exec_result.exit_code != 0:
+            raise ValueError(f"Functions failed to load: {exec_result.output}")
+
+        self._setup_functions_complete = True
+
     def execute_code_blocks(self, code_blocks: List[CodeBlock]) -> CommandLineCodeResult:
         """(Experimental) Execute the code blocks and return the result.
 
@@ -107,6 +187,13 @@ def execute_code_blocks(self, code_blocks: List[CodeBlock]) -> CommandLineCodeRe
 
         Returns:
             CommandLineCodeResult: The result of the code execution."""
+
+        if not self._setup_functions_complete:
+            self._setup_functions()
+
+        return self._execute_code_dont_check_setup(code_blocks)
+
+    def _execute_code_dont_check_setup(self, code_blocks: List[CodeBlock]) -> CommandLineCodeResult:
         logs_all = ""
         file_names = []
         for code_block in code_blocks:

setup.py

@@ -55,6 +55,7 @@
             "pre-commit",
             "pytest-asyncio",
             "pytest>=6.1.1,<8",
+            "pandas",
         ],
         "blendsearch": ["flaml[blendsearch]"],
         "mathchat": ["sympy", "pydantic==1.10.9", "wolframalpha"],