Support string based UDFs (#2195)

Author: jackgerrits

autogen/coding/func_with_reqs.py

@@ -5,12 +5,17 @@
 from typing_extensions import ParamSpec
 from textwrap import indent, dedent
 from dataclasses import dataclass, field
+from importlib.abc import SourceLoader
+import importlib
 
 T = TypeVar("T")
 P = ParamSpec("P")
 
 
-def _to_code(func: Union[FunctionWithRequirements[T, P], Callable[P, T]]) -> str:
+def _to_code(func: Union[FunctionWithRequirements[T, P], Callable[P, T], FunctionWithRequirementsStr]) -> str:
+    if isinstance(func, FunctionWithRequirementsStr):
+        return func.func
+
     code = inspect.getsource(func)
     # Strip the decorator
     if code.startswith("@"):
@@ -50,6 +55,57 @@ def to_str(i: Union[str, Alias]) -> str:
         return f"from {im.module} import {imports}"
 
 
+class _StringLoader(SourceLoader):
+    def __init__(self, data: str):
+        self.data = data
+
+    def get_source(self, fullname: str) -> str:
+        return self.data
+
+    def get_data(self, path: str) -> bytes:
+        return self.data.encode("utf-8")
+
+    def get_filename(self, fullname: str) -> str:
+        return "<not a real path>/" + fullname + ".py"
+
+
+@dataclass
+class FunctionWithRequirementsStr:
+    func: str
+    _compiled_func: Callable[..., Any]
+    _func_name: str
+    python_packages: List[str] = field(default_factory=list)
+    global_imports: List[Import] = field(default_factory=list)
+
+    def __init__(self, func: str, python_packages: List[str] = [], global_imports: List[Import] = []):
+        self.func = func
+        self.python_packages = python_packages
+        self.global_imports = global_imports
+
+        module_name = "func_module"
+        loader = _StringLoader(func)
+        spec = importlib.util.spec_from_loader(module_name, loader)
+        if spec is None:
+            raise ValueError("Could not create spec")
+        module = importlib.util.module_from_spec(spec)
+        if spec.loader is None:
+            raise ValueError("Could not create loader")
+
+        try:
+            spec.loader.exec_module(module)
+        except Exception as e:
+            raise ValueError(f"Could not compile function: {e}") from e
+
+        functions = inspect.getmembers(module, inspect.isfunction)
+        if len(functions) != 1:
+            raise ValueError("The string must contain exactly one function")
+
+        self._func_name, self._compiled_func = functions[0]
+
+    def __call__(self, *args: Any, **kwargs: Any) -> None:
+        raise NotImplementedError("String based function with requirement objects are not directly callable")
+
+
 @dataclass
 class FunctionWithRequirements(Generic[T, P]):
     func: Callable[P, T]
@@ -62,6 +118,12 @@ def from_callable(
     ) -> FunctionWithRequirements[T, P]:
         return cls(python_packages=python_packages, global_imports=global_imports, func=func)
 
+    @staticmethod
+    def from_str(
+        func: str, python_packages: List[str] = [], global_imports: List[Import] = []
+    ) -> FunctionWithRequirementsStr:
+        return FunctionWithRequirementsStr(func=func, python_packages=python_packages, global_imports=global_imports)
+
     # Type this based on F
     def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T:
         return self.func(*args, **kwargs)
@@ -91,11 +153,13 @@ def wrapper(func: Callable[P, T]) -> FunctionWithRequirements[T, P]:
     return wrapper
 
 
-def _build_python_functions_file(funcs: List[Union[FunctionWithRequirements[Any, P], Callable[..., Any]]]) -> str:
+def _build_python_functions_file(
+    funcs: List[Union[FunctionWithRequirements[Any, P], Callable[..., Any], FunctionWithRequirementsStr]]
+) -> str:
     # First collect all global imports
     global_imports = set()
     for func in funcs:
-        if isinstance(func, FunctionWithRequirements):
+        if isinstance(func, (FunctionWithRequirements, FunctionWithRequirementsStr)):
             global_imports.update(func.global_imports)
 
     content = "\n".join(map(_import_to_str, global_imports)) + "\n\n"
@@ -106,7 +170,7 @@ def _build_python_functions_file(funcs: List[Union[FunctionWithRequirements[Any,
     return content
 
 
-def to_stub(func: Callable[..., Any]) -> str:
+def to_stub(func: Union[Callable[..., Any], FunctionWithRequirementsStr]) -> str:
     """Generate a stub for a function as a string
 
     Args:
@@ -115,6 +179,9 @@ def to_stub(func: Callable[..., Any]) -> str:
     Returns:
         str: The stub for the function
     """
+    if isinstance(func, FunctionWithRequirementsStr):
+        return to_stub(func._compiled_func)
+
     content = f"def {func.__name__}{inspect.signature(func)}:\n"
     docstring = func.__doc__
 

autogen/coding/local_commandline_code_executor.py

@@ -6,7 +6,12 @@
 import warnings
 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 autogen.coding.func_with_reqs import (
+    FunctionWithRequirements,
+    FunctionWithRequirementsStr,
+    _build_python_functions_file,
+    to_stub,
+)
 
 from ..code_utils import TIMEOUT_MSG, WIN32, _cmd
 from .base import CodeBlock, CodeExecutor, CodeExtractor, CommandLineCodeResult
@@ -39,7 +44,7 @@ def __init__(
         self,
         timeout: int = 60,
         work_dir: Union[Path, str] = Path("."),
-        functions: List[Union[FunctionWithRequirements[Any, A], Callable[..., Any]]] = [],
+        functions: List[Union[FunctionWithRequirements[Any, A], Callable[..., Any], FunctionWithRequirementsStr]] = [],
     ):
         """(Experimental) A code executor class that executes code through a local command line
         environment.
@@ -104,7 +109,9 @@ def format_functions_for_prompt(self, prompt_template: str = FUNCTION_PROMPT_TEM
         )
 
     @property
-    def functions(self) -> List[Union[FunctionWithRequirements[Any, A], Callable[..., Any]]]:
+    def functions(
+        self,
+    ) -> List[Union[FunctionWithRequirements[Any, A], Callable[..., Any], FunctionWithRequirementsStr]]:
         """(Experimental) The functions that are available to the code executor."""
         return self._functions
 

test/coding/test_user_defined_functions.py

@@ -10,7 +10,7 @@
 else:
     skip = False
 
-from autogen.coding.func_with_reqs import with_requirements
+from autogen.coding.func_with_reqs import FunctionWithRequirements, with_requirements
 
 classes_to_test = [LocalCommandLineCodeExecutor]
 
@@ -137,7 +137,6 @@ def test_fails_for_function_incorrect_dep(cls) -> None:
 
 
 @pytest.mark.parametrize("cls", classes_to_test)
-@pytest.mark.skipif(skip, reason="pandas not installed")
 def test_formatted_prompt(cls) -> None:
     with tempfile.TemporaryDirectory() as temp_dir:
         executor = cls(work_dir=temp_dir, functions=[add_two_numbers])
@@ -149,3 +148,86 @@ def test_formatted_prompt(cls) -> None:
 '''
             in result
         )
+
+
+@pytest.mark.parametrize("cls", classes_to_test)
+def test_formatted_prompt_str_func(cls) -> None:
+    with tempfile.TemporaryDirectory() as temp_dir:
+        func = FunctionWithRequirements.from_str(
+            '''
+def add_two_numbers(a: int, b: int) -> int:
+    """Add two numbers together."""
+    return a + b
+'''
+        )
+        executor = cls(work_dir=temp_dir, functions=[func])
+
+        result = executor.format_functions_for_prompt()
+        assert (
+            '''def add_two_numbers(a: int, b: int) -> int:
+    """Add two numbers together."""
+'''
+            in result
+        )
+
+
+@pytest.mark.parametrize("cls", classes_to_test)
+def test_can_load_str_function_with_reqs(cls) -> None:
+    with tempfile.TemporaryDirectory() as temp_dir:
+        func = FunctionWithRequirements.from_str(
+            '''
+def add_two_numbers(a: int, b: int) -> int:
+    """Add two numbers together."""
+    return a + b
+'''
+        )
+
+        executor = cls(work_dir=temp_dir, functions=[func])
+        code = f"""from {cls.FUNCTIONS_MODULE} import add_two_numbers
+print(add_two_numbers(1, 2))"""
+
+        result = executor.execute_code_blocks(
+            code_blocks=[
+                CodeBlock(language="python", code=code),
+            ]
+        )
+        assert result.output == "3\n"
+        assert result.exit_code == 0
+
+
+@pytest.mark.parametrize("cls", classes_to_test)
+def test_cant_load_broken_str_function_with_reqs(cls) -> None:
+
+    with pytest.raises(ValueError):
+        _ = FunctionWithRequirements.from_str(
+            '''
+invaliddef add_two_numbers(a: int, b: int) -> int:
+    """Add two numbers together."""
+    return a + b
+'''
+        )
+
+
+@pytest.mark.parametrize("cls", classes_to_test)
+def test_cant_run_broken_str_function_with_reqs(cls) -> None:
+    with tempfile.TemporaryDirectory() as temp_dir:
+        func = FunctionWithRequirements.from_str(
+            '''
+def add_two_numbers(a: int, b: int) -> int:
+    """Add two numbers together."""
+    return a + b
+'''
+        )
+
+        code = f"""from {cls.FUNCTIONS_MODULE} import add_two_numbers
+print(add_two_numbers(object(), False))"""
+
+        executor = cls(work_dir=temp_dir, functions=[func])
+
+        result = executor.execute_code_blocks(
+            code_blocks=[
+                CodeBlock(language="python", code=code),
+            ]
+        )
+        assert "TypeError: unsupported operand type(s) for +:" in result.output
+        assert result.exit_code == 1