deepset-ai/haystack
1
0
Fork 0
mirror of https://github.com/deepset-ai/haystack.git synced 2026-10-12 02:12:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat!: add lifecycle management for Tools/Toolsets (#13024)

Browse Source
This commit is contained in:
Stefano Fiorucci
2026-10-05 12:13:43 +02:00
committed by GitHub
parent 410e4f677f
commit 9b8c3ba911
27 changed files with 591 additions and 725 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs-website/docs/tools/mcptoolset.mdx
+2
View File
@@ -24,6 +24,8 @@ description: "`MCPToolset` connects to an MCP-compliant server and automatically
MCPToolset is a subclass of `Toolset` that dynamically discovers and loads tools from any MCP-compliant server.
The Agent warms up the toolset automatically before a run. When using it directly with a Chat Generator, first call `warm_up_tools(toolset)` from `haystack.tools` to load its tools.
It supports:
- **Streamable HTTP** for connecting to HTTP servers
docs-website/docs/tools/toolset.mdx
+6
View File
@@ -130,12 +130,18 @@ Output:
### With a ChatGenerator
Lazy toolsets, such as [`MCPToolset`](mcptoolset.mdx), load their tools during warm-up. When using them directly with a Chat Generator, prepare them first as shown in the comments below.
You can pass a `Toolset` directly to a Chat Generator. The model prepares the tool calls; executing them (for example, with `Tool.invoke`) is up to you:
```python
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.dataclasses import ChatMessage
# For a lazy toolset such as MCPToolset, prepare it before passing it to the generator:
# from haystack.tools import warm_up_tools
# warm_up_tools(toolset)
chat_generator = OpenAIChatGenerator(model="gpt-5.4-nano", tools=math_toolset)
user_message = ChatMessage.from_user("What is 10 minus 5?")
haystack/components/agents/agent.py
+25 -6
View File
@@ -54,10 +54,13 @@ from haystack.tools import (
Toolset,
ToolsType,
_check_duplicate_tool_names,
close_tools,
close_tools_async,
deserialize_tools_or_toolset_inplace,
flatten_tools_or_toolsets,
serialize_tools_or_toolset,
warm_up_tools,
warm_up_tools_async,
)
from haystack.utils.async_utils import _execute_component_async
from haystack.utils.callable_serialization import deserialize_callable, serialize_callable
@@ -601,7 +604,7 @@ class Agent:
async def warm_up_async(self) -> None:
"""Warm up the tools, hooks, and the underlying chat generator on the serving event loop."""
warm_up_tools(tools=self.tools)
await warm_up_tools_async(tools=self.tools)
await warm_up_hooks_async(self.hooks)
if hasattr(self.chat_generator, "warm_up_async"):
await self.chat_generator.warm_up_async()
@@ -609,13 +612,15 @@ class Agent:
self.chat_generator.warm_up()
def close(self) -> None:
"""Release the hooks' and the underlying chat generator's resources."""
"""Release tools, hooks, and chat generator resources."""
close_tools(tools=self.tools)
close_hooks(self.hooks)
if hasattr(self.chat_generator, "close"):
self.chat_generator.close()
async def close_async(self) -> None:
"""Release the hooks' and the underlying chat generator's async resources."""
"""Release async tools, hooks, and chat generator resources."""
await close_tools_async(tools=self.tools)
await close_hooks_async(self.hooks)
if hasattr(self.chat_generator, "close_async"):
await self.chat_generator.close_async()
@@ -817,9 +822,6 @@ class Agent:
if isinstance(tools, (Toolset, list)):
selected = cast(ToolsType, tools) # mypy can't narrow the Union type from the isinstance checks
# Per-run tools are not covered by the Agent's own warm_up(), so warm them up here.
# warm_up() is expected to be idempotent, so re-warming on every run is cheap.
warm_up_tools(tools=selected)
return _spawn_tools(tools=selected)
raise TypeError(
@@ -847,6 +849,8 @@ class Agent:
precedence, keys set only at initialization are kept.
:param tools: Optional list of Tool objects, a Toolset, or list of tool names to use for this run.
When passing tool names, tools are selected from the Agent's originally configured tools.
Tool and Toolset objects passed here are warmed up automatically; the caller is responsible for
closing them if they hold resources.
:param hook_context: Optional dictionary of request-scoped resources made available to hooks via
`state.data.get("hook_context")`. Useful in web/server environments to provide per-request objects
(e.g., WebSocket connections, async queues, Redis pub/sub clients) that a hook can use, for
@@ -874,6 +878,12 @@ class Agent:
"""
agent_inputs = {"messages": messages, "streaming_callback": streaming_callback, **kwargs}
self.warm_up()
# warm up tools passed at runtime
if isinstance(tools, Toolset):
tools_to_warm_up: ToolsType = tools
else:
tools_to_warm_up = [tool for tool in tools or [] if not isinstance(tool, str)]
warm_up_tools(tools=tools_to_warm_up)
exe_context = self._initialize_fresh_execution(
messages=messages,
@@ -934,6 +944,9 @@ class Agent:
with the `generation_kwargs` passed at the chat generator's initialization: keys provided here take
precedence, keys set only at initialization are kept.
:param tools: Optional list of Tool objects, a Toolset, or list of tool names to use for this run.
When passing tool names, tools are selected from the Agent's originally configured tools.
Tool and Toolset objects passed here are warmed up automatically; the caller is responsible for
closing them if they hold resources.
:param hook_context: Optional dictionary of request-scoped resources made available to hooks via
`state.data.get("hook_context")`. Useful in web/server environments to provide per-request objects
(e.g., WebSocket connections, async queues, Redis pub/sub clients) that a hook can use, for
@@ -961,6 +974,12 @@ class Agent:
"""
agent_inputs = {"messages": messages, "streaming_callback": streaming_callback, **kwargs}
await self.warm_up_async()
# warm up tools passed at runtime
if isinstance(tools, Toolset):
tools_to_warm_up: ToolsType = tools
else:
tools_to_warm_up = [tool for tool in tools or [] if not isinstance(tool, str)]
await warm_up_tools_async(tools=tools_to_warm_up)
exe_context = self._initialize_fresh_execution(
messages=messages,
haystack/components/agents/utils.py
+1 -1
View File
@@ -115,7 +115,7 @@ def _select_tools_by_name(configured_tools: ToolsType, names: list[str]) -> list
Standalone Tools are kept when their name is requested. A Toolset with run-scoped state (one overriding
`spawn()`, such as SearchableToolset) is replaced by a per-run copy carrying the requested names, so its
dynamic behavior (search/lazy-loading) is preserved without mutating the shared, configured Toolset. Any
other Toolset is warmed up and reduced to the matching Tools.
other Toolset is reduced to the matching Tools. The caller must warm up the tools before selection.
:param configured_tools: The tools configured on the Agent.
:param names: The requested tool names.
haystack/components/generators/chat/azure.py
+2 -11
View File
@@ -18,7 +18,6 @@ from haystack.tools import (
deserialize_tools_or_toolset_inplace,
flatten_tools_or_toolsets,
serialize_tools_or_toolset,
warm_up_tools,
)
from haystack.utils import Secret, deserialize_callable, serialize_callable
from haystack.utils.http_client import init_http_client
@@ -234,7 +233,6 @@ class AzureOpenAIChatGenerator(OpenAIChatGenerator):
self.client: AzureOpenAI | None = None
self.async_client: AsyncAzureOpenAI | None = None
self._tools_warmed_up = False
def _client_kwargs(self) -> dict[str, Any]:
timeout = self.timeout if self.timeout is not None else float(os.environ.get("OPENAI_TIMEOUT", "30.0"))
@@ -260,16 +258,10 @@ class AzureOpenAIChatGenerator(OpenAIChatGenerator):
"azure_ad_token_provider": self.azure_ad_token_provider,
}
def _warm_up_tools(self) -> None:
if not self._tools_warmed_up:
warm_up_tools(self.tools)
self._tools_warmed_up = True
def warm_up(self) -> None:
"""
Warm up the tools and initialize the synchronous Azure OpenAI client.
Initialize the synchronous Azure OpenAI client.
"""
self._warm_up_tools()
if self.client is None:
# openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
# https://github.com/openai/openai-python/blob/main/httpx2.md
@@ -281,9 +273,8 @@ class AzureOpenAIChatGenerator(OpenAIChatGenerator):
async def warm_up_async(self) -> None: # noqa: RUF029
"""
Warm up the tools and initialize the asynchronous Azure OpenAI client on the serving event loop.
Initialize the asynchronous Azure OpenAI client on the serving event loop.
"""
self._warm_up_tools()
if self.async_client is None:
# openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
# https://github.com/openai/openai-python/blob/main/httpx2.md
haystack/components/generators/chat/openai.py
+2 -11
View File
@@ -47,7 +47,6 @@ from haystack.tools import (
deserialize_tools_or_toolset_inplace,
flatten_tools_or_toolsets,
serialize_tools_or_toolset,
warm_up_tools,
)
from haystack.utils import Secret, deserialize_callable, serialize_callable
from haystack.utils.http_client import init_http_client
@@ -212,7 +211,6 @@ class OpenAIChatGenerator:
self.client: OpenAI | None = None
self.async_client: AsyncOpenAI | None = None
self._tools_warmed_up = False
def _client_kwargs(self) -> dict[str, Any]:
timeout = self.timeout if self.timeout is not None else float(os.environ.get("OPENAI_TIMEOUT", "30.0"))
@@ -227,16 +225,10 @@ class OpenAIChatGenerator:
"max_retries": max_retries,
}
def _warm_up_tools(self) -> None:
if not self._tools_warmed_up:
warm_up_tools(self.tools)
self._tools_warmed_up = True
def warm_up(self) -> None:
"""
Warm up the tools and initialize the synchronous OpenAI client.
Initialize the synchronous OpenAI client.
"""
self._warm_up_tools()
if self.client is None:
# openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
# https://github.com/openai/openai-python/blob/main/httpx2.md
@@ -248,9 +240,8 @@ class OpenAIChatGenerator:
async def warm_up_async(self) -> None: # noqa: RUF029
"""
Warm up the tools and initialize the asynchronous OpenAI client on the serving event loop.
Initialize the asynchronous OpenAI client on the serving event loop.
"""
self._warm_up_tools()
if self.async_client is None:
# openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
# https://github.com/openai/openai-python/blob/main/httpx2.md
haystack/components/generators/chat/openai_responses.py
+2 -15
View File
@@ -37,7 +37,6 @@ from haystack.tools import (
deserialize_tools_or_toolset_inplace,
flatten_tools_or_toolsets,
serialize_tools_or_toolset,
warm_up_tools,
)
from haystack.utils import Secret, deserialize_callable, serialize_callable
from haystack.utils.http_client import init_http_client
@@ -232,7 +231,6 @@ class OpenAIResponsesChatGenerator:
self.client: OpenAI | None = None
self.async_client: AsyncOpenAI | None = None
self._tools_warmed_up = False
def _client_kwargs(self) -> dict[str, Any]:
timeout = self.timeout if self.timeout is not None else float(os.environ.get("OPENAI_TIMEOUT", "30.0"))
@@ -248,20 +246,10 @@ class OpenAIResponsesChatGenerator:
"max_retries": max_retries,
}
def _warm_up_tools(self) -> None:
if not self._tools_warmed_up:
is_openai_tool = isinstance(self.tools, list) and bool(self.tools) and isinstance(self.tools[0], dict)
# We only warm up Haystack tools, not OpenAI/MCP tools
# The type ignore is needed because mypy cannot infer the type correctly
if not is_openai_tool:
warm_up_tools(self.tools) # type: ignore[arg-type]
self._tools_warmed_up = True
def warm_up(self) -> None:
"""
Warm up the tools and initialize the synchronous OpenAI client.
Initialize the synchronous OpenAI client.
"""
self._warm_up_tools()
if self.client is None:
# openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
# https://github.com/openai/openai-python/blob/main/httpx2.md
@@ -273,9 +261,8 @@ class OpenAIResponsesChatGenerator:
async def warm_up_async(self) -> None: # noqa: RUF029
"""
Warm up the tools and initialize the asynchronous OpenAI client on the serving event loop.
Initialize the asynchronous OpenAI client on the serving event loop.
"""
self._warm_up_tools()
if self.async_client is None:
# openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
# https://github.com/openai/openai-python/blob/main/httpx2.md
haystack/tools/__init__.py
+4 -1
View File
@@ -31,7 +31,7 @@ _import_structure = {
"pipeline_tool": ["PipelineTool"],
"agent_tool": ["AgentTool"],
"serde_utils": ["deserialize_tools_or_toolset_inplace", "serialize_tools_or_toolset"],
"utils": ["flatten_tools_or_toolsets", "warm_up_tools"],
"utils": ["flatten_tools_or_toolsets", "warm_up_tools", "warm_up_tools_async", "close_tools", "close_tools_async"],
"tool_types": ["ToolsType"],
}
@@ -45,8 +45,11 @@ if TYPE_CHECKING:
from haystack.tools.skills import SkillToolset as SkillToolset
from haystack.tools.tool_types import ToolsType as ToolsType
from haystack.tools.toolset import Toolset as Toolset
from haystack.tools.utils import close_tools as close_tools
from haystack.tools.utils import close_tools_async as close_tools_async
from haystack.tools.utils import flatten_tools_or_toolsets as flatten_tools_or_toolsets
from haystack.tools.utils import warm_up_tools as warm_up_tools
from haystack.tools.utils import warm_up_tools_async as warm_up_tools_async
else:
sys.modules[__name__] = LazyImporter(
name=__name__,
haystack/tools/component_tool.py
+22 -8
View File
@@ -236,7 +236,6 @@ class ComponentTool(Tool):
# Store component before calling super().__init__() so _get_valid_outputs() can access it
self._component = component
self._is_warmed_up = False
# Create the Tool instance with the component invoker as the function to be called and the schema.
# When the wrapped component exposes a `run_async`, also pass the async invoker.
@@ -274,13 +273,28 @@ class ComponentTool(Tool):
return set(self._component.__haystack_output__._sockets_dict.keys()) # type: ignore[attr-defined]
def warm_up(self) -> None:
"""
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
"""Warm up the wrapped component."""
if hasattr(self._component, "warm_up"):
self._component.warm_up()
async def warm_up_async(self) -> None:
"""Warm up the wrapped component on the serving event loop."""
if self.async_function is not None and hasattr(self._component, "warm_up_async"):
await self._component.warm_up_async()
elif hasattr(self._component, "warm_up"):
self._component.warm_up()
def close(self) -> None:
"""Release the wrapped component's synchronous resources."""
if hasattr(self._component, "close"):
self._component.close()
async def close_async(self) -> None:
"""Release the wrapped component's async resources."""
if self.async_function is not None and hasattr(self._component, "close_async"):
await self._component.close_async()
elif hasattr(self._component, "close"):
self._component.close()
def to_dict(self) -> dict[str, Any]:
"""
haystack/tools/searchable_toolset.py
+47 -16
View File
@@ -14,7 +14,13 @@ from haystack.tools.from_function import create_tool_from_function
from haystack.tools.serde_utils import deserialize_tools_or_toolset_inplace, serialize_tools_or_toolset
from haystack.tools.tool import Tool, _check_duplicate_tool_names
from haystack.tools.toolset import Toolset
from haystack.tools.utils import flatten_tools_or_toolsets, warm_up_tools
from haystack.tools.utils import (
close_tools,
close_tools_async,
flatten_tools_or_toolsets,
warm_up_tools,
warm_up_tools_async,
)
if TYPE_CHECKING:
from haystack.tools import ToolsType
@@ -143,21 +149,32 @@ class SearchableToolset(Toolset):
def warm_up(self) -> None:
"""
Prepare the toolset for use.
Warms up the catalog (so lazy toolsets like MCPToolset can connect) and flattens it. Above the passthrough
threshold, it also indexes the catalog and creates the search_tools bootstrap tool.
This method is idempotent: it only warms up the toolset the first time it is called.
:raises ValueError: If the flattened catalog contains tools with duplicate names.
Warm up the catalog tools and initialize tool discovery.
"""
if self._passthrough is not None:
return
# Warm up the catalog first (triggers lazy connections like MCPToolset), then flatten — lazy toolsets will
# have their real tools available.
warm_up_tools(self._raw_catalog)
self._build_catalog()
async def warm_up_async(self) -> None:
"""
Warm up the catalog tools asynchronously and initialize tool discovery.
"""
await warm_up_tools_async(self._raw_catalog)
self._build_catalog()
def _build_catalog(self) -> None:
"""
Flatten the warmed-up catalog, validate tool names, and initialize discovery once.
The catalog must be warmed up first so lazy toolsets, such as MCPToolset, expose their tools.
Catalogs below the search threshold use passthrough mode. Otherwise, build the BM25 index
and create the search_tools bootstrap tool.
:raises ValueError: If the catalog contains tools with duplicate names.
"""
if self._passthrough is not None:
return
self._catalog = flatten_tools_or_toolsets(self._raw_catalog)
_check_duplicate_tool_names(self._catalog)
self._passthrough = len(self._catalog) < self._search_threshold
@@ -180,13 +197,25 @@ class SearchableToolset(Toolset):
Return the full catalog of tools that can be selected by name.
Iteration only exposes the search tool plus already-discovered tools, but name-based selection can target
any tool in the catalog, so this returns the entire flattened catalog (warming up first if needed).
any tool in the catalog, so this returns the entire flattened catalog. Prepare it with `warm_up()`
or `warm_up_async()` before selection.
:returns: The flattened catalog of tools.
"""
self.warm_up()
return list(self._catalog)
def close(self) -> None:
"""Close the tools and toolsets in the catalog."""
close_tools(self._raw_catalog)
self._passthrough = None
self._discovered_tools = {}
async def close_async(self) -> None:
"""Close the tools and toolsets in the catalog asynchronously."""
await close_tools_async(self._raw_catalog)
self._passthrough = None
self._discovered_tools = {}
def clear(self) -> None:
"""
Clear all discovered tools.
@@ -209,7 +238,8 @@ class SearchableToolset(Toolset):
restriction.
:returns: A run-scoped copy of this SearchableToolset.
"""
self.warm_up()
if self._passthrough is None:
self.warm_up()
new = copy.copy(self)
new._discovered_tools = {}
new._selected_tool_names = set(selected_tool_names) if selected_tool_names is not None else None
@@ -308,7 +338,8 @@ class SearchableToolset(Toolset):
# This toolset materializes everything (flattened catalog, bootstrap tool, passthrough decision) in warm_up.
# Without warming here, iterating before warm_up would yield nothing, so we warm up to make the toolset usable
# at all.
self.warm_up()
if self._passthrough is None:
self.warm_up()
if self._passthrough:
yield from (tool for tool in self._catalog if self._is_selected(tool.name))
else:
haystack/tools/skills/skill_toolset.py
+7
View File
@@ -103,6 +103,13 @@ class SkillToolset(Toolset):
self._load_skill_tool.description = self._load_skill_description()
self._is_warmed_up = True
def close(self) -> None:
"""Close the skill store and invalidate the cached catalog."""
if hasattr(self._store, "close"):
self._store.close()
self._skills = {}
self._is_warmed_up = False
def add(self, tool: Tool) -> None:
"""Adding tools is not supported: a SkillToolset's tools are fixed and defined by its store."""
raise NotImplementedError(
haystack/tools/tool.py
+18 -20
View File
@@ -25,9 +25,10 @@ class Tool:
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.
loading models, implement an optional `warm_up()` method. It is called before invocation
and during pipeline/agent setup, and must be idempotent. Native async tools can implement
`warm_up_async()` instead. Implement `close()` and/or `close_async()` to release owned resources;
these methods must be safe to call before warm-up and multiple times.
:param name:
Name of the Tool.
@@ -270,19 +271,9 @@ class Tool:
"""
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 synchronously with the provided keyword arguments.
Warm up the Tool if supported and invoke it synchronously with the provided keyword arguments.
:raises ToolInvocationError: If the Tool has no sync `function`, or if the underlying call
raises an exception.
@@ -295,6 +286,8 @@ class Tool:
)
try:
if hasattr(self, "warm_up"):
self.warm_up()
result = self.function(**kwargs)
except Exception as e:
raise ToolInvocationError(
@@ -306,16 +299,21 @@ class Tool:
"""
Invoke the Tool asynchronously with the provided keyword arguments.
If `async_function` is set, it is awaited directly. Otherwise the sync `function` is dispatched to a worker
thread via `asyncio.to_thread`, which propagates the current context to the worker.
If `async_function` is set, prepare with `warm_up_async` (falling back to `warm_up`) before awaiting it.
Otherwise, synchronous preparation and invocation run in a worker thread via `asyncio.to_thread`,
which propagates the current context to the worker.
:raises ToolInvocationError: If the underlying call raises an exception.
"""
if self.async_function is None:
return await asyncio.to_thread(self.invoke, **kwargs)
try:
if self.async_function is not None:
return await self.async_function(**kwargs)
# `function` is guaranteed to be set: __post_init__ enforces at least one of the two.
return await asyncio.to_thread(self.function, **kwargs) # type: ignore[arg-type]
if hasattr(self, "warm_up_async"):
await self.warm_up_async()
elif hasattr(self, "warm_up"):
self.warm_up()
return await self.async_function(**kwargs)
except Exception as e:
raise ToolInvocationError(
f"Failed to invoke Tool `{self.name}` with parameters {kwargs}. Error: {e}", tool_name=self.name
haystack/tools/toolset.py
+3 -35
View File
@@ -117,12 +117,12 @@ class Toolset:
"""
Return the tools available for name-based selection (e.g. via `Agent.run(tools=["tool_name"])`).
Warms up the Toolset first, so lazily loaded tools are selectable too. Subclasses whose iteration does
not surface every selectable tool (e.g. SearchableToolset) override this to return the full set.
Prepare lazy toolsets with `warm_up_tools` or `warm_up_tools_async` before selection.
Subclasses whose iteration does not surface every selectable tool (e.g. SearchableToolset)
override this to return the full set.
:returns: The list of tools available for name-based selection.
"""
self.warm_up()
return list(self.tools)
def spawn(self, selected_tool_names: set[str] | None = None) -> "Toolset": # noqa: ARG002
@@ -156,38 +156,6 @@ class Toolset:
return any(tool is item or tool == item for tool in self)
return False
def warm_up(self) -> None:
"""
Prepare the Toolset for use.
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
- Loading tools dynamically from an external source and assigning them to `self.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 and load the tools through it:
```python
class MCPToolset(Toolset):
def warm_up(self) -> None:
if self.mcp_connection is not None:
return
self.mcp_connection = establish_connection(self.server_url)
self.tools = self.mcp_connection.fetch_tools()
```
This method may be called multiple times (e.g. before every run): implementations are responsible for
their own idempotence, guarding on their own state as in the example above. The default implementation delegates
to the tools' own idempotent `warm_up()`.
"""
for tool in self.tools:
if hasattr(tool, "warm_up"):
tool.warm_up()
def add(self, tool: Tool) -> None:
"""
Add a new Tool to this Toolset.
haystack/tools/utils.py
+67 -15
View File
@@ -2,6 +2,7 @@
#
# SPDX-License-Identifier: Apache-2.0
from collections.abc import Sequence
from typing import TYPE_CHECKING
from haystack.tools.tool import Tool
@@ -11,30 +12,81 @@ if TYPE_CHECKING:
from haystack.tools import ToolsType
def _as_tool_sequence(tools: "ToolsType | None") -> Sequence[Tool | Toolset]:
if tools is None:
return []
if isinstance(tools, Toolset):
return [tools]
return tools
def warm_up_tools(tools: "ToolsType | None" = 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()
For Toolset objects, this delegates to warm_up() if implemented; otherwise,
it calls warm_up() on each tool that implements it. Toolset subclasses can implement 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.
:param tools: A sequence of Tool and/or Toolset objects, a single Toolset, or None.
"""
if tools is None:
return
for item in _as_tool_sequence(tools):
if hasattr(item, "warm_up"):
item.warm_up()
elif isinstance(item, Toolset):
warm_up_tools(tools=item.tools)
# 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 hasattr(item, "warm_up"):
item.warm_up()
async def warm_up_tools_async(tools: "ToolsType | None" = None) -> None:
"""
Warm up tools asynchronously from various formats (Tools, Toolsets, or mixed lists).
Call warm_up_async() if implemented, otherwise warm_up() if available.
If a Toolset implements neither method, this function warms up its tools.
:param tools: A sequence of Tool and/or Toolset objects, a single Toolset, or None.
"""
for item in _as_tool_sequence(tools):
if hasattr(item, "warm_up_async"):
await item.warm_up_async()
elif hasattr(item, "warm_up"):
item.warm_up()
elif isinstance(item, Toolset):
await warm_up_tools_async(tools=item.tools)
def close_tools(tools: "ToolsType | None" = None) -> None:
"""
Close tools from various formats (Tools, Toolsets, or mixed lists).
For Toolset objects, this delegates to close() if implemented; otherwise,
it calls close() on each tool that implements it.
:param tools: A sequence of Tool and/or Toolset objects, a single Toolset, or None.
"""
for item in _as_tool_sequence(tools):
if hasattr(item, "close"):
item.close()
elif isinstance(item, Toolset):
close_tools(tools=item.tools)
async def close_tools_async(tools: "ToolsType | None" = None) -> None:
"""
Close tools asynchronously from various formats (Tools, Toolsets, or mixed lists).
Call close_async() if implemented, otherwise close() if available.
If a Toolset implements neither method, this function closes its tools.
:param tools: A sequence of Tool and/or Toolset objects, a single Toolset, or None.
"""
for item in _as_tool_sequence(tools):
if hasattr(item, "close_async"):
await item.close_async()
elif hasattr(item, "close"):
item.close()
elif isinstance(item, Toolset):
await close_tools_async(tools=item.tools)
def flatten_tools_or_toolsets(tools: "ToolsType | None") -> list[Tool]:
releasenotes/notes/tool-resource-lifecycle-76155f6f1d431757.yaml
+18
View File
@@ -0,0 +1,18 @@
---
enhancements:
- |
Agent now uses its ``close()``/``close_async()`` methods to release its configured tools and supports asynchronous
tool preparation using ``warm_up_async()``.
Invoking a tool directly through ``tool.invoke()`` or ``tool.invoke_async()``
automatically calls its warm-up methods when available.
This enables tools to manage synchronous and asynchronous resources separately.
We'll update tools in Haystack Core Integrations to take advantage of this support.
upgrade:
- |
The base ``Tool.warm_up()`` and ``Toolset.warm_up()`` methods were removed.
OpenAI/Azure Chat Generators no longer warm up tools.
If you use tools through an Agent, no changes are needed. For manual preparation, use ``warm_up_tools`` or
``warm_up_tools_async``. Lazy toolsets such as ``MCPToolset`` must be warmed up before passing them to a Chat
Generator.
test/components/agents/test_agent.py
+135 -104
View File
@@ -34,7 +34,7 @@ from haystack.dataclasses.chat_message import ChatRole, TextContent
from haystack.dataclasses.streaming_chunk import StreamingChunk
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.hooks import hook
from haystack.tools import ComponentTool, Tool
from haystack.tools import ComponentTool, Tool, warm_up_tools, warm_up_tools_async
from haystack.tools.toolset import Toolset
from haystack.utils import Secret
@@ -1922,21 +1922,23 @@ class _TrackingTool(Tool):
was_warmed_up: bool = False
def warm_up(self) -> None:
super().warm_up()
self.was_warmed_up = True
class _TrackingToolset(Toolset):
was_warmed_up: bool = False
was_warmed_up_async: bool = False
def warm_up(self) -> None:
super().warm_up()
warm_up_tools(tools=self.tools)
self.was_warmed_up = True
async def warm_up_async(self) -> None:
await warm_up_tools_async(tools=self.tools)
self.was_warmed_up_async = True
class TestAgentWarmUp:
"""Tests that Agent.warm_up() correctly warms up tools and toolsets."""
class TestComponentLifecycle:
def _make_tracking_tool(self, name: str = "test_tool") -> _TrackingTool:
return _TrackingTool(
name=name,
@@ -1945,23 +1947,9 @@ class TestAgentWarmUp:
function=lambda: "result",
)
def _make_tracking_toolset(self, tools: list[Tool]) -> _TrackingToolset:
return _TrackingToolset(tools)
def test_warm_up_multiple_tools(self):
tool1 = self._make_tracking_tool("tool1")
tool2 = self._make_tracking_tool("tool2")
agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=[tool1, tool2])
assert not tool1.was_warmed_up
assert not tool2.was_warmed_up
agent.warm_up()
assert tool1.was_warmed_up
assert tool2.was_warmed_up
def test_warm_up_toolset(self):
inner_tool = self._make_tracking_tool()
toolset = self._make_tracking_toolset([inner_tool])
toolset = _TrackingToolset([inner_tool])
agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=toolset)
assert not toolset.was_warmed_up
@@ -1969,47 +1957,15 @@ class TestAgentWarmUp:
assert toolset.was_warmed_up
def test_warm_up_mixed_list_of_tools_and_toolsets(self):
tool1 = self._make_tracking_tool("standalone_tool1")
tool2 = self._make_tracking_tool("standalone_tool2")
tool3 = self._make_tracking_tool("toolset_tool1")
toolset1 = self._make_tracking_toolset([tool3])
tool4 = self._make_tracking_tool("toolset_tool2")
toolset2 = self._make_tracking_toolset([tool4])
standalone_tool = self._make_tracking_tool("standalone")
toolset = _TrackingToolset([self._make_tracking_tool("child")])
agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=[standalone_tool, toolset])
agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=[tool1, toolset1, tool2, toolset2])
assert not tool1.was_warmed_up
assert not tool2.was_warmed_up
assert not toolset1.was_warmed_up
assert not toolset2.was_warmed_up
assert not standalone_tool.was_warmed_up
assert not toolset.was_warmed_up
agent.warm_up()
assert tool1.was_warmed_up
assert tool2.was_warmed_up
assert toolset1.was_warmed_up
assert toolset2.was_warmed_up
def test_warm_up_rewarms_tools_on_every_call(self, monkeypatch):
call_count = {"n": 0}
tool = Tool(
name="counting_tool",
description="A tool that counts warm_up calls",
parameters={"type": "object", "properties": {}},
function=lambda: "test",
)
original = tool.warm_up
def counting_warm_up():
original()
call_count["n"] += 1
monkeypatch.setattr(tool, "warm_up", counting_warm_up)
agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=[tool])
agent.warm_up()
agent.warm_up()
agent.warm_up()
assert call_count["n"] == 3
assert standalone_tool.was_warmed_up
assert toolset.was_warmed_up
@pytest.mark.parametrize(
"initial_tools",
@@ -2105,14 +2061,55 @@ class TestAgentWarmUp:
assert result["messages"][2].tool_call_result.result == "2024-12-01T12:00:00Z"
assert result["last_message"].text == "done"
def test_run_warms_up_per_run_toolset(self):
def test_run_warms_empty_toolset_before_selection_by_name(self):
actual_tool = self._make_tracking_tool("get_time")
class LazyToolset(Toolset):
def warm_up(self):
self.tools = [actual_tool]
toolset = LazyToolset([])
generator = MockChatGenerator(
[ChatMessage.from_assistant(tool_calls=[ToolCall(tool_name="get_time", arguments={})]), "done"]
)
agent = Agent(chat_generator=generator, tools=toolset)
result = agent.run(messages=[ChatMessage.from_user("What time is it?")], tools=["get_time"])
assert toolset.tools == [actual_tool]
assert result["messages"][2].tool_call_result.result == "result"
assert result["last_message"].text == "done"
@pytest.mark.asyncio
async def test_run_async_warms_empty_toolset_before_selection_by_name(self):
actual_tool = self._make_tracking_tool("get_time")
class LazyToolset(Toolset):
async def warm_up_async(self):
self.tools = [actual_tool]
toolset = LazyToolset([])
generator = MockChatGenerator(
[ChatMessage.from_assistant(tool_calls=[ToolCall(tool_name="get_time", arguments={})]), "done"]
)
agent = Agent(chat_generator=generator, tools=toolset)
result = await agent.run_async(messages=[ChatMessage.from_user("What time is it?")], tools=["get_time"])
assert toolset.tools == [actual_tool]
assert result["messages"][2].tool_call_result.result == "result"
assert result["last_message"].text == "done"
def test_run_warms_up_but_does_not_close_per_run_toolset(self, monkeypatch):
"""Per-run tools passed to run() are not covered by Agent.warm_up() and must be warmed up at run time."""
init_tool = self._make_tracking_tool("init_tool")
agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=Toolset([init_tool]))
agent.warm_up()
per_run_tool = self._make_tracking_tool("per_run_tool")
per_run_toolset = self._make_tracking_toolset([per_run_tool])
per_run_toolset = _TrackingToolset([per_run_tool])
toolset_close = MagicMock()
monkeypatch.setattr(per_run_toolset, "close", toolset_close, raising=False)
assert not per_run_toolset.was_warmed_up
assert not per_run_tool.was_warmed_up
@@ -2120,6 +2117,8 @@ class TestAgentWarmUp:
assert per_run_toolset.was_warmed_up
assert per_run_tool.was_warmed_up
agent.close()
toolset_close.assert_not_called()
def test_run_warms_up_per_run_list_of_tools_and_toolsets(self):
"""A per-run list of Tools and Toolsets must be warmed up at run time."""
@@ -2129,7 +2128,7 @@ class TestAgentWarmUp:
per_run_tool = self._make_tracking_tool("per_run_tool")
toolset_tool = self._make_tracking_tool("toolset_tool")
per_run_toolset = self._make_tracking_toolset([toolset_tool])
per_run_toolset = _TrackingToolset([toolset_tool])
per_run_tools: list[Tool | Toolset] = [per_run_tool, per_run_toolset]
agent.run(messages=[ChatMessage.from_user("hi")], tools=per_run_tools)
@@ -2139,86 +2138,118 @@ class TestAgentWarmUp:
assert toolset_tool.was_warmed_up
@pytest.mark.asyncio
async def test_run_async_warms_up_per_run_toolset(self):
async def test_run_async_warms_up_but_does_not_close_per_run_toolset(self, monkeypatch):
"""The async run path must also warm up per-run tools."""
init_tool = self._make_tracking_tool("init_tool")
agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=Toolset([init_tool]))
agent.warm_up()
per_run_tool = self._make_tracking_tool("per_run_tool")
per_run_toolset = self._make_tracking_toolset([per_run_tool])
per_run_toolset = _TrackingToolset([per_run_tool])
toolset_close_async = AsyncMock()
monkeypatch.setattr(per_run_toolset, "close_async", toolset_close_async, raising=False)
await agent.run_async(messages=[ChatMessage.from_user("hi")], tools=per_run_toolset)
assert per_run_toolset.was_warmed_up
assert per_run_toolset.was_warmed_up_async
assert per_run_tool.was_warmed_up
assert not per_run_toolset.was_warmed_up
await agent.close_async()
toolset_close_async.assert_not_called()
class TestComponentLifecycle:
def test_warm_up_delegates_to_chat_generator(self, weather_tool, monkeypatch):
def test_warm_up_delegates_to_tools_and_chat_generator(self, monkeypatch):
chat_generator = MockChatGenerator("Hello")
warm_up_mock = MagicMock()
monkeypatch.setattr(chat_generator, "warm_up", warm_up_mock)
agent = Agent(chat_generator=chat_generator, tools=[weather_tool], system_prompt="This is a system prompt.")
generator_warm_up = MagicMock()
monkeypatch.setattr(chat_generator, "warm_up", generator_warm_up)
tool = MagicMock(spec=Tool, warm_up=MagicMock())
agent = Agent(chat_generator=chat_generator, tools=[tool], system_prompt="This is a system prompt.")
agent.warm_up()
warm_up_mock.assert_called_once()
generator_warm_up.assert_called_once_with()
tool.warm_up.assert_called_once_with()
warm_up_mock.reset_mock()
agent.run([ChatMessage.from_user("What is the weather in Berlin?")])
# warm_up runs twice here: the Agent delegates to the generator, and the generator's own run() self-warms
assert warm_up_mock.call_count == 2
agent.warm_up()
assert generator_warm_up.call_count == 2
assert tool.warm_up.call_count == 2
def test_run_warms_up_tools_and_chat_generator(self, monkeypatch, weather_tool):
chat_generator = MockChatGenerator("Hello")
generator_warm_up = MagicMock()
monkeypatch.setattr(chat_generator, "warm_up", generator_warm_up)
tool_warm_up = MagicMock()
monkeypatch.setattr(weather_tool, "warm_up", tool_warm_up, raising=False)
agent = Agent(chat_generator=chat_generator, tools=[weather_tool])
agent.run(messages=[ChatMessage.from_user("What is the weather in Berlin?")])
generator_warm_up.assert_called_with()
tool_warm_up.assert_called_once_with()
@pytest.mark.asyncio
async def test_warm_up_async_delegates_to_chat_generator(self, monkeypatch):
async def test_warm_up_async_delegates_to_tools_and_chat_generator(self, monkeypatch):
chat_generator = MockChatGenerator("Hello")
warm_up_async_mock = AsyncMock()
warm_up_mock = MagicMock()
monkeypatch.setattr(chat_generator, "warm_up_async", warm_up_async_mock, raising=False)
monkeypatch.setattr(chat_generator, "warm_up", warm_up_mock)
agent = Agent(chat_generator=chat_generator, tools=[])
generator_warm_up_async = AsyncMock()
generator_warm_up = MagicMock()
monkeypatch.setattr(chat_generator, "warm_up_async", generator_warm_up_async, raising=False)
monkeypatch.setattr(chat_generator, "warm_up", generator_warm_up)
tool = MagicMock(spec=Tool, warm_up=MagicMock(), warm_up_async=AsyncMock())
agent = Agent(chat_generator=chat_generator, tools=[tool])
await agent.warm_up_async()
warm_up_async_mock.assert_awaited_once()
warm_up_mock.assert_not_called()
generator_warm_up_async.assert_awaited_once_with()
generator_warm_up.assert_not_called()
tool.warm_up_async.assert_awaited_once_with()
tool.warm_up.assert_not_called()
@pytest.mark.asyncio
async def test_warm_up_async_falls_back_to_sync_warm_up(self, monkeypatch):
chat_generator = MockChatGeneratorWithoutRunAsync()
warm_up_mock = MagicMock()
monkeypatch.setattr(chat_generator, "warm_up", warm_up_mock, raising=False)
agent = Agent(chat_generator=chat_generator, tools=[])
generator_warm_up = MagicMock()
monkeypatch.setattr(chat_generator, "warm_up", generator_warm_up, raising=False)
tool = MagicMock(spec=Tool, warm_up=MagicMock())
agent = Agent(chat_generator=chat_generator, tools=[tool])
await agent.warm_up_async()
warm_up_mock.assert_called_once()
generator_warm_up.assert_called_once_with()
tool.warm_up.assert_called_once_with()
def test_close_delegates_to_chat_generator(self, monkeypatch):
def test_close_delegates_to_tools_and_chat_generator(self, monkeypatch):
chat_generator = MockChatGenerator("Hello")
close_mock = MagicMock()
monkeypatch.setattr(chat_generator, "close", close_mock, raising=False)
agent = Agent(chat_generator=chat_generator, tools=[])
generator_close = MagicMock()
monkeypatch.setattr(chat_generator, "close", generator_close, raising=False)
tool = MagicMock(spec=Tool, close=MagicMock())
agent = Agent(chat_generator=chat_generator, tools=[tool])
agent.close()
close_mock.assert_called_once()
generator_close.assert_called_once_with()
tool.close.assert_called_once_with()
@pytest.mark.asyncio
async def test_close_async_delegates_to_chat_generator(self, monkeypatch):
async def test_close_async_delegates_to_tools_and_chat_generator(self, monkeypatch):
chat_generator = MockChatGenerator("Hello")
close_async_mock = AsyncMock()
monkeypatch.setattr(chat_generator, "close_async", close_async_mock, raising=False)
agent = Agent(chat_generator=chat_generator, tools=[])
generator_close = MagicMock()
monkeypatch.setattr(chat_generator, "close", generator_close, raising=False)
generator_close_async = AsyncMock()
monkeypatch.setattr(chat_generator, "close_async", generator_close_async, raising=False)
tool = MagicMock(spec=Tool, close=MagicMock(), close_async=AsyncMock())
agent = Agent(chat_generator=chat_generator, tools=[tool])
await agent.close_async()
close_async_mock.assert_awaited_once()
generator_close_async.assert_awaited_once_with()
generator_close.assert_not_called()
tool.close_async.assert_awaited_once_with()
tool.close.assert_not_called()
@pytest.mark.asyncio
async def test_close_async_falls_back_to_sync_close(self, monkeypatch):
chat_generator = MockChatGenerator("Hello")
close_mock = MagicMock()
monkeypatch.setattr(chat_generator, "close", close_mock, raising=False)
agent = Agent(chat_generator=chat_generator, tools=[])
generator_close = MagicMock()
monkeypatch.setattr(chat_generator, "close", generator_close, raising=False)
tool = MagicMock(spec=Tool, close=MagicMock())
agent = Agent(chat_generator=chat_generator, tools=[tool])
await agent.close_async()
close_mock.assert_called_once()
generator_close.assert_called_once_with()
tool.close.assert_called_once_with()
@pytest.mark.asyncio
async def test_lifecycle_is_safe_when_chat_generator_lacks_methods(self):
agent = Agent(chat_generator=MockChatGeneratorWithoutRunAsync(), tools=[])
async def test_lifecycle_is_safe_when_tools_and_chat_generator_lack_methods(self, weather_tool):
agent = Agent(chat_generator=MockChatGeneratorWithoutRunAsync(), tools=[weather_tool])
agent.warm_up()
await agent.warm_up_async()
agent.close()
test/components/agents/test_utils.py
-19
View File
@@ -112,25 +112,6 @@ class TestSelectToolsByName:
selected = _select_tools_by_name([second_tool, toolset], [first_tool.name, second_tool.name])
assert selected == [second_tool, first_tool]
def test_warms_up_lazy_toolsets_to_resolve_names(self, first_tool: Tool, second_tool: Tool) -> None:
class LazyToolset(Toolset):
"""A Toolset that loads its tools on warm_up(), like toolsets backed by external services."""
def __init__(self, tools_to_load):
self._tools_to_load = tools_to_load
super().__init__([]) # no tools until warm_up
def warm_up(self):
if not self.tools:
self.tools = list(self._tools_to_load)
toolset = LazyToolset([first_tool, second_tool])
assert list(toolset) == [] # not loaded yet
# Name resolution warms the toolset first, so lazily loaded tools are selectable.
selected = _select_tools_by_name([toolset], [first_tool.name])
assert selected == [first_tool]
def test_spawns_toolsets_without_mutating_them(self, first_tool: Tool, second_tool: Tool) -> None:
class RunScopedToolset(Toolset):
"""A Toolset overriding spawn(), signaling run-scoped state."""
test/components/generators/chat/test_azure.py
-36
View File
@@ -708,42 +708,6 @@ class TestComponentLifecycle:
with pytest.raises(OpenAIError):
generator.warm_up()
def test_warm_up_warms_tools_once(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "fake-api-key")
warm_up_calls = []
class MockTool(Tool):
def __init__(self, tool_name):
super().__init__(
name=tool_name,
description=f"Mock tool {tool_name}",
parameters={"type": "object", "properties": {"x": {"type": "string"}}, "required": ["x"]},
function=lambda x: x,
)
def warm_up(self):
warm_up_calls.append(self.name)
generator = AzureOpenAIChatGenerator(
azure_endpoint="some-non-existing-endpoint", tools=[MockTool("tool1"), MockTool("tool2")]
)
assert not generator._tools_warmed_up
generator.warm_up()
assert sorted(warm_up_calls) == ["tool1", "tool2"]
assert generator._tools_warmed_up
generator.warm_up()
assert sorted(warm_up_calls) == ["tool1", "tool2"]
def test_warm_up_with_no_tools_does_not_raise(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "fake-api-key")
generator = AzureOpenAIChatGenerator(azure_endpoint="some-non-existing-endpoint")
generator.warm_up()
assert generator._tools_warmed_up
def test_sync_lifecycle(self, mock_azure_clients: tuple[MagicMock, MagicMock]) -> None:
sync_cls, _ = mock_azure_clients
test/components/generators/chat/test_azure_responses.py
-36
View File
@@ -417,42 +417,6 @@ class TestSerDe:
class TestComponentLifecycle:
def test_warm_up_warms_tools_once(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
warm_up_calls = []
class MockTool(Tool):
def __init__(self, tool_name):
super().__init__(
name=tool_name,
description=f"Mock tool {tool_name}",
parameters={"type": "object", "properties": {"x": {"type": "string"}}, "required": ["x"]},
function=lambda x: x,
)
def warm_up(self):
warm_up_calls.append(self.name)
component = AzureOpenAIResponsesChatGenerator(
azure_endpoint="some-non-existing-endpoint", tools=[MockTool("tool1"), MockTool("tool2")]
)
assert not component._tools_warmed_up
component.warm_up()
assert sorted(warm_up_calls) == ["tool1", "tool2"]
assert component._tools_warmed_up
component.warm_up()
assert sorted(warm_up_calls) == ["tool1", "tool2"]
def test_warm_up_with_no_tools_does_not_raise(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
component = AzureOpenAIResponsesChatGenerator(azure_endpoint="some-non-existing-endpoint")
component.warm_up()
assert component._tools_warmed_up
def test_sync_lifecycle(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
test/components/generators/chat/test_openai.py
-34
View File
@@ -1791,40 +1791,6 @@ class TestComponentLifecycle:
with pytest.raises(ValueError, match="None of the .* environment variables are set"):
generator.warm_up()
def test_warm_up_warms_tools_once(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key")
warm_up_calls = []
class MockTool(Tool):
def __init__(self, tool_name):
super().__init__(
name=tool_name,
description=f"Mock tool {tool_name}",
parameters={"type": "object", "properties": {"x": {"type": "string"}}, "required": ["x"]},
function=lambda x: x,
)
def warm_up(self):
warm_up_calls.append(self.name)
generator = OpenAIChatGenerator(tools=[MockTool("tool1"), MockTool("tool2")])
assert not generator._tools_warmed_up
generator.warm_up()
assert sorted(warm_up_calls) == ["tool1", "tool2"]
assert generator._tools_warmed_up
generator.warm_up()
assert sorted(warm_up_calls) == ["tool1", "tool2"]
def test_warm_up_with_no_tools_does_not_raise(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key")
generator = OpenAIChatGenerator()
generator.warm_up()
assert generator._tools_warmed_up
def test_sync_lifecycle(self, mock_openai_clients: tuple[MagicMock, MagicMock]) -> None:
sync_cls, _ = mock_openai_clients
test/components/generators/chat/test_openai_responses.py
-61
View File
@@ -436,67 +436,6 @@ class TestComponentLifecycle:
with pytest.raises(ValueError, match="None of the .* environment variables are set"):
generator.warm_up()
def test_warm_up_warms_tools_once(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key")
warm_up_calls = []
class MockTool(Tool):
def __init__(self, tool_name):
super().__init__(
name=tool_name,
description=f"Mock tool {tool_name}",
parameters={"type": "object", "properties": {"x": {"type": "string"}}, "required": ["x"]},
function=lambda x: x,
)
def warm_up(self):
warm_up_calls.append(self.name)
generator = OpenAIResponsesChatGenerator(tools=[MockTool("tool1"), MockTool("tool2")])
assert not generator._tools_warmed_up
generator.warm_up()
assert sorted(warm_up_calls) == ["tool1", "tool2"]
assert generator._tools_warmed_up
generator.warm_up()
assert sorted(warm_up_calls) == ["tool1", "tool2"]
def test_warm_up_with_no_tools_does_not_raise(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key")
generator = OpenAIResponsesChatGenerator()
generator.warm_up()
assert generator._tools_warmed_up
def test_warm_up_with_empty_tools_list_does_not_raise(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key")
# An empty list is a valid ``tools`` value (e.g. when built programmatically); warming up
# must not index ``tools[0]`` unguarded.
generator = OpenAIResponsesChatGenerator(tools=[])
generator.warm_up()
assert generator._tools_warmed_up
def test_warm_up_with_openai_tools_does_not_raise(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key")
generator = OpenAIResponsesChatGenerator(
tools=[
{"type": "web_search_preview"},
{
"type": "mcp",
"server_label": "dmcp",
"server_description": "A Dungeons and Dragons MCP server to assist with dice rolling.",
"server_url": "https://dmcp-server.deno.dev/sse",
"require_approval": "never",
},
]
)
generator.warm_up()
assert generator._tools_warmed_up
def test_sync_lifecycle(self, mock_openai_clients: tuple[MagicMock, MagicMock]) -> None:
sync_cls, _ = mock_openai_clients
test/tools/skills/test_skill_toolset.py
+28 -5
View File
@@ -5,6 +5,7 @@
import threading
from concurrent.futures import ThreadPoolExecutor
from unittest.mock import Mock
import pytest
@@ -53,7 +54,8 @@ class TestSkillToolset:
_write_skill(tmp_path, "pdf-forms", description="Use to fill PDF forms.")
_write_skill(tmp_path, "excel", description="Use to edit spreadsheets.")
toolset = SkillToolset(FileSystemSkillStore(tmp_path))
store = FileSystemSkillStore(tmp_path)
toolset = SkillToolset(store)
# The catalog is only scanned on warm_up.
assert toolset._is_warmed_up is False
@@ -61,6 +63,7 @@ class TestSkillToolset:
toolset.warm_up()
assert toolset._is_warmed_up is True
assert store._is_warmed_up is True
assert set(toolset.skills) == {"pdf-forms", "excel"}
assert toolset.skills["pdf-forms"].description == "Use to fill PDF forms."
@@ -71,19 +74,39 @@ class TestSkillToolset:
assert set(toolset.skills) == {"pdf-forms"}
assert toolset._is_warmed_up is True
def test_warm_up_is_idempotent(self, tmp_path):
def test_warm_up_is_idempotent(self, tmp_path, monkeypatch):
_write_skill(tmp_path, "pdf-forms", description="Use to fill PDF forms.")
toolset = SkillToolset(FileSystemSkillStore(tmp_path))
list_skills = Mock(wraps=toolset._store.list_skills)
monkeypatch.setattr(toolset._store, "list_skills", list_skills)
toolset.warm_up()
toolset.warm_up()
assert set(toolset.skills) == {"pdf-forms"}
list_skills.assert_called_once_with()
def test_warm_up_warms_up_the_store(self, tmp_path):
_write_skill(tmp_path, "pdf-forms", description="Use to fill PDF forms.")
def test_close_is_safe_before_warm_up_and_when_repeated(self):
store = Mock(spec=FileSystemSkillStore, close=Mock())
toolset = SkillToolset(store)
toolset.close()
store.close.assert_called_once_with()
toolset.close()
assert store.close.call_count == 2
def test_close_releases_store_and_invalidates_catalog(self, tmp_path, monkeypatch):
_write_skill(tmp_path, "first", description="First skill")
store = FileSystemSkillStore(tmp_path)
store_close, list_skills = Mock(), Mock(wraps=store.list_skills)
monkeypatch.setattr(store, "close", store_close, raising=False)
monkeypatch.setattr(store, "list_skills", list_skills)
toolset = SkillToolset(store)
toolset.warm_up()
assert store._is_warmed_up is True
list_skills.assert_called_once_with()
toolset.close()
store_close.assert_called_once_with()
assert toolset._skills == {}
toolset.warm_up()
assert "first" in toolset.skills
assert list_skills.call_count == 2
def test_concurrent_warm_up(self, tmp_path):
# Concurrent first use (e.g. parallel requests hitting a shared Agent) must produce a complete,
test/tools/test_component_tool.py
+31 -14
View File
@@ -6,7 +6,7 @@ import json
import os
from dataclasses import dataclass
from typing import Any
from unittest.mock import patch
from unittest.mock import AsyncMock, Mock, patch
import pytest
from openai.types.chat import ChatCompletion, ChatCompletionMessage
@@ -471,22 +471,19 @@ class TestComponentTool:
assert result["output_a"] == "A processed: test input"
assert result["output_b"] == "B processed: test input"
def test_warm_up_is_idempotent(self):
"""Test that calling warm_up multiple times only warms up the component once."""
from unittest.mock import MagicMock, patch
def test_warm_up_delegates_to_component(self):
component = SimpleComponent()
tool = ComponentTool(component=component)
with patch.object(component, "warm_up") as component_warm_up:
tool.warm_up()
tool.warm_up()
assert component_warm_up.call_count == 2
with patch.object(component, "warm_up", MagicMock()) as mock_warm_up:
# Call warm_up multiple times
tool.warm_up()
tool.warm_up()
tool.warm_up()
# Component's warm_up should only be called once
mock_warm_up.assert_called_once()
def test_close_delegates_to_component(self):
component = SimpleComponent()
with patch.object(component, "close", create=True) as component_close:
ComponentTool(component=component).close()
component_close.assert_called_once_with()
def test_from_component_with_callable_params_skipped(self, monkeypatch):
monkeypatch.setenv("OPENAI_API_KEY", "test-api-key")
@@ -922,3 +919,23 @@ class TestComponentToolAsync:
def test_invoke_uses_run_on_dual_mode_component(self, dual_tool):
assert dual_tool.invoke(text="hi") == {"reply": "sync:hi"}
async def test_async_lifecycle_delegates_to_component(self, dual_tool, monkeypatch):
component_warm_up_async, component_close_async = AsyncMock(), AsyncMock()
monkeypatch.setattr(dual_tool._component, "warm_up_async", component_warm_up_async, raising=False)
monkeypatch.setattr(dual_tool._component, "close_async", component_close_async, raising=False)
await dual_tool.warm_up_async()
await dual_tool.close_async()
component_warm_up_async.assert_awaited_once_with()
component_close_async.assert_awaited_once_with()
async def test_sync_component_uses_sync_lifecycle(self, sync_tool, monkeypatch):
component_warm_up, component_close = Mock(), Mock()
monkeypatch.setattr(sync_tool._component, "warm_up", component_warm_up, raising=False)
monkeypatch.setattr(sync_tool._component, "close", component_close, raising=False)
monkeypatch.setattr(sync_tool._component, "warm_up_async", AsyncMock(side_effect=AssertionError), raising=False)
monkeypatch.setattr(sync_tool._component, "close_async", AsyncMock(side_effect=AssertionError), raising=False)
await sync_tool.warm_up_async()
await sync_tool.close_async()
component_warm_up.assert_called_once_with()
component_close.assert_called_once_with()
test/tools/test_searchable_toolset.py
+47 -10
View File
@@ -6,6 +6,7 @@
import os
from collections.abc import Callable
from typing import Annotated, Any
from unittest.mock import AsyncMock, Mock
import pytest
@@ -517,9 +518,7 @@ class TestSearchableToolsetWithToolset:
assert any(t.name == "multiply_numbers" for t in search_toolset._catalog)
class TestSearchableToolsetWarmUp:
"""Tests for warm_up behavior."""
class TestSearchableToolsetLifecycle:
def test_not_warmed_up_after_agent_init(self, large_catalog, monkeypatch):
"""Initializing an Agent with a SearchableToolset must not warm it up (no premature flatten/connect)."""
monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
@@ -534,6 +533,7 @@ class TestSearchableToolsetWarmUp:
"""Test that warm_up can be called multiple times safely."""
toolset = SearchableToolset(catalog=large_catalog)
assert toolset._bootstrap_tool is None
toolset.warm_up()
first_bootstrap = toolset._bootstrap_tool
@@ -554,12 +554,6 @@ class TestSearchableToolsetWarmUp:
toolset.warm_up()
assert len(list(toolset)) == len(small_catalog)
def test_bootstrap_tool_before_warm_up(self, large_catalog):
"""Test that bootstrap tool is None before warm_up."""
toolset = SearchableToolset(catalog=large_catalog)
assert toolset._bootstrap_tool is None
def test_warm_up_raises_on_duplicate_tool_names(self):
"""Test that warm_up raises when the flattened catalog has duplicate tool names."""
params = {"type": "object", "properties": {}}
@@ -570,6 +564,48 @@ class TestSearchableToolsetWarmUp:
with pytest.raises(ValueError, match="Duplicate tool names found"):
toolset.warm_up()
async def test_spawn_and_iteration_do_not_repeat_warm_up_after_async_warm_up(self, large_catalog, monkeypatch):
catalog = Toolset(large_catalog)
catalog_warm_up_async = AsyncMock()
monkeypatch.setattr(catalog, "warm_up_async", catalog_warm_up_async, raising=False)
monkeypatch.setattr(catalog, "warm_up", Mock(side_effect=AssertionError), raising=False)
toolset = SearchableToolset(catalog=catalog)
await toolset.warm_up_async()
spawn = toolset.spawn()
search_tool = next(iter(spawn))
assert search_tool.name == "search_tools"
catalog_warm_up_async.assert_awaited_once_with()
def test_close_clears_discovery_and_allows_warm_up_again(self, large_catalog, monkeypatch):
catalog = Toolset(large_catalog)
catalog_close = Mock()
monkeypatch.setattr(catalog, "close", catalog_close, raising=False)
toolset = SearchableToolset(catalog=catalog)
search_tool = next(iter(toolset))
search_tool.invoke(tool_keywords="weather")
assert toolset._discovered_tools
toolset.close()
catalog_close.assert_called_once_with()
assert toolset._discovered_tools == {}
catalog.tools = [large_catalog[0]]
toolset.warm_up()
assert list(toolset) == catalog.tools
async def test_close_async_releases_catalog_and_allows_warm_up_again(self, small_catalog, monkeypatch):
tool_close_async = AsyncMock()
monkeypatch.setattr(small_catalog[0], "close_async", tool_close_async, raising=False)
toolset = SearchableToolset(catalog=small_catalog, search_threshold=1)
await toolset.warm_up_async()
search_tool = next(iter(toolset))
search_tool.invoke(tool_keywords="weather")
assert toolset._discovered_tools
await toolset.close_async()
tool_close_async.assert_awaited_once_with()
assert toolset._discovered_tools == {}
small_catalog.pop()
await toolset.warm_up_async()
assert toolset.get_selectable_tools() == small_catalog
class TestSearchableToolsetEdgeCases:
"""Tests for edge cases and error handling."""
@@ -831,6 +867,7 @@ class TestSearchableToolsetAgentToolSelection:
toolset = SearchableToolset(catalog=large_catalog, search_threshold=3) # 8 tools -> search mode
agent = Agent(chat_generator=OpenAIChatGenerator(), tools=toolset)
toolset.warm_up()
selected = agent._select_tools(["get_weather", "add_numbers"])
# An isolated spawn is returned with the selection; the configured toolset is not mutated.
@@ -852,7 +889,6 @@ class TestSearchableToolsetAgentToolSelection:
def test_spawns_have_independent_discovered_tools_and_selection(self, large_catalog):
"""Two spawns of one SearchableToolset don't share discovered tools or collide on the active selection."""
toolset = SearchableToolset(catalog=large_catalog, search_threshold=3)
toolset.warm_up()
spawn_a = toolset.spawn(selected_tool_names={"get_weather"})
spawn_b = toolset.spawn()
@@ -879,6 +915,7 @@ class TestSearchableToolsetAgentToolSelection:
toolset = SearchableToolset(catalog=large_catalog, search_threshold=20) # 8 < 20 -> passthrough
agent = Agent(chat_generator=OpenAIChatGenerator(), tools=toolset)
toolset.warm_up()
selected = agent._select_tools(["get_weather", "add_numbers"])
assert selected == [toolset]
test/tools/test_tool.py
+32 -6
View File
@@ -3,7 +3,9 @@
# SPDX-License-Identifier: Apache-2.0
import re
import threading
from typing import Any
from unittest.mock import AsyncMock, Mock
import pytest
@@ -139,12 +141,17 @@ class TestTool:
assert tool.tool_spec == {"name": "weather", "description": "Get weather report", "parameters": parameters}
def test_invoke(self):
tool = Tool(
name="weather", description="Get weather report", parameters=parameters, function=get_weather_report
)
def test_invoke(self, monkeypatch):
calls = []
def invoke(city):
calls.append("invoke")
return get_weather_report(city)
tool = Tool(name="weather", description="Get weather report", parameters=parameters, function=invoke)
monkeypatch.setattr(tool, "warm_up", lambda: calls.append("warm_up"), raising=False)
assert tool.invoke(city="Berlin") == "Weather report for Berlin: 20°C, sunny"
assert calls == ["warm_up", "invoke"]
def test_invoke_fail(self):
tool = Tool(
@@ -376,19 +383,33 @@ def sync_tool():
class TestToolAsync:
@pytest.mark.asyncio
async def test_invoke_async_awaits_async_function(self, async_tool):
async def test_invoke_async_awaits_async_function(self, async_tool, monkeypatch):
warm_up_async = AsyncMock()
monkeypatch.setattr(async_tool, "warm_up", Mock(side_effect=AssertionError), raising=False)
monkeypatch.setattr(async_tool, "warm_up_async", warm_up_async, raising=False)
assert async_tool.function is None
assert async_tool.async_function is async_get_weather
assert await async_tool.invoke_async(city="Berlin") == "Weather report for Berlin: 20°C, sunny"
warm_up_async.assert_awaited_once_with()
async def test_async_invocation_falls_back_to_sync_warm_up(self, async_tool, monkeypatch):
warm_up = Mock()
monkeypatch.setattr(async_tool, "warm_up", warm_up, raising=False)
await async_tool.invoke_async(city="Berlin")
warm_up.assert_called_once_with()
def test_invoke_on_async_only_tool_raises(self, async_tool):
with pytest.raises(ToolInvocationError, match=re.escape("has no sync `function`")):
async_tool.invoke(city="Berlin")
@pytest.mark.asyncio
async def test_invoke_async_falls_back_to_sync_function(self, sync_tool):
async def test_invoke_async_falls_back_to_sync_function(self, sync_tool, monkeypatch):
threads = []
monkeypatch.setattr(sync_tool, "warm_up", lambda: threads.append(threading.get_ident()), raising=False)
# Sync-only tool: invoke_async dispatches to a worker thread via asyncio.to_thread.
assert await sync_tool.invoke_async(city="Berlin") == "Weather report for Berlin: 20°C, sunny"
assert len(threads) == 1
assert threads[0] != threading.get_ident()
@pytest.mark.asyncio
async def test_async_function_is_preferred_when_both_set(self):
@@ -419,6 +440,11 @@ class TestToolAsync:
with pytest.raises(ToolInvocationError, match="kaboom"):
await tool.invoke_async(city="Berlin")
async def test_warm_up_error_is_wrapped(self, sync_tool, monkeypatch):
monkeypatch.setattr(sync_tool, "warm_up", Mock(side_effect=ValueError("setup failed")), raising=False)
with pytest.raises(ToolInvocationError, match="setup failed"):
await sync_tool.invoke_async(city="Berlin")
@pytest.mark.parametrize(
"kwargs, expected_function, expected_async_function",
[
test/tools/test_tools_utils.py
+86 -216
View File
@@ -2,9 +2,12 @@
#
# SPDX-License-Identifier: Apache-2.0
from unittest.mock import AsyncMock, Mock
import pytest
from haystack.tools import Tool, Toolset, flatten_tools_or_toolsets, warm_up_tools
from haystack.tools.utils import close_tools, close_tools_async, warm_up_tools_async
def add_numbers(a: int, b: int) -> int:
@@ -173,226 +176,93 @@ class TestFlattenToolsOrToolsets:
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
def test_ignores_none_and_tools_without_warm_up(self, add_tool):
warm_up_tools(None)
warm_up_tools([add_tool])
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
def test_warms_up_tools_in_mixed_list(self, add_tool, multiply_tool, subtract_tool, monkeypatch):
add_warm_up, multiply_warm_up, subtract_warm_up = Mock(), Mock(), Mock()
monkeypatch.setattr(add_tool, "warm_up", add_warm_up, raising=False)
monkeypatch.setattr(multiply_tool, "warm_up", multiply_warm_up, raising=False)
monkeypatch.setattr(subtract_tool, "warm_up", subtract_warm_up, raising=False)
warm_up_tools([add_tool, Toolset([multiply_tool]), Toolset([subtract_tool])])
add_warm_up.assert_called_once_with()
multiply_warm_up.assert_called_once_with()
subtract_warm_up.assert_called_once_with()
def test_delegates_each_call_to_custom_toolset(self):
tool = Mock(spec=Tool, warm_up=Mock())
toolset = Mock(spec=Toolset, tools=[tool], warm_up=Mock())
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 did real work."""
def __init__(self, tools):
super().__init__(tools)
self.warm_up_count = 0
self._loaded = False
def warm_up(self):
if self._loaded:
return
self.warm_up_count += 1
self._loaded = True
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)
toolset.warm_up.assert_called_once_with()
warm_up_tools(toolset)
assert toolset.warm_up.call_count == 2
tool.warm_up.assert_not_called()
# warm_up is idempotent, so the toolset and its tools are only warmed up once
assert toolset.warm_up_count == 1
assert tool.warm_up_count == 1
class TestWarmUpToolsAsync:
async def test_ignores_none_and_tools_without_warm_up(self, add_tool):
await warm_up_tools_async(None)
await warm_up_tools_async([add_tool])
async def test_prefers_async_warm_up_for_children(self):
tool = Mock(spec=Tool, warm_up=Mock(), warm_up_async=AsyncMock())
tool.name = "lookup"
await warm_up_tools_async(Toolset([tool]))
tool.warm_up_async.assert_awaited_once_with()
tool.warm_up.assert_not_called()
async def test_falls_back_to_sync_warm_up_for_children(self):
tool = Mock(spec=Tool, warm_up=Mock())
tool.name = "lookup"
await warm_up_tools_async(Toolset([tool]))
tool.warm_up.assert_called_once_with()
async def test_falls_back_to_custom_toolset_warm_up(self):
tool = Mock(spec=Tool, warm_up=Mock())
toolset = Mock(spec=Toolset, tools=[tool], warm_up=Mock())
await warm_up_tools_async(toolset)
toolset.warm_up.assert_called_once_with()
tool.warm_up.assert_not_called()
class TestCloseTools:
def test_ignores_none_and_tools_without_close(self, add_tool):
close_tools(None)
close_tools([add_tool])
def test_closes_children_of_plain_toolset(self):
tool = Mock(spec=Tool, close=Mock())
tool.name = "lookup"
close_tools(Toolset([tool]))
tool.close.assert_called_once_with()
def test_custom_toolset_owns_cleanup(self):
tool = Mock(spec=Tool, close=Mock())
toolset = Mock(spec=Toolset, tools=[tool], close=Mock())
close_tools(toolset)
toolset.close.assert_called_once_with()
tool.close.assert_not_called()
class TestCloseToolsAsync:
async def test_ignores_none_and_tools_without_close(self, add_tool):
await close_tools_async(None)
await close_tools_async([add_tool])
async def test_prefers_async_close_for_children(self):
tool = Mock(spec=Tool, close=Mock(), close_async=AsyncMock())
tool.name = "lookup"
await close_tools_async(Toolset([tool]))
tool.close_async.assert_awaited_once_with()
tool.close.assert_not_called()
async def test_falls_back_to_custom_toolset_close(self):
tool = Mock(spec=Tool, close=Mock())
toolset = Mock(spec=Toolset, tools=[tool], close=Mock())
await close_tools_async(toolset)
toolset.close.assert_called_once_with()
tool.close.assert_not_called()
test/tools/test_toolset.py
+6 -45
View File
@@ -3,6 +3,7 @@
# SPDX-License-Identifier: Apache-2.0
from typing import Annotated, Any
from unittest.mock import Mock
import pytest
@@ -343,35 +344,13 @@ class TestToolsetWithAgent:
assert new_pipeline.to_dict() == pipeline_dict
class TestToolsetWarmUp:
"""Stress tests for Toolset warm_up behavior."""
def test_warm_up_warms_all_tools(self):
t1, t2 = WarmUpCountingTool("a"), WarmUpCountingTool("b")
toolset = Toolset([t1, t2])
assert t1.warm_up_count == 0
assert t2.warm_up_count == 0
toolset.warm_up()
assert t1.warm_up_count == 1
assert t2.warm_up_count == 1
def test_warm_up_can_be_called_multiple_times(self):
t1 = WarmUpCountingTool("a")
toolset = Toolset([t1])
toolset.warm_up()
toolset.warm_up()
toolset.warm_up()
assert t1.warm_up_count == 3
class TestToolsetAdd:
def test_add_never_warms_the_new_tool(self):
existing = WarmUpCountingTool("a")
toolset = Toolset([existing])
toolset.warm_up()
new_tool = WarmUpCountingTool("b")
toolset.add(new_tool)
assert new_tool.warm_up_count == 0
toolset.warm_up()
assert new_tool.warm_up_count == 1
def test_add_toolset_raises(self):
toolset = Toolset([WarmUpCountingTool("a")])
@@ -393,27 +372,9 @@ class TestToolsetSpawn:
class TestToolsetToolSelection:
"""Tests for get_selectable_tools()."""
def test_get_selectable_tools_returns_all_tools(self, add_tool, multiply_tool):
def test_get_selectable_tools_returns_all_tools(self, add_tool, multiply_tool, monkeypatch):
toolset = Toolset([add_tool, multiply_tool])
warm_up = Mock()
monkeypatch.setattr(toolset, "warm_up", warm_up, raising=False)
assert toolset.get_selectable_tools() == [add_tool, multiply_tool]
def test_get_selectable_tools_warms_up_lazy_toolset(self, add_tool, multiply_tool):
"""get_selectable_tools() warms up a lazy toolset so its lazily loaded tools are available for selection."""
class LazyToolset(Toolset):
def __init__(self):
self._loaded = False
super().__init__([]) # no tools until warm_up
def warm_up(self):
if self._loaded:
return
self._loaded = True
self.tools = [add_tool, multiply_tool]
toolset = LazyToolset()
assert toolset.tools == [] # not loaded yet
selectable = toolset.get_selectable_tools()
assert [tool.name for tool in selectable] == ["add", "multiply"]
warm_up.assert_not_called()