fix: warm up individual tools inside Toolsets in warm_up_tools()

haystack/tools/toolset.py

@@ -189,10 +189,30 @@ def warm_up(self) -> None:
         """
         Prepare the Toolset for use.
 
-        Override this method to set up shared resources like database connections or HTTP sessions.
+        By default, this method iterates through and warms up all tools in the Toolset.
+        Subclasses can override this method to customize initialization behavior, such as:
+
+        - Setting up shared resources (database connections, HTTP sessions) instead of
+          warming individual tools
+        - Implementing custom initialization logic for dynamically loaded tools
+        - Controlling when and how tools are initialized
+
+        For example, a Toolset that manages tools from an external service (like MCPToolset)
+        might override this to initialize a shared connection rather than warming up
+        individual tools:
+
+        ```python
+        class MCPToolset(Toolset):
+            def warm_up(self) -> None:
+                # Only warm up the shared MCP connection, not individual tools
+                self.mcp_connection = establish_connection(self.server_url)
+        ```
+
         This method should be idempotent, as it may be called multiple times.
         """
-        pass
+        for tool in self.tools:
+            if hasattr(tool, "warm_up"):
+                tool.warm_up()
 
     def add(self, tool: Union[Tool, "Toolset"]) -> None:
         """

haystack/tools/utils.py

@@ -15,21 +15,25 @@ def warm_up_tools(tools: "Optional[ToolsType]" = None) -> None:
     """
     Warm up tools from various formats (Tools, Toolsets, or mixed lists).
 
+    For Toolset objects, this delegates to Toolset.warm_up(), which by default
+    warms up all tools in the Toolset. Toolset subclasses can override warm_up()
+    to customize initialization behavior (e.g., setting up shared resources).
+
     :param tools: A list of Tool and/or Toolset objects, a single Toolset, or None.
     """
     if tools is None:
         return
 
-    # If tools is a single Toolset, warm up the toolset itself
-    if isinstance(tools, Toolset):
+    # If tools is a single Toolset or Tool, warm it up
+    if isinstance(tools, (Toolset, Tool)):
         if hasattr(tools, "warm_up"):
             tools.warm_up()
         return
 
     # If tools is a list, warm up each item (Tool or Toolset)
     if isinstance(tools, list):
         for item in tools:
-            if isinstance(item, (Toolset, Tool)) and hasattr(item, "warm_up"):
+            if hasattr(item, "warm_up"):
                 item.warm_up()
 
 

releasenotes/notes/fix-toolset-warm-up-3e985e6a5e95a55d.yaml

@@ -0,0 +1,7 @@
+enhancements:
+  - |
+    Improved Toolset warm-up architecture for better encapsulation. The base
+    Toolset.warm_up() method now warms up all tools by default, while subclasses
+    can override it to customize initialization (e.g., setting up shared resources
+    instead of warming individual tools). The warm_up_tools() utility function has
+    been simplified to delegate to Toolset.warm_up().

test/components/tools/test_tool_invoker.py

@@ -136,6 +136,8 @@ def __init__(self, tools):
 
     def warm_up(self):
         self.was_warmed_up = True
+        # Call parent to warm up individual tools
+        super().warm_up()
 
 
 class TestToolInvokerCore:

test/tools/test_tools_utils.py

@@ -4,7 +4,7 @@
 
 import pytest
 
-from haystack.tools import Tool, Toolset, flatten_tools_or_toolsets
+from haystack.tools import Tool, Toolset, flatten_tools_or_toolsets, warm_up_tools
 
 
 def add_numbers(a: int, b: int) -> int:
@@ -171,3 +171,225 @@ def test_flatten_multiple_toolsets(self, add_tool, multiply_tool, subtract_tool)
         assert result[0].name == "add"
         assert result[1].name == "multiply"
         assert result[2].name == "subtract"
+
+
+class WarmupTrackingTool(Tool):
+    """A tool that tracks whether warm_up was called."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.was_warmed_up = False
+
+    def warm_up(self):
+        self.was_warmed_up = True
+
+
+class WarmupTrackingToolset(Toolset):
+    """A toolset that tracks whether warm_up was called."""
+
+    def __init__(self, tools):
+        super().__init__(tools)
+        self.was_warmed_up = False
+
+    def warm_up(self):
+        self.was_warmed_up = True
+        # Call parent to warm up individual tools
+        super().warm_up()
+
+
+class TestWarmUpTools:
+    """Tests for the warm_up_tools() function"""
+
+    def test_warm_up_tools_with_none(self):
+        """Test that warm_up_tools with None does nothing."""
+        # Should not raise any errors
+        warm_up_tools(None)
+
+    def test_warm_up_tools_with_single_tool(self):
+        """Test that warm_up_tools works with a single tool in a list."""
+        tool = WarmupTrackingTool(
+            name="test_tool",
+            description="A test tool",
+            parameters={"type": "object", "properties": {}},
+            function=lambda: "test",
+        )
+
+        assert not tool.was_warmed_up
+        warm_up_tools([tool])
+        assert tool.was_warmed_up
+
+    def test_warm_up_tools_with_single_toolset(self):
+        """
+        Test that when passing a single Toolset, both the Toolset.warm_up()
+        and each individual tool's warm_up() are called.
+        """
+        tool1 = WarmupTrackingTool(
+            name="tool1",
+            description="First tool",
+            parameters={"type": "object", "properties": {}},
+            function=lambda: "tool1",
+        )
+        tool2 = WarmupTrackingTool(
+            name="tool2",
+            description="Second tool",
+            parameters={"type": "object", "properties": {}},
+            function=lambda: "tool2",
+        )
+
+        toolset = WarmupTrackingToolset([tool1, tool2])
+
+        assert not toolset.was_warmed_up
+        assert not tool1.was_warmed_up
+        assert not tool2.was_warmed_up
+
+        warm_up_tools(toolset)
+
+        # Both the toolset itself and individual tools should be warmed up
+        assert toolset.was_warmed_up
+        assert tool1.was_warmed_up
+        assert tool2.was_warmed_up
+
+    def test_warm_up_tools_with_list_containing_toolset(self):
+        """Test that when a Toolset is in a list, individual tools inside get warmed up."""
+        tool1 = WarmupTrackingTool(
+            name="tool1",
+            description="First tool",
+            parameters={"type": "object", "properties": {}},
+            function=lambda: "tool1",
+        )
+        tool2 = WarmupTrackingTool(
+            name="tool2",
+            description="Second tool",
+            parameters={"type": "object", "properties": {}},
+            function=lambda: "tool2",
+        )
+
+        toolset = WarmupTrackingToolset([tool1, tool2])
+
+        assert not toolset.was_warmed_up
+        assert not tool1.was_warmed_up
+        assert not tool2.was_warmed_up
+
+        warm_up_tools([toolset])
+
+        # Both the toolset itself and individual tools should be warmed up
+        assert toolset.was_warmed_up
+        assert tool1.was_warmed_up
+        assert tool2.was_warmed_up
+
+    def test_warm_up_tools_with_multiple_toolsets(self):
+        """Test multiple Toolsets in a list."""
+        tool1 = WarmupTrackingTool(
+            name="tool1",
+            description="First tool",
+            parameters={"type": "object", "properties": {}},
+            function=lambda: "tool1",
+        )
+        tool2 = WarmupTrackingTool(
+            name="tool2",
+            description="Second tool",
+            parameters={"type": "object", "properties": {}},
+            function=lambda: "tool2",
+        )
+        tool3 = WarmupTrackingTool(
+            name="tool3",
+            description="Third tool",
+            parameters={"type": "object", "properties": {}},
+            function=lambda: "tool3",
+        )
+
+        toolset1 = WarmupTrackingToolset([tool1])
+        toolset2 = WarmupTrackingToolset([tool2, tool3])
+
+        assert not toolset1.was_warmed_up
+        assert not toolset2.was_warmed_up
+        assert not tool1.was_warmed_up
+        assert not tool2.was_warmed_up
+        assert not tool3.was_warmed_up
+
+        warm_up_tools([toolset1, toolset2])
+
+        # Both toolsets and all individual tools should be warmed up
+        assert toolset1.was_warmed_up
+        assert toolset2.was_warmed_up
+        assert tool1.was_warmed_up
+        assert tool2.was_warmed_up
+        assert tool3.was_warmed_up
+
+    def test_warm_up_tools_with_mixed_tools_and_toolsets(self):
+        """Test list with both Tool objects and Toolsets."""
+        standalone_tool = WarmupTrackingTool(
+            name="standalone",
+            description="Standalone tool",
+            parameters={"type": "object", "properties": {}},
+            function=lambda: "standalone",
+        )
+        toolset_tool1 = WarmupTrackingTool(
+            name="toolset_tool1",
+            description="Tool in toolset",
+            parameters={"type": "object", "properties": {}},
+            function=lambda: "toolset_tool1",
+        )
+        toolset_tool2 = WarmupTrackingTool(
+            name="toolset_tool2",
+            description="Another tool in toolset",
+            parameters={"type": "object", "properties": {}},
+            function=lambda: "toolset_tool2",
+        )
+
+        toolset = WarmupTrackingToolset([toolset_tool1, toolset_tool2])
+
+        assert not standalone_tool.was_warmed_up
+        assert not toolset.was_warmed_up
+        assert not toolset_tool1.was_warmed_up
+        assert not toolset_tool2.was_warmed_up
+
+        warm_up_tools([standalone_tool, toolset])
+
+        # All tools and the toolset should be warmed up
+        assert standalone_tool.was_warmed_up
+        assert toolset.was_warmed_up
+        assert toolset_tool1.was_warmed_up
+        assert toolset_tool2.was_warmed_up
+
+    def test_warm_up_tools_idempotency(self):
+        """Test that calling warm_up_tools() multiple times is safe."""
+
+        class WarmupCountingTool(Tool):
+            """A tool that counts how many times warm_up was called."""
+
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+                self.warm_up_count = 0
+
+            def warm_up(self):
+                self.warm_up_count += 1
+
+        class WarmupCountingToolset(Toolset):
+            """A toolset that counts how many times warm_up was called."""
+
+            def __init__(self, tools):
+                super().__init__(tools)
+                self.warm_up_count = 0
+
+            def warm_up(self):
+                self.warm_up_count += 1
+                super().warm_up()  # Also warm up individual tools
+
+        tool = WarmupCountingTool(
+            name="counting_tool",
+            description="A counting tool",
+            parameters={"type": "object", "properties": {}},
+            function=lambda: "test",
+        )
+        toolset = WarmupCountingToolset([tool])
+
+        # Call warm_up_tools multiple times
+        warm_up_tools(toolset)
+        warm_up_tools(toolset)
+        warm_up_tools(toolset)
+
+        # warm_up_tools itself doesn't prevent multiple calls,
+        # but verify the calls actually happen multiple times
+        assert toolset.warm_up_count == 3
+        assert tool.warm_up_count == 3