feat: Add Toolset to Agent (#9284)

* Add Toolset to Agent

* Add reno

Author: sjrl

haystack/components/agents/agent.py

@@ -4,7 +4,7 @@
 
 import inspect
 from copy import deepcopy
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from haystack import component, default_from_dict, default_to_dict, logging, tracing
 from haystack.components.generators.chat.types import ChatGenerator
@@ -16,7 +16,7 @@
 from haystack.dataclasses.state import State, _schema_from_dict, _schema_to_dict, _validate_schema
 from haystack.dataclasses.state_utils import merge_lists
 from haystack.dataclasses.streaming_chunk import StreamingCallbackT
-from haystack.tools import Tool, deserialize_tools_or_toolset_inplace
+from haystack.tools import Tool, Toolset, deserialize_tools_or_toolset_inplace, serialize_tools_or_toolset
 from haystack.utils.callable_serialization import deserialize_callable, serialize_callable
 from haystack.utils.deserialization import deserialize_chatgenerator_inplace
 
@@ -61,7 +61,7 @@ def __init__(
         self,
         *,
         chat_generator: ChatGenerator,
-        tools: Optional[List[Tool]] = None,
+        tools: Optional[Union[List[Tool], Toolset]] = None,
         system_prompt: Optional[str] = None,
         exit_conditions: Optional[List[str]] = None,
         state_schema: Optional[Dict[str, Any]] = None,
@@ -73,7 +73,7 @@ def __init__(
         Initialize the agent component.
 
         :param chat_generator: An instance of the chat generator that your agent should use. It must support tools.
-        :param tools: List of Tool objects available to the agent
+        :param tools: List of Tool objects or a Toolset that the agent can use.
         :param system_prompt: System prompt for the agent.
         :param exit_conditions: List of conditions that will cause the agent to return.
             Can include "text" if the agent should return when it generates a message without tool calls,
@@ -166,7 +166,7 @@ def to_dict(self) -> Dict[str, Any]:
         return default_to_dict(
             self,
             chat_generator=component_to_dict(obj=self.chat_generator, name="chat_generator"),
-            tools=[t.to_dict() for t in self.tools],
+            tools=serialize_tools_or_toolset(self.tools),
             system_prompt=self.system_prompt,
             exit_conditions=self.exit_conditions,
             # We serialize the original state schema, not the resolved one to reflect the original user input

haystack/components/tools/tool_invoker.py

@@ -184,15 +184,17 @@ def __init__(
 
         # Convert Toolset to list for internal use
         if isinstance(tools, Toolset):
-            tools = list(tools)
+            converted_tools = list(tools)
+        else:
+            converted_tools = tools
 
-        _check_duplicate_tool_names(tools)
-        tool_names = [tool.name for tool in tools]
+        _check_duplicate_tool_names(converted_tools)
+        tool_names = [tool.name for tool in converted_tools]
         duplicates = {name for name in tool_names if tool_names.count(name) > 1}
         if duplicates:
             raise ValueError(f"Duplicate tool names found: {duplicates}")
 
-        self._tools_with_names = dict(zip(tool_names, tools))
+        self._tools_with_names = dict(zip(tool_names, converted_tools))
         self.raise_on_failure = raise_on_failure
         self.convert_result_to_json_string = convert_result_to_json_string
 

releasenotes/notes/add-toolset-to-agent-d4a14822c4aed001.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    Agent now supports a List of Tools or a Toolset as input.

test/components/agents/test_agent.py

@@ -239,6 +239,62 @@ def test_to_dict(self, weather_tool, component_tool, monkeypatch):
             },
         }
 
+    def test_to_dict_with_toolset(self, monkeypatch, weather_tool):
+        monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
+        toolset = Toolset(tools=[weather_tool])
+        agent = Agent(chat_generator=OpenAIChatGenerator(), tools=toolset)
+        serialized_agent = agent.to_dict()
+        assert serialized_agent == {
+            "type": "haystack.components.agents.agent.Agent",
+            "init_parameters": {
+                "chat_generator": {
+                    "type": "haystack.components.generators.chat.openai.OpenAIChatGenerator",
+                    "init_parameters": {
+                        "model": "gpt-4o-mini",
+                        "streaming_callback": None,
+                        "api_base_url": None,
+                        "organization": None,
+                        "generation_kwargs": {},
+                        "api_key": {"type": "env_var", "env_vars": ["OPENAI_API_KEY"], "strict": True},
+                        "timeout": None,
+                        "max_retries": None,
+                        "tools": None,
+                        "tools_strict": False,
+                        "http_client_kwargs": None,
+                    },
+                },
+                "tools": {
+                    "type": "haystack.tools.toolset.Toolset",
+                    "data": {
+                        "tools": [
+                            {
+                                "type": "haystack.tools.tool.Tool",
+                                "data": {
+                                    "name": "weather_tool",
+                                    "description": "Provides weather information for a given location.",
+                                    "parameters": {
+                                        "type": "object",
+                                        "properties": {"location": {"type": "string"}},
+                                        "required": ["location"],
+                                    },
+                                    "function": "test_agent.weather_function",
+                                    "outputs_to_string": None,
+                                    "inputs_from_state": None,
+                                    "outputs_to_state": None,
+                                },
+                            }
+                        ]
+                    },
+                },
+                "system_prompt": None,
+                "exit_conditions": ["text"],
+                "state_schema": {},
+                "max_agent_steps": 100,
+                "raise_on_tool_invocation_failure": False,
+                "streaming_callback": None,
+            },
+        }
+
     def test_from_dict(self, weather_tool, component_tool, monkeypatch):
         monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
         data = {
@@ -318,6 +374,67 @@ def test_from_dict(self, weather_tool, component_tool, monkeypatch):
             "messages": {"handler": merge_lists, "type": List[ChatMessage]},
         }
 
+    def test_from_dict_with_toolset(self, monkeypatch):
+        monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
+        data = {
+            "type": "haystack.components.agents.agent.Agent",
+            "init_parameters": {
+                "chat_generator": {
+                    "type": "haystack.components.generators.chat.openai.OpenAIChatGenerator",
+                    "init_parameters": {
+                        "model": "gpt-4o-mini",
+                        "streaming_callback": None,
+                        "api_base_url": None,
+                        "organization": None,
+                        "generation_kwargs": {},
+                        "api_key": {"type": "env_var", "env_vars": ["OPENAI_API_KEY"], "strict": True},
+                        "timeout": None,
+                        "max_retries": None,
+                        "tools": None,
+                        "tools_strict": False,
+                        "http_client_kwargs": None,
+                    },
+                },
+                "tools": {
+                    "type": "haystack.tools.toolset.Toolset",
+                    "data": {
+                        "tools": [
+                            {
+                                "type": "haystack.tools.tool.Tool",
+                                "data": {
+                                    "name": "weather_tool",
+                                    "description": "Provides weather information for a given location.",
+                                    "parameters": {
+                                        "type": "object",
+                                        "properties": {"location": {"type": "string"}},
+                                        "required": ["location"],
+                                    },
+                                    "function": "test_agent.weather_function",
+                                    "outputs_to_string": None,
+                                    "inputs_from_state": None,
+                                    "outputs_to_state": None,
+                                },
+                            }
+                        ]
+                    },
+                },
+                "system_prompt": None,
+                "exit_conditions": ["text"],
+                "state_schema": {},
+                "max_agent_steps": 100,
+                "raise_on_tool_invocation_failure": False,
+                "streaming_callback": None,
+            },
+        }
+        agent = Agent.from_dict(data)
+        assert isinstance(agent, Agent)
+        assert isinstance(agent.chat_generator, OpenAIChatGenerator)
+        assert agent.chat_generator.model == "gpt-4o-mini"
+        assert agent.chat_generator.api_key == Secret.from_env_var("OPENAI_API_KEY")
+        assert isinstance(agent.tools, Toolset)
+        assert agent.tools[0].function is weather_function
+        assert agent.exit_conditions == ["text"]
+
     def test_serde(self, weather_tool, component_tool, monkeypatch):
         monkeypatch.setenv("FAKE_OPENAI_KEY", "fake-key")
         generator = OpenAIChatGenerator(api_key=Secret.from_env_var("FAKE_OPENAI_KEY"))

test/components/tools/test_tool_invoker.py

@@ -9,7 +9,7 @@
 from haystack.components.tools.tool_invoker import ToolInvoker, ToolNotFoundException, StringConversionError
 from haystack.dataclasses import ChatMessage, ToolCall, ToolCallResult, ChatRole
 from haystack.dataclasses.state import State
-from haystack.tools import ComponentTool, Tool
+from haystack.tools import ComponentTool, Tool, Toolset
 from haystack.tools.errors import ToolInvocationError
 
 
@@ -54,6 +54,34 @@ def faulty_tool_func(location):
     )
 
 
+def add_function(num1: int, num2: int):
+    return num1 + num2
+
+
+@pytest.fixture
+def tool_set():
+    return Toolset(
+        tools=[
+            Tool(
+                name="weather_tool",
+                description="Provides weather information for a given location.",
+                parameters=weather_parameters,
+                function=weather_function,
+            ),
+            Tool(
+                name="addition_tool",
+                description="A tool that adds two numbers.",
+                parameters={
+                    "type": "object",
+                    "properties": {"num1": {"type": "integer"}, "num2": {"type": "integer"}},
+                    "required": ["num1", "num2"],
+                },
+                function=add_function,
+            ),
+        ]
+    )
+
+
 @pytest.fixture
 def invoker(weather_tool):
     return ToolInvoker(tools=[weather_tool], raise_on_failure=True, convert_result_to_json_string=False)
@@ -73,6 +101,11 @@ def test_init(self, weather_tool):
         assert invoker.raise_on_failure
         assert not invoker.convert_result_to_json_string
 
+    def test_init_with_toolset(self, tool_set):
+        tool_invoker = ToolInvoker(tools=tool_set)
+        assert tool_invoker.tools == tool_set
+        assert tool_invoker._tools_with_names == {"weather_tool": tool_set.tools[0], "addition_tool": tool_set.tools[1]}
+
     def test_init_fails_wo_tools(self):
         with pytest.raises(ValueError):
             ToolInvoker(tools=[])
@@ -149,6 +182,26 @@ def test_run(self, invoker):
         assert tool_call_result.origin == tool_call
         assert not tool_call_result.error
 
+    def test_run_with_toolset(self, tool_set):
+        tool_invoker = ToolInvoker(tools=tool_set, raise_on_failure=True, convert_result_to_json_string=False)
+        tool_call = ToolCall(tool_name="addition_tool", arguments={"num1": 5, "num2": 3})
+        message = ChatMessage.from_assistant(tool_calls=[tool_call])
+
+        result = tool_invoker.run(messages=[message])
+        assert "tool_messages" in result
+        assert len(result["tool_messages"]) == 1
+
+        tool_message = result["tool_messages"][0]
+        assert isinstance(tool_message, ChatMessage)
+        assert tool_message.is_from(ChatRole.TOOL)
+        assert tool_message.tool_call_results
+
+        tool_call_result = tool_message.tool_call_result
+        assert isinstance(tool_call_result, ToolCallResult)
+        assert tool_call_result.result == str(8)
+        assert tool_call_result.origin == tool_call
+        assert not tool_call_result.error
+
     def test_run_no_messages(self, invoker):
         result = invoker.run(messages=[])
         assert result["tool_messages"] == []