Skip to content

Conversation

@vblagoje
Copy link
Member

@vblagoje vblagoje commented Oct 7, 2025

Why:

Enables tools and toolsets to perform initialization during Agent/ToolInvoker warmup. This supports tools that need to establish remote connections, load models, initialize connection pools, or perform other setup operations before use.

What:

  • Added warm_up() method to Tool and Toolset as extension points for initialization
  • Created warm_up_tools() utility function to warm up lists of tools or Toolsets
  • Modified ToolInvoker to warm up tools during initialization
  • Optimized duplicate checking in Toolset.__add__() to avoid redundant validation
  • Added focused tests covering warmup delegation and integration

How can it be used:

from haystack.tools import Toolset
from haystack.components.agents import Agent

# Custom toolset with initialization needs
class DatabaseToolset(Toolset):
    def __init__(self, connection_string):
        self.connection_string = connection_string
        self.pool = None
        super().__init__([query_tool, update_tool])
        
    def warm_up(self):
        # Initialize connection pool
        self.pool = create_connection_pool(self.connection_string)

How did you test it:

  • Added comprehensive tests in test/tools/test_toolset_warmup_and_wrapper.py
  • Verified warmup delegation through _ToolsetWrapper to each wrapped toolset
  • Manual tests with this PR and mcp-haystack with warm_up implemented in itinerary agent

Notes for the reviewer:

  • _ToolsetWrapper.warm_up() delegates to each wrapped toolset - critical for preserving individual initialization behavior when combining toolsets
  • No breaking changes - warm_up() defaults to no-op in base classes
  • Works with previous versions of Tool implementation "out there". We check if there is warm_up attr and call it only if found.

@vblagoje vblagoje added the ignore-for-release-notes PRs with this flag won't be included in the release notes. label Oct 7, 2025
@github-actions github-actions bot added the type:documentation Improvements on the docs label Oct 7, 2025
@coveralls
Copy link
Collaborator

coveralls commented Oct 7, 2025

Pull Request Test Coverage Report for Build 18647136217

Warning: This coverage report may be inaccurate.

This pull request's base commit is no longer the HEAD commit of its target branch. This means it includes changes from outside the original pull request, including, potentially, unrelated coverage changes.

Details

  • 0 of 0 changed or added relevant lines in 0 files are covered.
  • 36 unchanged lines in 7 files lost coverage.
  • Overall coverage decreased (-0.04%) to 92.179%

Files with Coverage Reduction New Missed Lines %
tools/utils.py 1 96.55%
tools/component_tool.py 2 94.25%
core/pipeline/async_pipeline.py 3 65.88%
tools/tool.py 3 96.55%
components/agents/agent.py 7 96.83%
components/tools/tool_invoker.py 9 95.98%
tools/toolset.py 11 84.34%
Totals Coverage Status
Change from base Build 18645317230: -0.04%
Covered Lines: 13425
Relevant Lines: 14564

💛 - Coveralls

@vblagoje vblagoje marked this pull request as ready for review October 8, 2025 10:03
@vblagoje vblagoje requested a review from a team as a code owner October 8, 2025 10:03
@vblagoje vblagoje requested review from anakin87, julian-risch and sjrl and removed request for a team October 8, 2025 10:03
@vblagoje
Copy link
Member Author

vblagoje commented Oct 8, 2025

Note to @anakin87 @julian-risch - @sjrl and I talked about invoking tool warm up in all chat generators via their warm_up method. Implementation in this PR will make Agent work because now ToolInvoker warms up tools. But to have pipelines working with tools warmup (ChatGenerator + ToolInvoker pipeline without Agent) we need to add warm_up to all ChatGenerators. However, that is scope for another PR.

@vblagoje vblagoje changed the title Tools warmup Add Tools warm_up Oct 8, 2025
@vblagoje
Copy link
Member Author

@sjrl @julian-risch @anakin87 can I 🙏 have some feedback on this PR so we can push it forward!

Copy link
Member

@julian-risch julian-risch left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Left some comments. My my concern is the renaming of serde_utils.py to utils.py, which is a breaking change. Other than that looks quite good to me already.

@sjrl
Copy link
Contributor

sjrl commented Oct 13, 2025

@vblagoje I believe we should also update the Agent's warm up method to warm up it's tools as well right? Possibly I think would also make sense to update the ChatGenerators in main to do so as well. What do you think?

Comment on lines +297 to +303
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.
"""
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@vblagoje is it possible that this wrapper could also solve this issue deepset-ai/haystack-core-integrations#2369 if we add to_dict and from_dict methods?

@vblagoje
Copy link
Member Author

vblagoje commented Oct 14, 2025

I've tested this version of haystack with mcp and warmup and it works ok. However, as we discussed above, Agent warm_up needs to be called to warm up tools. Is that ok for no code envs? cc @julian-risch @sjrl

@sjrl
Copy link
Contributor

sjrl commented Oct 14, 2025

@vblagoje how do you mean by this comment?

Agent warm_up needs to be called to warm up tools. Is that ok for no code envs?

I think this is fine. Not entirely sure how this applies for no code envs

@sjrl
Copy link
Contributor

sjrl commented Oct 14, 2025

Looking good @vblagoje I have a one more comment in addition to the few above

Could we create a warm_up method for ComponentTool to call the underlying component's warm_up method?

@vblagoje vblagoje removed the ignore-for-release-notes PRs with this flag won't be included in the release notes. label Oct 14, 2025
vblagoje and others added 2 commits October 15, 2025 16:09
Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com>
Copy link
Contributor

@sjrl sjrl left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@vblagoje looks good to me!

@julian-risch julian-risch dismissed their stale review October 17, 2025 13:40

Sebastian re-reviewed and approved. My earlier comments have been addressed.

@vblagoje
Copy link
Member Author

@sjrl Please have another look for changes since merge. I left _ToolsetWrapper because it doesn't affect main APIs but we can re-eval if we want to support + operations between Tool and Toolset

@sjrl
Copy link
Contributor

sjrl commented Oct 20, 2025

@vblagoje looks good!

@vblagoje vblagoje merged commit 9bc59c3 into main Oct 20, 2025
35 of 39 checks passed
@vblagoje vblagoje deleted the tools_warmup_new branch October 20, 2025 10:11
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

topic:tests type:documentation Improvements on the docs

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants