From 1fb62483afc2a6a437e36b9355203cde01227827 Mon Sep 17 00:00:00 2001
From: AgentGenie <peterag2ai@gmail.com>
Date: Mon, 20 Jan 2025 12:55:37 -0800
Subject: [PATCH] Tools capability (#526)

* Add tools capability

* Add unit test

* Change to use Tool interface

* Update

* Added to docstring and test

Signed-off-by: Mark Sze <mark@sze.family>

---------

Signed-off-by: Mark Sze <mark@sze.family>
Co-authored-by: Mark Sze <mark@sze.family>
---
 .../contrib/capabilities/tools_capability.py  | 24 +++++++++++
 autogen/tools/tool.py                         | 15 +++++++
 .../capabilities/test_tools_capability.py     | 43 +++++++++++++++++++
 3 files changed, 82 insertions(+)
 create mode 100644 autogen/agentchat/contrib/capabilities/tools_capability.py
 create mode 100644 test/agentchat/contrib/capabilities/test_tools_capability.py

diff --git a/autogen/agentchat/contrib/capabilities/tools_capability.py b/autogen/agentchat/contrib/capabilities/tools_capability.py
new file mode 100644
index 0000000000..b3bf13b138
--- /dev/null
+++ b/autogen/agentchat/contrib/capabilities/tools_capability.py
@@ -0,0 +1,24 @@
+# Copyright (c) 2023 - 2025, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from autogen.agentchat import ConversableAgent
+from autogen.tools import Tool
+
+
+class ToolsCapability:
+    """Adding a list of tools as composable capabilities to a single agent.
+    This class can be inherited from to allow code to run at the point of creating or adding the capability.
+
+    Note: both caller and executor of the tools are the same agent.
+    """
+
+    def __init__(self, tool_list: list[Tool]):
+        self.tools = [tool for tool in tool_list]
+
+    def add_to_agent(self, agent: ConversableAgent):
+        """
+        Add tools to the given agent.
+        """
+        for tool in self.tools:
+            tool.register_tool(agent=agent)
diff --git a/autogen/tools/tool.py b/autogen/tools/tool.py
index d5391f1fd4..245b9ce76a 100644
--- a/autogen/tools/tool.py
+++ b/autogen/tools/tool.py
@@ -89,6 +89,21 @@ def register_for_execution(self, agent: "ConversableAgent") -> None:
         """
         agent.register_for_execution()(self)
 
+    def register_tool(self, agent: "ConversableAgent") -> None:
+        """
+        Register a tool to be both proposed and executed by an agent.
+
+        Equivalent to calling both `register_for_llm` and `register_for_execution` with the same agent.
+
+        Note: This will not make the agent recommend and execute the call in the one step. If the agent
+        recommends the tool, it will need to be the next agent to speak in order to execute the tool.
+
+        Args:
+            agent (ConversableAgent): The agent to which the tool will be registered.
+        """
+        self.register_for_llm(agent)
+        self.register_for_execution(agent)
+
     def __call__(self, *args: Any, **kwargs: Any) -> Any:
         """Execute the tool by calling its underlying function with the provided arguments.
 
diff --git a/test/agentchat/contrib/capabilities/test_tools_capability.py b/test/agentchat/contrib/capabilities/test_tools_capability.py
new file mode 100644
index 0000000000..c30ee73182
--- /dev/null
+++ b/test/agentchat/contrib/capabilities/test_tools_capability.py
@@ -0,0 +1,43 @@
+# 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 ToolsCapability
+from autogen.tools import Tool
+
+
+@pytest.fixture
+def add_tools():
+    def add(x: int, y: int) -> int:
+        return x + y
+
+    return Tool(
+        name="add_function",
+        description="Provide add function to two argument and return sum.",
+        func_or_tool=add,
+    )
+
+
+@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_tools, test_agent) -> None:
+        # Arrange
+        tools_capability = ToolsCapability(tool_list=[add_tools])
+        assert "tools" not in test_agent.llm_config
+        # Act
+        tools_capability.add_to_agent(agent=test_agent)
+        # Assert that the tool was added for LLM and Execution
+        assert len(test_agent.llm_config["tools"]) == 1  # LLM
+        assert len(test_agent.function_map) == 1  # Execution