From 4444247b00445ba82c346d22d5e0fbe272bb15b0 Mon Sep 17 00:00:00 2001 From: Victor Dibia Date: Mon, 13 Jan 2025 22:59:56 -0800 Subject: [PATCH 01/10] vi1 for declarative tools --- .../src/autogen_core/tools/_base.py | 32 ++++++++++++++++--- 1 file changed, 27 insertions(+), 5 deletions(-) diff --git a/python/packages/autogen-core/src/autogen_core/tools/_base.py b/python/packages/autogen-core/src/autogen_core/tools/_base.py index 7c4042e9afd6..f66c4eb025a5 100644 --- a/python/packages/autogen-core/src/autogen_core/tools/_base.py +++ b/python/packages/autogen-core/src/autogen_core/tools/_base.py @@ -9,6 +9,7 @@ from .. import CancellationToken from .._function_utils import normalize_annotated_type +from .._component_config import ComponentBase, Component T = TypeVar("T", bound=BaseModel, contravariant=True) @@ -26,7 +27,10 @@ class ToolSchema(TypedDict): @runtime_checkable -class Tool(Protocol): +class Tool(Protocol, ComponentBase[BaseModel]): + + component_type = "tool" + @property def name(self) -> str: ... @@ -44,7 +48,8 @@ def state_type(self) -> Type[BaseModel] | None: ... def return_value_as_string(self, value: Any) -> str: ... - async def run_json(self, args: Mapping[str, Any], cancellation_token: CancellationToken) -> Any: ... + async def run_json( + self, args: Mapping[str, Any], cancellation_token: CancellationToken) -> Any: ... def save_state_json(self) -> Mapping[str, Any]: ... @@ -56,7 +61,14 @@ def load_state_json(self, state: Mapping[str, Any]) -> None: ... StateT = TypeVar("StateT", bound=BaseModel) -class BaseTool(ABC, Tool, Generic[ArgsT, ReturnT]): +class BaseToolConfig(BaseModel): + name: str + description: str + args_type: Type[Any] + return_type: Type[Any] + + +class BaseTool(ABC, Tool, Generic[ArgsT, ReturnT], Component[BaseToolConfig]): def __init__( self, args_type: Type[ArgsT], @@ -75,7 +87,8 @@ def schema(self) -> ToolSchema: model_schema: Dict[str, Any] = self._args_type.model_json_schema() if "$defs" in model_schema: - model_schema = cast(Dict[str, Any], jsonref.replace_refs(obj=model_schema, proxies=False)) # type: ignore + model_schema = cast(Dict[str, Any], jsonref.replace_refs( + obj=model_schema, proxies=False)) # type: ignore del model_schema["$defs"] tool_schema = ToolSchema( @@ -119,7 +132,8 @@ def return_value_as_string(self, value: Any) -> str: return str(value) @abstractmethod - async def run(self, args: ArgsT, cancellation_token: CancellationToken) -> ReturnT: ... + async def run(self, args: ArgsT, + cancellation_token: CancellationToken) -> ReturnT: ... async def run_json(self, args: Mapping[str, Any], cancellation_token: CancellationToken) -> Any: return_value = await self.run(self._args_type.model_validate(args), cancellation_token) @@ -131,6 +145,14 @@ def save_state_json(self) -> Mapping[str, Any]: def load_state_json(self, state: Mapping[str, Any]) -> None: pass + def _to_config(self) -> BaseToolConfig: + return BaseToolConfig( + name=self._name, + description=self._description, + args_type=self._args_type, + return_type=self._return_type, + ) + class BaseToolWithState(BaseTool[ArgsT, ReturnT], ABC, Generic[ArgsT, ReturnT, StateT]): def __init__( From 31a8bb15decbfe81606c0ffa3539c4930c827459 Mon Sep 17 00:00:00 2001 From: Victor Dibia Date: Tue, 14 Jan 2025 15:38:05 -0800 Subject: [PATCH 02/10] make functtools declarative --- .../code_executor/_func_with_reqs.py | 10 ++- .../src/autogen_core/tools/__init__.py | 2 + .../src/autogen_core/tools/_base.py | 37 +++-------- .../src/autogen_core/tools/_function_tool.py | 66 ++++++++++++++++++- .../models/openai/_openai_client.py | 2 +- 5 files changed, 85 insertions(+), 32 deletions(-) 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 77fc0b831427..50a8c5280935 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 @@ -42,7 +42,7 @@ class ImportFromModule: module: str imports: Tuple[Union[str, Alias], ...] - ## backward compatibility + # backward compatibility def __init__( self, module: str, @@ -214,3 +214,11 @@ def to_stub(func: Union[Callable[..., Any], FunctionWithRequirementsStr]) -> str content += " ..." return content + + +def to_code(func: Union[FunctionWithRequirements[T, P], Callable[P, T], FunctionWithRequirementsStr]) -> str: + return _to_code(func) + + +def import_to_str(im: Import) -> str: + return _import_to_str(im) diff --git a/python/packages/autogen-core/src/autogen_core/tools/__init__.py b/python/packages/autogen-core/src/autogen_core/tools/__init__.py index 52a9d725f7d6..250fd273862c 100644 --- a/python/packages/autogen-core/src/autogen_core/tools/__init__.py +++ b/python/packages/autogen-core/src/autogen_core/tools/__init__.py @@ -1,3 +1,4 @@ +from ..code_executor import ImportFromModule from ._base import BaseTool, BaseToolWithState, ParametersSchema, Tool, ToolSchema from ._function_tool import FunctionTool @@ -8,4 +9,5 @@ "BaseTool", "BaseToolWithState", "FunctionTool", + "ImportFromModule", ] diff --git a/python/packages/autogen-core/src/autogen_core/tools/_base.py b/python/packages/autogen-core/src/autogen_core/tools/_base.py index f66c4eb025a5..b484ef84f3e9 100644 --- a/python/packages/autogen-core/src/autogen_core/tools/_base.py +++ b/python/packages/autogen-core/src/autogen_core/tools/_base.py @@ -8,8 +8,8 @@ from typing_extensions import NotRequired from .. import CancellationToken +from .._component_config import ComponentBase from .._function_utils import normalize_annotated_type -from .._component_config import ComponentBase, Component T = TypeVar("T", bound=BaseModel, contravariant=True) @@ -27,10 +27,7 @@ class ToolSchema(TypedDict): @runtime_checkable -class Tool(Protocol, ComponentBase[BaseModel]): - - component_type = "tool" - +class Tool(Protocol): @property def name(self) -> str: ... @@ -48,8 +45,7 @@ def state_type(self) -> Type[BaseModel] | None: ... def return_value_as_string(self, value: Any) -> str: ... - async def run_json( - self, args: Mapping[str, Any], cancellation_token: CancellationToken) -> Any: ... + async def run_json(self, args: Mapping[str, Any], cancellation_token: CancellationToken) -> Any: ... def save_state_json(self) -> Mapping[str, Any]: ... @@ -61,14 +57,9 @@ def load_state_json(self, state: Mapping[str, Any]) -> None: ... StateT = TypeVar("StateT", bound=BaseModel) -class BaseToolConfig(BaseModel): - name: str - description: str - args_type: Type[Any] - return_type: Type[Any] - +class BaseTool(ABC, Tool, Generic[ArgsT, ReturnT], ComponentBase[BaseModel]): + component_type = "tool" -class BaseTool(ABC, Tool, Generic[ArgsT, ReturnT], Component[BaseToolConfig]): def __init__( self, args_type: Type[ArgsT], @@ -87,8 +78,7 @@ def schema(self) -> ToolSchema: model_schema: Dict[str, Any] = self._args_type.model_json_schema() if "$defs" in model_schema: - model_schema = cast(Dict[str, Any], jsonref.replace_refs( - obj=model_schema, proxies=False)) # type: ignore + model_schema = cast(Dict[str, Any], jsonref.replace_refs(obj=model_schema, proxies=False)) # type: ignore del model_schema["$defs"] tool_schema = ToolSchema( @@ -132,8 +122,7 @@ def return_value_as_string(self, value: Any) -> str: return str(value) @abstractmethod - async def run(self, args: ArgsT, - cancellation_token: CancellationToken) -> ReturnT: ... + async def run(self, args: ArgsT, cancellation_token: CancellationToken) -> ReturnT: ... async def run_json(self, args: Mapping[str, Any], cancellation_token: CancellationToken) -> Any: return_value = await self.run(self._args_type.model_validate(args), cancellation_token) @@ -145,16 +134,8 @@ def save_state_json(self) -> Mapping[str, Any]: def load_state_json(self, state: Mapping[str, Any]) -> None: pass - def _to_config(self) -> BaseToolConfig: - return BaseToolConfig( - name=self._name, - description=self._description, - args_type=self._args_type, - return_type=self._return_type, - ) - -class BaseToolWithState(BaseTool[ArgsT, ReturnT], ABC, Generic[ArgsT, ReturnT, StateT]): +class BaseToolWithState(BaseTool[ArgsT, ReturnT], ABC, Generic[ArgsT, ReturnT, StateT], ComponentBase[BaseModel]): def __init__( self, args_type: Type[ArgsT], @@ -166,6 +147,8 @@ def __init__( super().__init__(args_type, return_type, name, description) self._state_type = state_type + component_type = "tool" + @abstractmethod def save_state(self) -> StateT: ... diff --git a/python/packages/autogen-core/src/autogen_core/tools/_function_tool.py b/python/packages/autogen-core/src/autogen_core/tools/_function_tool.py index 026fc845e9c2..26a1b059ce99 100644 --- a/python/packages/autogen-core/src/autogen_core/tools/_function_tool.py +++ b/python/packages/autogen-core/src/autogen_core/tools/_function_tool.py @@ -1,18 +1,31 @@ import asyncio import functools -from typing import Any, Callable +from typing import Any, Callable, Sequence from pydantic import BaseModel +from typing_extensions import Self from .. import CancellationToken +from .._component_config import Component from .._function_utils import ( args_base_model_from_signature, get_typed_signature, ) +from ..code_executor._func_with_reqs import Import, import_to_str, to_code from ._base import BaseTool -class FunctionTool(BaseTool[BaseModel, BaseModel]): +class FunctionToolConfig(BaseModel): + """Configuration for a function tool.""" + + source_code: str + name: str + description: str + global_imports: Sequence[Import] + has_cancellation_support: bool + + +class FunctionTool(BaseTool[BaseModel, BaseModel], Component[FunctionToolConfig]): """ Create custom tools by wrapping standard Python functions. @@ -64,8 +77,14 @@ async def example(): asyncio.run(example()) """ - def __init__(self, func: Callable[..., Any], description: str, name: str | None = None) -> None: + component_provider_override = "autogen_core.tools.FunctionTool" + component_config_schema = FunctionToolConfig + + def __init__( + self, func: Callable[..., Any], description: str, name: str | None = None, global_imports: Sequence[Import] = [] + ) -> None: self._func = func + self._global_imports = global_imports signature = get_typed_signature(func) func_name = name or func.__name__ args_model = args_base_model_from_signature(func_name + "args", signature) @@ -98,3 +117,44 @@ async def run(self, args: BaseModel, cancellation_token: CancellationToken) -> A result = await future return result + + def _to_config(self) -> FunctionToolConfig: + return FunctionToolConfig( + source_code=to_code(self._func), + global_imports=self._global_imports, + name=self.name, + description=self.description, + has_cancellation_support=self._has_cancellation_support, + ) + + @classmethod + def _from_config(cls, config: FunctionToolConfig) -> Self: + exec_globals: dict[str, Any] = {} + + # Execute imports first + for import_stmt in config.global_imports: + import_code = import_to_str(import_stmt) + try: + exec(import_code, exec_globals) + except ModuleNotFoundError as e: + raise ModuleNotFoundError( + f"Failed to import {import_code}: Module not found. Please ensure the module is installed." + ) from e + except ImportError as e: + raise ImportError(f"Failed to import {import_code}: {str(e)}") from e + except Exception as e: + raise RuntimeError(f"Unexpected error while importing {import_code}: {str(e)}") from e + + # Execute function code + try: + exec(config.source_code, exec_globals) + func_name = config.source_code.split("def ")[1].split("(")[0] + except Exception as e: + raise ValueError(f"Could not compile and load function: {e}") from e + + # Get function and verify it's callable + func: Callable[..., Any] = exec_globals[func_name] + if not callable(func): + raise TypeError(f"Expected function but got {type(func)}") + + return cls(func, "", None) diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py index b525e6340fd0..04512d7ce9b2 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py @@ -30,13 +30,13 @@ Image, MessageHandlerContext, ) -from autogen_core.models import FinishReasons from autogen_core.logging import LLMCallEvent from autogen_core.models import ( AssistantMessage, ChatCompletionClient, ChatCompletionTokenLogprob, CreateResult, + FinishReasons, FunctionExecutionResultMessage, LLMMessage, ModelCapabilities, # type: ignore From f0375e86642d10f059bf3545e80881bf432e158d Mon Sep 17 00:00:00 2001 From: Victor Dibia Date: Tue, 14 Jan 2025 16:14:01 -0800 Subject: [PATCH 03/10] add tests --- .../tests/test_declarative_components.py | 59 ++++++++++++++++++- .../src/autogen_core/tools/_function_tool.py | 3 +- 2 files changed, 60 insertions(+), 2 deletions(-) diff --git a/python/packages/autogen-agentchat/tests/test_declarative_components.py b/python/packages/autogen-agentchat/tests/test_declarative_components.py index 35cf54f86416..64bf362e2ede 100644 --- a/python/packages/autogen-agentchat/tests/test_declarative_components.py +++ b/python/packages/autogen-agentchat/tests/test_declarative_components.py @@ -10,7 +10,8 @@ TimeoutTermination, TokenUsageTermination, ) -from autogen_core import ComponentLoader, ComponentModel +from autogen_core import ComponentLoader, ComponentModel, CancellationToken +from autogen_core.tools import FunctionTool, ImportFromModule @pytest.mark.asyncio @@ -92,3 +93,59 @@ async def test_termination_declarative() -> None: # Test loading complex composition loaded_composite = ComponentLoader.load_component(composite_config) assert isinstance(loaded_composite, AndTerminationCondition) + + +@pytest.mark.asyncio +async def test_function_tool() -> None: + """Test FunctionTool with different function types and features.""" + + # Test sync and async functions + def sync_func(x: int, y: str) -> str: + return y * x + + async def async_func(x: float, y: float, cancellation_token: CancellationToken) -> float: + if cancellation_token.is_cancelled(): + raise Exception("Cancelled") + return x + y + + # Create tools with different configurations + sync_tool = FunctionTool( + func=sync_func, description="Multiply string", global_imports=[ImportFromModule("typing", ("Dict",))] + ) + async_tool = FunctionTool( + func=async_func, + description="Add numbers", + name="custom_adder", + global_imports=[ImportFromModule("autogen_core", ("CancellationToken",))], + ) + + # Test serialization and config + + sync_config = sync_tool.dump_component() + assert isinstance(sync_config, ComponentModel) + assert sync_config.config["name"] == "sync_func" + assert len(sync_config.config["global_imports"]) == 1 + assert not sync_config.config["has_cancellation_support"] + + async_config = async_tool.dump_component() + assert async_config.config["name"] == "custom_adder" + assert async_config.config["has_cancellation_support"] + + # Test deserialization and execution + loaded_sync = FunctionTool.load_component(sync_config, FunctionTool) + loaded_async = FunctionTool.load_component(async_config, FunctionTool) + + # Test execution and validation + token = CancellationToken() + assert await loaded_sync.run_json({"x": 2, "y": "test"}, token) == "testtest" + assert await loaded_async.run_json({"x": 1.5, "y": 2.5}, token) == 4.0 + + # Test error cases + with pytest.raises(ValueError): + # Type error + await loaded_sync.run_json({"x": "invalid", "y": "test"}, token) + + cancelled_token = CancellationToken() + cancelled_token.cancel() + with pytest.raises(Exception, match="Cancelled"): + await loaded_async.run_json({"x": 1.0, "y": 2.0}, cancelled_token) diff --git a/python/packages/autogen-core/src/autogen_core/tools/_function_tool.py b/python/packages/autogen-core/src/autogen_core/tools/_function_tool.py index 26a1b059ce99..6b861292f249 100644 --- a/python/packages/autogen-core/src/autogen_core/tools/_function_tool.py +++ b/python/packages/autogen-core/src/autogen_core/tools/_function_tool.py @@ -1,5 +1,6 @@ import asyncio import functools +from textwrap import dedent from typing import Any, Callable, Sequence from pydantic import BaseModel @@ -120,7 +121,7 @@ async def run(self, args: BaseModel, cancellation_token: CancellationToken) -> A def _to_config(self) -> FunctionToolConfig: return FunctionToolConfig( - source_code=to_code(self._func), + source_code=dedent(to_code(self._func)), global_imports=self._global_imports, name=self.name, description=self.description, From 8d4408c0d6640e8030db7e44c38b22de4d3c5c8e Mon Sep 17 00:00:00 2001 From: Victor Dibia Date: Tue, 14 Jan 2025 16:27:59 -0800 Subject: [PATCH 04/10] update imports --- .../autogen-agentchat/tests/test_declarative_components.py | 3 ++- .../packages/autogen-core/src/autogen_core/tools/__init__.py | 3 +-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/python/packages/autogen-agentchat/tests/test_declarative_components.py b/python/packages/autogen-agentchat/tests/test_declarative_components.py index 64bf362e2ede..acc4b8fd6710 100644 --- a/python/packages/autogen-agentchat/tests/test_declarative_components.py +++ b/python/packages/autogen-agentchat/tests/test_declarative_components.py @@ -11,7 +11,8 @@ TokenUsageTermination, ) from autogen_core import ComponentLoader, ComponentModel, CancellationToken -from autogen_core.tools import FunctionTool, ImportFromModule +from autogen_core.tools import FunctionTool +from autogen_core.code_executor import ImportFromModule @pytest.mark.asyncio diff --git a/python/packages/autogen-core/src/autogen_core/tools/__init__.py b/python/packages/autogen-core/src/autogen_core/tools/__init__.py index 250fd273862c..ed15ddcbe8b7 100644 --- a/python/packages/autogen-core/src/autogen_core/tools/__init__.py +++ b/python/packages/autogen-core/src/autogen_core/tools/__init__.py @@ -1,4 +1,4 @@ -from ..code_executor import ImportFromModule + from ._base import BaseTool, BaseToolWithState, ParametersSchema, Tool, ToolSchema from ._function_tool import FunctionTool @@ -9,5 +9,4 @@ "BaseTool", "BaseToolWithState", "FunctionTool", - "ImportFromModule", ] From 5bc9bc14cc9f8326ebd2b1b977edff14e814613a Mon Sep 17 00:00:00 2001 From: Victor Dibia Date: Tue, 14 Jan 2025 18:28:08 -0800 Subject: [PATCH 05/10] update formatting --- python/packages/autogen-core/src/autogen_core/tools/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/packages/autogen-core/src/autogen_core/tools/__init__.py b/python/packages/autogen-core/src/autogen_core/tools/__init__.py index ed15ddcbe8b7..52a9d725f7d6 100644 --- a/python/packages/autogen-core/src/autogen_core/tools/__init__.py +++ b/python/packages/autogen-core/src/autogen_core/tools/__init__.py @@ -1,4 +1,3 @@ - from ._base import BaseTool, BaseToolWithState, ParametersSchema, Tool, ToolSchema from ._function_tool import FunctionTool From e97129ed26b8fe6e068394ab7805dbf6f0376841 Mon Sep 17 00:00:00 2001 From: Victor Dibia Date: Tue, 14 Jan 2025 19:21:10 -0800 Subject: [PATCH 06/10] move tests, format fixes --- .../tests/test_declarative_components.py | 56 ----------------- .../tests/test_component_config.py | 60 ++++++++++++++++++- 2 files changed, 59 insertions(+), 57 deletions(-) diff --git a/python/packages/autogen-agentchat/tests/test_declarative_components.py b/python/packages/autogen-agentchat/tests/test_declarative_components.py index acc4b8fd6710..6d28753c8ef6 100644 --- a/python/packages/autogen-agentchat/tests/test_declarative_components.py +++ b/python/packages/autogen-agentchat/tests/test_declarative_components.py @@ -94,59 +94,3 @@ async def test_termination_declarative() -> None: # Test loading complex composition loaded_composite = ComponentLoader.load_component(composite_config) assert isinstance(loaded_composite, AndTerminationCondition) - - -@pytest.mark.asyncio -async def test_function_tool() -> None: - """Test FunctionTool with different function types and features.""" - - # Test sync and async functions - def sync_func(x: int, y: str) -> str: - return y * x - - async def async_func(x: float, y: float, cancellation_token: CancellationToken) -> float: - if cancellation_token.is_cancelled(): - raise Exception("Cancelled") - return x + y - - # Create tools with different configurations - sync_tool = FunctionTool( - func=sync_func, description="Multiply string", global_imports=[ImportFromModule("typing", ("Dict",))] - ) - async_tool = FunctionTool( - func=async_func, - description="Add numbers", - name="custom_adder", - global_imports=[ImportFromModule("autogen_core", ("CancellationToken",))], - ) - - # Test serialization and config - - sync_config = sync_tool.dump_component() - assert isinstance(sync_config, ComponentModel) - assert sync_config.config["name"] == "sync_func" - assert len(sync_config.config["global_imports"]) == 1 - assert not sync_config.config["has_cancellation_support"] - - async_config = async_tool.dump_component() - assert async_config.config["name"] == "custom_adder" - assert async_config.config["has_cancellation_support"] - - # Test deserialization and execution - loaded_sync = FunctionTool.load_component(sync_config, FunctionTool) - loaded_async = FunctionTool.load_component(async_config, FunctionTool) - - # Test execution and validation - token = CancellationToken() - assert await loaded_sync.run_json({"x": 2, "y": "test"}, token) == "testtest" - assert await loaded_async.run_json({"x": 1.5, "y": 2.5}, token) == 4.0 - - # Test error cases - with pytest.raises(ValueError): - # Type error - await loaded_sync.run_json({"x": "invalid", "y": "test"}, token) - - cancelled_token = CancellationToken() - cancelled_token.cancel() - with pytest.raises(Exception, match="Cancelled"): - await loaded_async.run_json({"x": 1.0, "y": 2.0}, cancelled_token) diff --git a/python/packages/autogen-core/tests/test_component_config.py b/python/packages/autogen-core/tests/test_component_config.py index d59fde59c1b6..8cae3cedce2c 100644 --- a/python/packages/autogen-core/tests/test_component_config.py +++ b/python/packages/autogen-core/tests/test_component_config.py @@ -4,12 +4,14 @@ from typing import Any, Dict import pytest -from autogen_core import Component, ComponentBase, ComponentLoader, ComponentModel +from autogen_core import Component, ComponentBase, ComponentLoader, ComponentModel, CancellationToken from autogen_core._component_config import _type_to_provider_str # type: ignore from autogen_core.models import ChatCompletionClient from autogen_test_utils import MyInnerComponent, MyOuterComponent from pydantic import BaseModel, ValidationError from typing_extensions import Self +from autogen_core.tools import FunctionTool +from autogen_core.code_executor import ImportFromModule class MyConfig(BaseModel): @@ -283,3 +285,59 @@ def test_component_version_from_dict() -> None: assert comp.info == "test" assert comp.__class__ == ComponentNonOneVersionWithUpgrade assert comp.dump_component().version == 2 + + +@pytest.mark.asyncio +async def test_function_tool() -> None: + """Test FunctionTool with different function types and features.""" + + # Test sync and async functions + def sync_func(x: int, y: str) -> str: + return y * x + + async def async_func(x: float, y: float, cancellation_token: CancellationToken) -> float: + if cancellation_token.is_cancelled(): + raise Exception("Cancelled") + return x + y + + # Create tools with different configurations + sync_tool = FunctionTool( + func=sync_func, description="Multiply string", global_imports=[ImportFromModule("typing", ("Dict",))] + ) + async_tool = FunctionTool( + func=async_func, + description="Add numbers", + name="custom_adder", + global_imports=[ImportFromModule("autogen_core", ("CancellationToken",))], + ) + + # Test serialization and config + + sync_config = sync_tool.dump_component() + assert isinstance(sync_config, ComponentModel) + assert sync_config.config["name"] == "sync_func" + assert len(sync_config.config["global_imports"]) == 1 + assert not sync_config.config["has_cancellation_support"] + + async_config = async_tool.dump_component() + assert async_config.config["name"] == "custom_adder" + assert async_config.config["has_cancellation_support"] + + # Test deserialization and execution + loaded_sync = FunctionTool.load_component(sync_config, FunctionTool) + loaded_async = FunctionTool.load_component(async_config, FunctionTool) + + # Test execution and validation + token = CancellationToken() + assert await loaded_sync.run_json({"x": 2, "y": "test"}, token) == "testtest" + assert await loaded_async.run_json({"x": 1.5, "y": 2.5}, token) == 4.0 + + # Test error cases + with pytest.raises(ValueError): + # Type error + await loaded_sync.run_json({"x": "invalid", "y": "test"}, token) + + cancelled_token = CancellationToken() + cancelled_token.cancel() + with pytest.raises(Exception, match="Cancelled"): + await loaded_async.run_json({"x": 1.0, "y": 2.0}, cancelled_token) From 036a1353c3cea1c018b111a96d8cde66bd0f8b4b Mon Sep 17 00:00:00 2001 From: Victor Dibia Date: Tue, 14 Jan 2025 19:44:44 -0800 Subject: [PATCH 07/10] format updates --- .../autogen-agentchat/tests/test_declarative_components.py | 4 +--- python/packages/autogen-core/tests/test_component_config.py | 6 +++--- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/python/packages/autogen-agentchat/tests/test_declarative_components.py b/python/packages/autogen-agentchat/tests/test_declarative_components.py index 6d28753c8ef6..35cf54f86416 100644 --- a/python/packages/autogen-agentchat/tests/test_declarative_components.py +++ b/python/packages/autogen-agentchat/tests/test_declarative_components.py @@ -10,9 +10,7 @@ TimeoutTermination, TokenUsageTermination, ) -from autogen_core import ComponentLoader, ComponentModel, CancellationToken -from autogen_core.tools import FunctionTool -from autogen_core.code_executor import ImportFromModule +from autogen_core import ComponentLoader, ComponentModel @pytest.mark.asyncio diff --git a/python/packages/autogen-core/tests/test_component_config.py b/python/packages/autogen-core/tests/test_component_config.py index 8cae3cedce2c..4071638117a7 100644 --- a/python/packages/autogen-core/tests/test_component_config.py +++ b/python/packages/autogen-core/tests/test_component_config.py @@ -4,14 +4,14 @@ from typing import Any, Dict import pytest -from autogen_core import Component, ComponentBase, ComponentLoader, ComponentModel, CancellationToken +from autogen_core import CancellationToken, Component, ComponentBase, ComponentLoader, ComponentModel from autogen_core._component_config import _type_to_provider_str # type: ignore +from autogen_core.code_executor import ImportFromModule from autogen_core.models import ChatCompletionClient +from autogen_core.tools import FunctionTool from autogen_test_utils import MyInnerComponent, MyOuterComponent from pydantic import BaseModel, ValidationError from typing_extensions import Self -from autogen_core.tools import FunctionTool -from autogen_core.code_executor import ImportFromModule class MyConfig(BaseModel): From 78fdfede7c8335f18993efae6d21a8a985f17905 Mon Sep 17 00:00:00 2001 From: Victor Dibia Date: Tue, 14 Jan 2025 20:45:27 -0800 Subject: [PATCH 08/10] update test --- .../packages/autogen-core/tests/test_component_config.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/python/packages/autogen-core/tests/test_component_config.py b/python/packages/autogen-core/tests/test_component_config.py index 4071638117a7..1f78e907a447 100644 --- a/python/packages/autogen-core/tests/test_component_config.py +++ b/python/packages/autogen-core/tests/test_component_config.py @@ -304,6 +304,15 @@ async def async_func(x: float, y: float, cancellation_token: CancellationToken) sync_tool = FunctionTool( func=sync_func, description="Multiply string", global_imports=[ImportFromModule("typing", ("Dict",))] ) + invalid_import_sync_tool = FunctionTool( + func=sync_func, description="Multiply string", global_imports=[ImportFromModule("invalid_module (", ("Dict",))] + ) + + invalid_import_config = invalid_import_sync_tool.dump_component() + # check that invalid import raises an error + with pytest.raises(RuntimeError): + _ = FunctionTool.load_component(invalid_import_config, FunctionTool) + async_tool = FunctionTool( func=async_func, description="Add numbers", From 1374eefb94701effdc886c03ad306868d1f3b9a9 Mon Sep 17 00:00:00 2001 From: Victor Dibia Date: Thu, 23 Jan 2025 08:53:07 -0800 Subject: [PATCH 09/10] add user warning to _from_config --- .../src/autogen_core/tools/_function_tool.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/python/packages/autogen-core/src/autogen_core/tools/_function_tool.py b/python/packages/autogen-core/src/autogen_core/tools/_function_tool.py index 6b861292f249..b43e061350e5 100644 --- a/python/packages/autogen-core/src/autogen_core/tools/_function_tool.py +++ b/python/packages/autogen-core/src/autogen_core/tools/_function_tool.py @@ -2,6 +2,7 @@ import functools from textwrap import dedent from typing import Any, Callable, Sequence +import warnings from pydantic import BaseModel from typing_extensions import Self @@ -130,6 +131,14 @@ def _to_config(self) -> FunctionToolConfig: @classmethod def _from_config(cls, config: FunctionToolConfig) -> Self: + warnings.warn( + "\n⚠️ SECURITY WARNING ⚠️\n" + "Loading a FunctionTool from config will execute code to import the provided global imports and and function code.\n" + "Only load configs from TRUSTED sources to prevent arbitrary code execution.", + UserWarning, + stacklevel=2, + ) + exec_globals: dict[str, Any] = {} # Execute imports first From b351529e523c86c6b24bda13267d24a8bca86249 Mon Sep 17 00:00:00 2001 From: Victor Dibia Date: Thu, 23 Jan 2025 09:32:19 -0800 Subject: [PATCH 10/10] add warning on load_component to docs --- .../serialize-components.ipynb | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/serialize-components.ipynb b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/serialize-components.ipynb index 5a3855f48080..ff29efa9100a 100644 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/serialize-components.ipynb +++ b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/serialize-components.ipynb @@ -6,21 +6,20 @@ "source": [ "# Serializing Components \n", "\n", - "AutoGen provides a {py:class}`~autogen_core.Component` configuration class that defines behaviours for to serialize/deserialize component into declarative specifications. This is useful for debugging, visualizing, and even for sharing your work with others. In this notebook, we will demonstrate how to serialize multiple components to a declarative specification like a JSON file. \n", + "AutoGen provides a {py:class}`~autogen_core.Component` configuration class that defines behaviours to serialize/deserialize component into declarative specifications. We can accomplish this by calling `.dump_component()` and `.load_component()` respectively. This is useful for debugging, visualizing, and even for sharing your work with others. In this notebook, we will demonstrate how to serialize multiple components to a declarative specification like a JSON file. \n", "\n", "\n", - "```{note}\n", - "This is work in progress\n", - "``` \n", + "```{warning}\n", "\n", - "We will be implementing declarative support for the following components:\n", + "ONLY LOAD COMPONENTS FROM TRUSTED SOURCES.\n", "\n", - "- Termination conditions ✔️\n", - "- Tools \n", - "- Agents \n", - "- Teams \n", + "With serilized components, each component implements the logic for how it is serialized and deserialized - i.e., how the declarative specification is generated and how it is converted back to an object. \n", "\n", + "In some cases, creating an object may include executing code (e.g., a serialized function). ONLY LOAD COMPONENTS FROM TRUSTED SOURCES. \n", + " \n", + "```\n", "\n", + " \n", "### Termination Condition Example \n", "\n", "In the example below, we will define termination conditions (a part of an agent team) in python, export this to a dictionary/json and also demonstrate how the termination condition object can be loaded from the dictionary/json. \n",