Add Tools warm_up (#9856)

* Tools warmup initial

* Fix lint

* Improve pydocs for warm_up

* Further improve pydocs for warm_up

* No need to warm_up tools in Agent as they are warmed up by ToolInvoker

* Simplify Toolset __add__ logic

* Simplify _ToolsetWrapper

* Add unit tests

* ToolInvoker warm_up

* Improve Tool pydoc

* Resurrect serde_utils.py

* Update tests

* Call ToolInvoker warm_up in agent warm_up

* Lint

* Move warm_up tests to ToolInvoker

* Update tests

* Remove tests

* Pydoc nit

* PR feedback

* ToolInvoker's warm_up is idempotent

* Add reno note

* Update releasenotes/notes/tools-warm-up-support-e16cc043fed3653f.yaml

Co-authored-by: Sebastian Husch Lee <[email protected]>

* Make ComponentTool warm_up idempotent

* Update warm_up_tools to use ToolsType

* Linting

* Add warm up test for mixed list of Tool/Toolset instances

---------

Co-authored-by: Sebastian Husch Lee <[email protected]>

Author: vblagoje

haystack/components/agents/agent.py

@@ -203,6 +203,8 @@ def warm_up(self) -> None:
         if not self._is_warmed_up:
             if hasattr(self.chat_generator, "warm_up"):
                 self.chat_generator.warm_up()
+            if hasattr(self._tool_invoker, "warm_up") and self._tool_invoker is not None:
+                self._tool_invoker.warm_up()
             self._is_warmed_up = True
 
     def to_dict(self) -> dict[str, Any]:

haystack/components/tools/tool_invoker.py

@@ -24,6 +24,7 @@
     deserialize_tools_or_toolset_inplace,
     flatten_tools_or_toolsets,
     serialize_tools_or_toolset,
+    warm_up_tools,
 )
 from haystack.tools.errors import ToolInvocationError
 from haystack.tracing.utils import _serializable_value
