Skip to content

Commit e7fb4f5

Browse files
HamidOnaHamidOna13
andauthored
fix: warm up individual tools inside Toolsets in warm_up_tools() (#10002)
* fix: warm up individual tools inside Toolsets in warm_up_tools() Related Issues: * Follows up on PR #9942 (feat: Add warm_up() method to ChatGenerators) * Addresses bug discovered during implementation of PR #9942 for issue #9907 Proposed Changes: The warm_up_tools() utility function was only calling warm_up() on Toolset objects themselves, but not on the individual Tool instances contained within them. This meant tools inside a Toolset were not properly initialized before use. This PR modifies warm_up_tools() to iterate through Toolsets and call warm_up() on each individual tool, in addition to calling warm_up() on the Toolset itself. Changes: - Modified warm_up_tools() in haystack/tools/utils.py to iterate through Toolsets when encountered (both as single argument and within lists) - Added iteration to call warm_up() on each individual Tool inside Toolsets - Added comprehensive test class TestWarmUpTools with 7 test cases How did you test it: - Added 7 comprehensive unit tests in test/tools/test_tools_utils.py: * test_warm_up_tools_with_none - handles None input * test_warm_up_tools_with_single_tool - single tool in list * test_warm_up_tools_with_single_toolset - KEY TEST: verifies both Toolset and individual tools are warmed * test_warm_up_tools_with_list_containing_toolset - toolset within list * test_warm_up_tools_with_multiple_toolsets - multiple toolsets * test_warm_up_tools_with_mixed_tools_and_toolsets - mixed scenarios * test_warm_up_tools_idempotency - safe to call multiple times Notes for the reviewer: I discovered this bug while implementing PR #9942 (for issue #9907). When a Toolset object is passed to a component's tools parameter, the warm_up_tools() function only calls Toolset.warm_up(), which is a no-op. It doesn't iterate through the individual tools inside the Toolset to warm them up. acknowledged by @vblagoje and @sjrl This implementation: - Modified warm_up_tools() to iterate through Toolsets and call warm_up() on each individual tool - Added comprehensive tests for Toolset warming behavior - Verified both the Toolset and its contained tools are warmed up Checklist: I have read the contributors guidelines and the code of conduct I have updated the related issue with new insights and changes I added unit tests and updated the docstrings I've used one of the conventional commit types for my PR title: fix: I documented my code I ran pre-commit hooks and fixed any issue * added release note * refactor: move tool warm-up iteration to Toolset.warm_up() Addresses architectural feedback - moved iteration logic from warm_up_tools() to base Toolset.warm_up() for better encapsulation. Subclasses can now override warm_up() to customize initialization without breaking the contract. - Toolset.warm_up() now iterates and warms tools by default - warm_up_tools() simplified to delegate to warm_up() - Updated tests and release notes --------- Co-authored-by: HamidOna13 <abdulhamid.onawole@aizatron.com>
1 parent 1e04b57 commit e7fb4f5

File tree

5 files changed

+261
-6
lines changed

5 files changed

+261
-6
lines changed

haystack/tools/toolset.py

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -189,10 +189,30 @@ def warm_up(self) -> None:
189189
"""
190190
Prepare the Toolset for use.
191191
192-
Override this method to set up shared resources like database connections or HTTP sessions.
192+
By default, this method iterates through and warms up all tools in the Toolset.
193+
Subclasses can override this method to customize initialization behavior, such as:
194+
195+
- Setting up shared resources (database connections, HTTP sessions) instead of
196+
warming individual tools
197+
- Implementing custom initialization logic for dynamically loaded tools
198+
- Controlling when and how tools are initialized
199+
200+
For example, a Toolset that manages tools from an external service (like MCPToolset)
201+
might override this to initialize a shared connection rather than warming up
202+
individual tools:
203+
204+
```python
205+
class MCPToolset(Toolset):
206+
def warm_up(self) -> None:
207+
# Only warm up the shared MCP connection, not individual tools
208+
self.mcp_connection = establish_connection(self.server_url)
209+
```
210+
193211
This method should be idempotent, as it may be called multiple times.
194212
"""
195-
pass
213+
for tool in self.tools:
214+
if hasattr(tool, "warm_up"):
215+
tool.warm_up()
196216

197217
def add(self, tool: Union[Tool, "Toolset"]) -> None:
198218
"""

haystack/tools/utils.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,21 +15,25 @@ def warm_up_tools(tools: "Optional[ToolsType]" = None) -> None:
1515
"""
1616
Warm up tools from various formats (Tools, Toolsets, or mixed lists).
1717
18+
For Toolset objects, this delegates to Toolset.warm_up(), which by default
19+
warms up all tools in the Toolset. Toolset subclasses can override warm_up()
20+
to customize initialization behavior (e.g., setting up shared resources).
21+
1822
:param tools: A list of Tool and/or Toolset objects, a single Toolset, or None.
1923
"""
2024
if tools is None:
2125
return
2226

23-
# If tools is a single Toolset, warm up the toolset itself
24-
if isinstance(tools, Toolset):
27+
# If tools is a single Toolset or Tool, warm it up
28+
if isinstance(tools, (Toolset, Tool)):
2529
if hasattr(tools, "warm_up"):
2630
tools.warm_up()
2731
return
2832

2933
# If tools is a list, warm up each item (Tool or Toolset)
3034
if isinstance(tools, list):
3135
for item in tools:
32-
if isinstance(item, (Toolset, Tool)) and hasattr(item, "warm_up"):
36+
if hasattr(item, "warm_up"):
3337
item.warm_up()
3438

3539

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
enhancements:
2+
- |
3+
Improved Toolset warm-up architecture for better encapsulation. The base
4+
Toolset.warm_up() method now warms up all tools by default, while subclasses
5+
can override it to customize initialization (e.g., setting up shared resources
6+
instead of warming individual tools). The warm_up_tools() utility function has
7+
been simplified to delegate to Toolset.warm_up().

test/components/tools/test_tool_invoker.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,8 @@ def __init__(self, tools):
136136

137137
def warm_up(self):
138138
self.was_warmed_up = True
139+
# Call parent to warm up individual tools
140+
super().warm_up()
139141

140142

141143
class TestToolInvokerCore:

test/tools/test_tools_utils.py

Lines changed: 223 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
import pytest
66

7-
from haystack.tools import Tool, Toolset, flatten_tools_or_toolsets
7+
from haystack.tools import Tool, Toolset, flatten_tools_or_toolsets, warm_up_tools
88

99

1010
def add_numbers(a: int, b: int) -> int:
@@ -171,3 +171,225 @@ def test_flatten_multiple_toolsets(self, add_tool, multiply_tool, subtract_tool)
171171
assert result[0].name == "add"
172172
assert result[1].name == "multiply"
173173
assert result[2].name == "subtract"
174+
175+
176+
class WarmupTrackingTool(Tool):
177+
"""A tool that tracks whether warm_up was called."""
178+
179+
def __init__(self, *args, **kwargs):
180+
super().__init__(*args, **kwargs)
181+
self.was_warmed_up = False
182+
183+
def warm_up(self):
184+
self.was_warmed_up = True
185+
186+
187+
class WarmupTrackingToolset(Toolset):
188+
"""A toolset that tracks whether warm_up was called."""
189+
190+
def __init__(self, tools):
191+
super().__init__(tools)
192+
self.was_warmed_up = False
193+
194+
def warm_up(self):
195+
self.was_warmed_up = True
196+
# Call parent to warm up individual tools
197+
super().warm_up()
198+
199+
200+
class TestWarmUpTools:
201+
"""Tests for the warm_up_tools() function"""
202+
203+
def test_warm_up_tools_with_none(self):
204+
"""Test that warm_up_tools with None does nothing."""
205+
# Should not raise any errors
206+
warm_up_tools(None)
207+
208+
def test_warm_up_tools_with_single_tool(self):
209+
"""Test that warm_up_tools works with a single tool in a list."""
210+
tool = WarmupTrackingTool(
211+
name="test_tool",
212+
description="A test tool",
213+
parameters={"type": "object", "properties": {}},
214+
function=lambda: "test",
215+
)
216+
217+
assert not tool.was_warmed_up
218+
warm_up_tools([tool])
219+
assert tool.was_warmed_up
220+
221+
def test_warm_up_tools_with_single_toolset(self):
222+
"""
223+
Test that when passing a single Toolset, both the Toolset.warm_up()
224+
and each individual tool's warm_up() are called.
225+
"""
226+
tool1 = WarmupTrackingTool(
227+
name="tool1",
228+
description="First tool",
229+
parameters={"type": "object", "properties": {}},
230+
function=lambda: "tool1",
231+
)
232+
tool2 = WarmupTrackingTool(
233+
name="tool2",
234+
description="Second tool",
235+
parameters={"type": "object", "properties": {}},
236+
function=lambda: "tool2",
237+
)
238+
239+
toolset = WarmupTrackingToolset([tool1, tool2])
240+
241+
assert not toolset.was_warmed_up
242+
assert not tool1.was_warmed_up
243+
assert not tool2.was_warmed_up
244+
245+
warm_up_tools(toolset)
246+
247+
# Both the toolset itself and individual tools should be warmed up
248+
assert toolset.was_warmed_up
249+
assert tool1.was_warmed_up
250+
assert tool2.was_warmed_up
251+
252+
def test_warm_up_tools_with_list_containing_toolset(self):
253+
"""Test that when a Toolset is in a list, individual tools inside get warmed up."""
254+
tool1 = WarmupTrackingTool(
255+
name="tool1",
256+
description="First tool",
257+
parameters={"type": "object", "properties": {}},
258+
function=lambda: "tool1",
259+
)
260+
tool2 = WarmupTrackingTool(
261+
name="tool2",
262+
description="Second tool",
263+
parameters={"type": "object", "properties": {}},
264+
function=lambda: "tool2",
265+
)
266+
267+
toolset = WarmupTrackingToolset([tool1, tool2])
268+
269+
assert not toolset.was_warmed_up
270+
assert not tool1.was_warmed_up
271+
assert not tool2.was_warmed_up
272+
273+
warm_up_tools([toolset])
274+
275+
# Both the toolset itself and individual tools should be warmed up
276+
assert toolset.was_warmed_up
277+
assert tool1.was_warmed_up
278+
assert tool2.was_warmed_up
279+
280+
def test_warm_up_tools_with_multiple_toolsets(self):
281+
"""Test multiple Toolsets in a list."""
282+
tool1 = WarmupTrackingTool(
283+
name="tool1",
284+
description="First tool",
285+
parameters={"type": "object", "properties": {}},
286+
function=lambda: "tool1",
287+
)
288+
tool2 = WarmupTrackingTool(
289+
name="tool2",
290+
description="Second tool",
291+
parameters={"type": "object", "properties": {}},
292+
function=lambda: "tool2",
293+
)
294+
tool3 = WarmupTrackingTool(
295+
name="tool3",
296+
description="Third tool",
297+
parameters={"type": "object", "properties": {}},
298+
function=lambda: "tool3",
299+
)
300+
301+
toolset1 = WarmupTrackingToolset([tool1])
302+
toolset2 = WarmupTrackingToolset([tool2, tool3])
303+
304+
assert not toolset1.was_warmed_up
305+
assert not toolset2.was_warmed_up
306+
assert not tool1.was_warmed_up
307+
assert not tool2.was_warmed_up
308+
assert not tool3.was_warmed_up
309+
310+
warm_up_tools([toolset1, toolset2])
311+
312+
# Both toolsets and all individual tools should be warmed up
313+
assert toolset1.was_warmed_up
314+
assert toolset2.was_warmed_up
315+
assert tool1.was_warmed_up
316+
assert tool2.was_warmed_up
317+
assert tool3.was_warmed_up
318+
319+
def test_warm_up_tools_with_mixed_tools_and_toolsets(self):
320+
"""Test list with both Tool objects and Toolsets."""
321+
standalone_tool = WarmupTrackingTool(
322+
name="standalone",
323+
description="Standalone tool",
324+
parameters={"type": "object", "properties": {}},
325+
function=lambda: "standalone",
326+
)
327+
toolset_tool1 = WarmupTrackingTool(
328+
name="toolset_tool1",
329+
description="Tool in toolset",
330+
parameters={"type": "object", "properties": {}},
331+
function=lambda: "toolset_tool1",
332+
)
333+
toolset_tool2 = WarmupTrackingTool(
334+
name="toolset_tool2",
335+
description="Another tool in toolset",
336+
parameters={"type": "object", "properties": {}},
337+
function=lambda: "toolset_tool2",
338+
)
339+
340+
toolset = WarmupTrackingToolset([toolset_tool1, toolset_tool2])
341+
342+
assert not standalone_tool.was_warmed_up
343+
assert not toolset.was_warmed_up
344+
assert not toolset_tool1.was_warmed_up
345+
assert not toolset_tool2.was_warmed_up
346+
347+
warm_up_tools([standalone_tool, toolset])
348+
349+
# All tools and the toolset should be warmed up
350+
assert standalone_tool.was_warmed_up
351+
assert toolset.was_warmed_up
352+
assert toolset_tool1.was_warmed_up
353+
assert toolset_tool2.was_warmed_up
354+
355+
def test_warm_up_tools_idempotency(self):
356+
"""Test that calling warm_up_tools() multiple times is safe."""
357+
358+
class WarmupCountingTool(Tool):
359+
"""A tool that counts how many times warm_up was called."""
360+
361+
def __init__(self, *args, **kwargs):
362+
super().__init__(*args, **kwargs)
363+
self.warm_up_count = 0
364+
365+
def warm_up(self):
366+
self.warm_up_count += 1
367+
368+
class WarmupCountingToolset(Toolset):
369+
"""A toolset that counts how many times warm_up was called."""
370+
371+
def __init__(self, tools):
372+
super().__init__(tools)
373+
self.warm_up_count = 0
374+
375+
def warm_up(self):
376+
self.warm_up_count += 1
377+
super().warm_up() # Also warm up individual tools
378+
379+
tool = WarmupCountingTool(
380+
name="counting_tool",
381+
description="A counting tool",
382+
parameters={"type": "object", "properties": {}},
383+
function=lambda: "test",
384+
)
385+
toolset = WarmupCountingToolset([tool])
386+
387+
# Call warm_up_tools multiple times
388+
warm_up_tools(toolset)
389+
warm_up_tools(toolset)
390+
warm_up_tools(toolset)
391+
392+
# warm_up_tools itself doesn't prevent multiple calls,
393+
# but verify the calls actually happen multiple times
394+
assert toolset.warm_up_count == 3
395+
assert tool.warm_up_count == 3

0 commit comments

Comments
 (0)