Add unit test

autogen/agentchat/contrib/capabilities/tools_capability.py

@@ -4,20 +4,20 @@
 from dataclasses import dataclass
 from typing import Any, Callable, Optional
 
-from autogen.agentchat.assistant_agent import ConversableAgent
+from autogen.agentchat import ConversableAgent, register_function
 
 
 @dataclass
 class ToolSpecs:
     tool_func: Callable[..., Any]
-    caller: ConversableAgent
-    executor: ConversableAgent
     tool_description: str
     tool_name: Optional[str] = None
 
 
 class ToolsCapability:
-    """Adding a list of tools as composable capabilities to an agent."""
+    """Adding a list of tools as composable capabilities to a single agent.
+    Note: both caller and executor of the tools are the same agent.
+    """
 
     def __init__(self, tool_list: list[ToolSpecs]):
         self.tools = [specs for specs in tool_list]
@@ -27,10 +27,10 @@ def add_to_agent(self, agent: ConversableAgent):
         Add tools to the given agent.
         """
         for specs in self.tools:
-            agent.register_function(
+            register_function(
                 f=specs.tool_func,
-                caller=specs.caller,
-                executor=specs.executor,
-                tool_description=specs.tool_description,
-                tool_name=specs.tool_name,
+                caller=agent,
+                executor=agent,
+                description=specs.tool_description,
+                name=specs.tool_name,
             )

test/agentchat/contrib/capabilities/test_tools_capability.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+
+from autogen import AssistantAgent
+from autogen.agentchat.contrib.capabilities.tools_capability import ToolSpecs, ToolsCapability
+
+
+@pytest.fixture
+def add_tool_spec():
+    def add(x: int, y: int) -> int:
+        return x + y
+
+    return ToolSpecs(
+        tool_func=add,
+        tool_description="Provide add function to two argument and return sum.",
+        tool_name="add_function",
+    )
+
+
+@pytest.fixture
+def test_agent():
+    return AssistantAgent(
+        name="test_agent",
+        llm_config={
+            "config_list": [{"model": "gpt-4O", "api_key": "sk-proj-ABC"}],
+        },
+    )
+
+
+class TestToolsCapability:
+    def test_add_capability(self, add_tool_spec, test_agent) -> None:
+        # Arrange
+        tools_capability = ToolsCapability(tool_list=[add_tool_spec])
+        assert "tools" not in test_agent.llm_config
+        # Act
+        tools_capability.add_to_agent(agent=test_agent)
+        # Assert
+        assert len(test_agent.llm_config["tools"]) == 1