diff --git a/python/packages/autogen-core/src/autogen_core/code_executor/_func_with_reqs.py b/python/packages/autogen-core/src/autogen_core/code_executor/_func_with_reqs.py index b7f4fcaef..77fc0b831 100644 --- a/python/packages/autogen-core/src/autogen_core/code_executor/_func_with_reqs.py +++ b/python/packages/autogen-core/src/autogen_core/code_executor/_func_with_reqs.py @@ -9,7 +9,7 @@ from dataclasses import dataclass, field from importlib.abc import SourceLoader from importlib.util import module_from_spec, spec_from_loader from textwrap import dedent, indent -from typing import Any, Callable, Generic, List, Sequence, Set, TypeVar, Union +from typing import Any, Callable, Generic, List, Sequence, Set, Tuple, TypeVar, Union from typing_extensions import ParamSpec @@ -21,23 +21,38 @@ def _to_code(func: Union[FunctionWithRequirements[T, P], Callable[P, T], Functio if isinstance(func, FunctionWithRequirementsStr): return func.func - code = inspect.getsource(func) + if isinstance(func, FunctionWithRequirements): + code = inspect.getsource(func.func) + else: + code = inspect.getsource(func) # Strip the decorator if code.startswith("@"): code = code[code.index("\n") + 1 :] return code -@dataclass +@dataclass(frozen=True) class Alias: name: str alias: str -@dataclass +@dataclass(frozen=True) class ImportFromModule: module: str - imports: List[Union[str, Alias]] + imports: Tuple[Union[str, Alias], ...] + + ## backward compatibility + def __init__( + self, + module: str, + imports: Union[Tuple[Union[str, Alias], ...], List[Union[str, Alias]]], + ): + object.__setattr__(self, "module", module) + if isinstance(imports, list): + object.__setattr__(self, "imports", tuple(imports)) + else: + object.__setattr__(self, "imports", imports) Import = Union[str, ImportFromModule, Alias] diff --git a/python/packages/autogen-core/tests/test_code_executor.py b/python/packages/autogen-core/tests/test_code_executor.py new file mode 100644 index 000000000..a8412c58d --- /dev/null +++ b/python/packages/autogen-core/tests/test_code_executor.py @@ -0,0 +1,53 @@ +import textwrap + +import pytest +from autogen_core.code_executor import ( + Alias, + FunctionWithRequirements, + FunctionWithRequirementsStr, + ImportFromModule, +) +from autogen_core.code_executor._func_with_reqs import build_python_functions_file +from pandas import DataFrame, concat + + +def template_function() -> DataFrame: # type: ignore + data1 = { + "name": ["John", "Anna"], + "location": ["New York", "Paris"], + "age": [24, 13], + } + data2 = { + "name": ["Peter", "Linda"], + "location": ["Berlin", "London"], + "age": [53, 33], + } + df1 = DataFrame.from_dict(data1) # type: ignore + df2 = DataFrame.from_dict(data2) # type: ignore + return concat([df1, df2]) # type: ignore + + +@pytest.mark.asyncio +async def test_hashability_Import() -> None: + function = FunctionWithRequirements.from_callable( # type: ignore + template_function, + ["pandas"], + [ImportFromModule("pandas", ["DataFrame", "concat"])], + ) + functions_module = build_python_functions_file([function]) # type: ignore + + assert "from pandas import DataFrame, concat" in functions_module + + function2: FunctionWithRequirementsStr = FunctionWithRequirements.from_str( + textwrap.dedent( + """ + def template_function2(): + return pd.Series([1, 2]) + """ + ), + "pandas", + [Alias("pandas", "pd")], + ) + functions_module2 = build_python_functions_file([function2]) + + assert "import pandas as pd" in functions_module2