mirror of
https://github.com/deepset-ai/haystack.git
synced 2026-01-01 17:47:19 +00:00
e7fb4f5ac5
* 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>
396 lines
13 KiB
Python
396 lines
13 KiB
Python
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import pytest
|
|
|
|
from haystack.tools import Tool, Toolset, flatten_tools_or_toolsets, warm_up_tools
|
|
|
|
|
|
def add_numbers(a: int, b: int) -> int:
|
|
"""Add two numbers."""
|
|
return a + b
|
|
|
|
|
|
def multiply_numbers(a: int, b: int) -> int:
|
|
"""Multiply two numbers."""
|
|
return a * b
|
|
|
|
|
|
def subtract_numbers(a: int, b: int) -> int:
|
|
"""Subtract b from a."""
|
|
return a - b
|
|
|
|
|
|
@pytest.fixture
|
|
def add_tool():
|
|
return Tool(
|
|
name="add",
|
|
description="Add two numbers",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=add_numbers,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def multiply_tool():
|
|
return Tool(
|
|
name="multiply",
|
|
description="Multiply two numbers",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=multiply_numbers,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def subtract_tool():
|
|
return Tool(
|
|
name="subtract",
|
|
description="Subtract two numbers",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=subtract_numbers,
|
|
)
|
|
|
|
|
|
class TestFlattenToolsOrToolsets:
|
|
def test_flatten_none(self):
|
|
"""Test that None returns an empty list."""
|
|
result = flatten_tools_or_toolsets(None)
|
|
assert result == []
|
|
|
|
def test_flatten_empty_list(self):
|
|
"""Test that an empty list returns an empty list."""
|
|
result = flatten_tools_or_toolsets([])
|
|
assert result == []
|
|
|
|
def test_flatten_list_of_tools(self, add_tool, multiply_tool):
|
|
"""Test that a list of Tool instances is returned as-is."""
|
|
tools = [add_tool, multiply_tool]
|
|
result = flatten_tools_or_toolsets(tools)
|
|
assert result == tools
|
|
assert len(result) == 2
|
|
assert result[0].name == "add"
|
|
assert result[1].name == "multiply"
|
|
|
|
def test_flatten_single_toolset(self, add_tool, multiply_tool):
|
|
"""Test that a single Toolset is converted to a list of Tools."""
|
|
toolset = Toolset([add_tool, multiply_tool])
|
|
result = flatten_tools_or_toolsets(toolset)
|
|
assert isinstance(result, list)
|
|
assert len(result) == 2
|
|
assert all(isinstance(t, Tool) for t in result)
|
|
assert result[0].name == "add"
|
|
assert result[1].name == "multiply"
|
|
|
|
def test_flatten_list_of_toolsets(self, add_tool, multiply_tool, subtract_tool):
|
|
"""Test that a list of Toolset instances is flattened to a single list of Tools."""
|
|
toolset1 = Toolset([add_tool])
|
|
toolset2 = Toolset([multiply_tool, subtract_tool])
|
|
|
|
result = flatten_tools_or_toolsets([toolset1, toolset2])
|
|
assert isinstance(result, list)
|
|
assert len(result) == 3
|
|
assert all(isinstance(t, Tool) for t in result)
|
|
assert result[0].name == "add"
|
|
assert result[1].name == "multiply"
|
|
assert result[2].name == "subtract"
|
|
|
|
def test_flatten_list_with_mixed_tools_and_toolsets(self, add_tool, multiply_tool, subtract_tool):
|
|
"""Test that a mixed list of Tool and Toolset instances is flattened correctly."""
|
|
toolset = Toolset([multiply_tool])
|
|
mixed_list = [add_tool, toolset, subtract_tool]
|
|
|
|
result = flatten_tools_or_toolsets(mixed_list)
|
|
assert isinstance(result, list)
|
|
assert len(result) == 3
|
|
assert all(isinstance(t, Tool) for t in result)
|
|
assert result[0].name == "add"
|
|
assert result[1].name == "multiply"
|
|
assert result[2].name == "subtract"
|
|
|
|
def test_flatten_empty_toolset(self):
|
|
"""Test that an empty Toolset returns an empty list."""
|
|
toolset = Toolset([])
|
|
result = flatten_tools_or_toolsets(toolset)
|
|
assert result == []
|
|
|
|
def test_flatten_list_with_empty_toolsets(self, add_tool):
|
|
"""Test that a list with empty Toolsets handles correctly."""
|
|
toolset1 = Toolset([])
|
|
toolset2 = Toolset([add_tool])
|
|
toolset3 = Toolset([])
|
|
|
|
result = flatten_tools_or_toolsets([toolset1, toolset2, toolset3])
|
|
assert isinstance(result, list)
|
|
assert len(result) == 1
|
|
assert result[0].name == "add"
|
|
|
|
def test_flatten_invalid_type_in_list(self):
|
|
"""Test that invalid types in the list raise TypeError."""
|
|
with pytest.raises(TypeError, match="Items in the tools list must be Tool or Toolset instances"):
|
|
flatten_tools_or_toolsets(["not_a_tool"])
|
|
|
|
with pytest.raises(TypeError, match="Items in the tools list must be Tool or Toolset instances"):
|
|
flatten_tools_or_toolsets([123])
|
|
|
|
with pytest.raises(TypeError, match="Items in the tools list must be Tool or Toolset instances"):
|
|
flatten_tools_or_toolsets([{"key": "value"}])
|
|
|
|
def test_flatten_invalid_type(self):
|
|
"""Test that invalid root types raise TypeError."""
|
|
with pytest.raises(TypeError, match="tools must be list\\[Union\\[Tool, Toolset\\]\\], Toolset, or None"):
|
|
flatten_tools_or_toolsets("not_valid")
|
|
|
|
with pytest.raises(TypeError, match="tools must be list\\[Union\\[Tool, Toolset\\]\\], Toolset, or None"):
|
|
flatten_tools_or_toolsets(123)
|
|
|
|
with pytest.raises(TypeError, match="tools must be list\\[Union\\[Tool, Toolset\\]\\], Toolset, or None"):
|
|
flatten_tools_or_toolsets({"key": "value"})
|
|
|
|
def test_flatten_multiple_toolsets(self, add_tool, multiply_tool, subtract_tool):
|
|
"""Test flattening a list of multiple Toolsets."""
|
|
toolset1 = Toolset([add_tool])
|
|
toolset2 = Toolset([multiply_tool])
|
|
toolset3 = Toolset([subtract_tool])
|
|
|
|
# List of three separate toolsets
|
|
result = flatten_tools_or_toolsets([toolset1, toolset2, toolset3])
|
|
assert len(result) == 3
|
|
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
|