@@ -216,6 +217,7 @@ def __init__(
         self.convert_result_to_json_string = convert_result_to_json_string
 
         self._tools_with_names = self._validate_and_prepare_tools(tools)
+        self._is_warmed_up = False
 
     @staticmethod
     def _make_context_bound_invoke(tool_to_invoke: Tool, final_args: dict[str, Any]) -> Callable[[], Any]:
@@ -485,6 +487,17 @@ def _prepare_tool_call_params(
 
         return tool_calls, tool_call_params, error_messages
 
+    def warm_up(self):
+        """
+        Warm up the tool invoker.
+
+        This will warm up the tools registered in the tool invoker.
+        This method is idempotent and will only warm up the tools once.
+        """
+        if not self._is_warmed_up:
+            warm_up_tools(self.tools)
+            self._is_warmed_up = True
+
     @component.output_types(tool_messages=list[ChatMessage], state=State)
     def run(
         self,

haystack/tools/__init__.py

@@ -13,7 +13,7 @@
 from haystack.tools.component_tool import ComponentTool
 from haystack.tools.pipeline_tool import PipelineTool
 from haystack.tools.serde_utils import deserialize_tools_or_toolset_inplace, serialize_tools_or_toolset
-from haystack.tools.utils import flatten_tools_or_toolsets
+from haystack.tools.utils import flatten_tools_or_toolsets, warm_up_tools
 
 # Type alias for tools parameter - allows mixing Tools and Toolsets in a list
 ToolsType = Union[list[Union[Tool, Toolset]], Toolset]
@@ -27,7 +27,8 @@
     "PipelineTool",
     "serialize_tools_or_toolset",
     "Tool",
+    "ToolsType",
     "Toolset",
     "tool",
-    "ToolsType",
+    "warm_up_tools",
 ]

haystack/tools/component_tool.py

@@ -211,6 +211,16 @@ def component_invoker(**kwargs):
             outputs_to_string=outputs_to_string,
         )
         self._component = component
+        self._is_warmed_up = False
+
+    def warm_up(self):
+        """
+        Prepare the ComponentTool for use.
+        """
+        if not self._is_warmed_up:
+            if hasattr(self._component, "warm_up"):
+                self._component.warm_up()
+            self._is_warmed_up = True
 
     def to_dict(self) -> dict[str, Any]:
         """

haystack/tools/tool.py

@@ -21,6 +21,11 @@ class Tool:
     Accurate definitions of the textual attributes such as `name` and `description`
     are important for the Language Model to correctly prepare the call.
 
+    For resource-intensive operations like establishing connections to remote services or
+    loading models, override the `warm_up()` method. This method is called before the Tool
+    is used and should be idempotent, as it may be called multiple times during
+    pipeline/agent setup.
+
     :param name:
         Name of the Tool.
     :param description:
@@ -98,6 +103,16 @@ def tool_spec(self) -> dict[str, Any]:
         """
         return {"name": self.name, "description": self.description, "parameters": self.parameters}
 
+    def warm_up(self) -> None:
+        """
+        Prepare the Tool for use.
+
+        Override this method to establish connections to remote services, load models,
+        or perform other resource-intensive initialization. This method should be idempotent,
+        as it may be called multiple times.
+        """
+        pass
+
     def invoke(self, **kwargs: Any) -> Any:
         """
         Invoke the Tool with the provided keyword arguments.

haystack/tools/toolset.py

@@ -185,6 +185,15 @@ def __contains__(self, item: Any) -> bool:
             return item in self.tools
         return False
 
+    def warm_up(self) -> None:
+        """
+        Prepare the Toolset for use.
+
+        Override this method to set up shared resources like database connections or HTTP sessions.
+        This method should be idempotent, as it may be called multiple times.
+        """
+        pass
+
     def add(self, tool: Union[Tool, "Toolset"]) -> None:
         """
         Add a new Tool or merge another Toolset.
@@ -260,18 +269,12 @@ def __add__(self, other: Union[Tool, "Toolset", list[Tool]]) -> "Toolset":
         :raises ValueError: If the combination would result in duplicate tool names
         """
         if isinstance(other, Tool):
-            combined_tools = self.tools + [other]
-        elif isinstance(other, Toolset):
-            combined_tools = self.tools + list(other)
-        elif isinstance(other, list) and all(isinstance(item, Tool) for item in other):
-            combined_tools = self.tools + other
-        else:
-            raise TypeError(f"Cannot add {type(other).__name__} to Toolset")
-
-        # Check for duplicates
-        _check_duplicate_tool_names(combined_tools)
-
-        return Toolset(tools=combined_tools)
+            return Toolset(tools=self.tools + [other])
+        if isinstance(other, Toolset):
+            return _ToolsetWrapper([self, other])
+        if isinstance(other, list) and all(isinstance(item, Tool) for item in other):
+            return Toolset(tools=self.tools + other)
+        raise TypeError(f"Cannot add {type(other).__name__} to Toolset")
 
     def __len__(self) -> int:
         """
@@ -289,3 +292,52 @@ def __getitem__(self, index):
         :returns: The Tool at the specified index
         """
         return self.tools[index]
+
+
+class _ToolsetWrapper(Toolset):
+    """
+    A wrapper that holds multiple toolsets and provides a unified interface.
+
+    This is used internally when combining different types of toolsets to preserve
+    their individual configurations while still being usable with ToolInvoker.
+    """
+
+    def __init__(self, toolsets: list[Toolset]):
+        super().__init__([tool for toolset in toolsets for tool in toolset])
+        self.toolsets = toolsets
+
+    def __iter__(self):
+        """Iterate over all tools from all toolsets."""
+        for toolset in self.toolsets:
+            yield from toolset
+
+    def __contains__(self, item):
+        """Check if a tool is in any of the toolsets."""
+        return any(item in toolset for toolset in self.toolsets)
+
+    def warm_up(self):
+        """Warm up all toolsets."""
+        for toolset in self.toolsets:
+            toolset.warm_up()
+
+    def __len__(self):
+        """Return total number of tools across all toolsets."""
+        return sum(len(toolset) for toolset in self.toolsets)
+
+    def __getitem__(self, index):
+        """Get a tool by index across all toolsets."""
+        # Leverage iteration instead of manual index tracking
+        for i, tool in enumerate(self):
+            if i == index:
+                return tool
+        raise IndexError("ToolsetWrapper index out of range")
+
+    def __add__(self, other):
+        """Add another toolset or tool to this wrapper."""
+        if isinstance(other, Toolset):
+            return _ToolsetWrapper(self.toolsets + [other])
+        if isinstance(other, Tool):
+            return _ToolsetWrapper(self.toolsets + [Toolset([other])])
+        if isinstance(other, list) and all(isinstance(item, Tool) for item in other):
+            return _ToolsetWrapper(self.toolsets + [Toolset(other)])
+        raise TypeError(f"Cannot add {type(other).__name__} to ToolsetWrapper")

haystack/tools/utils.py

@@ -11,6 +11,28 @@
     from haystack.tools import ToolsType
 
 
+def warm_up_tools(tools: "Optional[ToolsType]" = None) -> None:
+    """
+    Warm up tools from various formats (Tools, Toolsets, or mixed lists).
+
+    :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 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"):
+                item.warm_up()
+
+
 def flatten_tools_or_toolsets(tools: "Optional[ToolsType]") -> list[Tool]:
     """
     Flatten tools from various formats into a list of Tool instances.

releasenotes/notes/tools-warm-up-support-e16cc043fed3653f.yaml

@@ -0,0 +1,8 @@
+---
+features:
+  - |
+    Added a `warm_up()` function to the `Tool` dataclass, allowing tools to perform resource-intensive initialization
+    before execution. Tools and Toolsets can now override the `warm_up()` method to establish connections to
+    remote services, load models, or perform other preparatory operations. The `ToolInvoker` and `Agent`
+    automatically call `warm_up()` on their tools during their own warm-up phase, ensuring tools are ready
+    before use.