diff --git a/docs-website/docs/tools/toolset.mdx b/docs-website/docs/tools/toolset.mdx index 33ce1a8f6cc..9c5792dbaa0 100644 --- a/docs-website/docs/tools/toolset.mdx +++ b/docs-website/docs/tools/toolset.mdx @@ -73,21 +73,28 @@ def multiply_numbers( math_toolset.add(multiply_numbers) +``` + +### Combining Toolsets + +To use multiple Toolsets together, pass them as a list wherever tools are accepted: -# or, you can merge toolsets together -math_toolset.add(another_toolset) +```python +agent = Agent( + chat_generator=OpenAIChatGenerator(), tools=[math_toolset, another_toolset] +) ``` ### Run-Scoped Copies and Tool Selection -A `Toolset` is never mutated in place during an [`Agent`](../pipeline-components/agents-1/agent.mdx) run. Each run operates on an isolated, run-scoped copy of the configured `Toolset`, created with the `spawn()` method. This makes concurrent runs that share the same `Toolset` instance safe: per-run state, such as an active tool-name selection or a [`SearchableToolset`](searchabletoolset.mdx)'s discovered tools, cannot leak or collide across runs. +An [`Agent`](../pipeline-components/agents-1/agent.mdx) run never modifies your configured `Toolset`. A `Toolset` with per-run state, such as a [`SearchableToolset`](searchabletoolset.mdx), is copied for each run through its `spawn()` method, so concurrent runs cannot leak state (like discovered tools) into each other. A plain `Toolset` has no per-run state and is shared as is; just avoid adding or removing tools while runs are in progress. -You can also restrict an `Agent` to a subset of tools at runtime by passing tool names, for example `agent.run(tools=["tool_a", "tool_b"])`. When a `Toolset` is configured, the selection is applied to the live (run-scoped) `Toolset` rather than flattening it into a static list, so dynamic behavior like a `SearchableToolset`'s search and lazy loading keeps working over the selected subset. +You can also restrict an `Agent` to a subset of tools at runtime by passing tool names, for example `agent.run(tools=["tool_a", "tool_b"])`. The selection applies only to that run, and dynamic behavior like a `SearchableToolset`'s search keeps working over the selected subset. Two methods support this and can be overridden when subclassing: -- `get_selectable_tools()`: Returns every tool available for name-based selection, ignoring any active selection restriction. Override it if your subclass's iteration does not surface every selectable tool. -- `spawn()`: Returns an isolated, run-scoped copy of the `Toolset`. Override it if your subclass holds additional run-scoped state. +- `get_selectable_tools()`: Returns every tool available for name-based selection. Override it if your subclass's iteration does not surface every selectable tool. +- `spawn()`: Returns the `Toolset` itself, which has no run-scoped state to isolate. Override it to return an isolated, run-scoped copy if your subclass holds run-scoped state. ## Usage diff --git a/haystack/components/agents/agent.py b/haystack/components/agents/agent.py index f3f3147ca4e..f883f7953ca 100644 --- a/haystack/components/agents/agent.py +++ b/haystack/components/agents/agent.py @@ -51,7 +51,6 @@ warm_up_hooks_async, ) from haystack.tools import ( - Tool, Toolset, ToolsType, _check_duplicate_tool_names, @@ -478,7 +477,6 @@ def __init__( # noqa: PLR0913 self.tool_concurrency_limit = tool_concurrency_limit self.tool_streaming_callback_passthrough = tool_streaming_callback_passthrough self.hooks = hooks - self._tools_warmed_up = False self._hooks_warmed_up = False # --- State schema --- @@ -569,13 +567,6 @@ def _register_prompt_variables(self) -> None: else: component.set_input_type(self, name=var_name, type=Any, default=None) - def _warm_up_tools(self) -> None: - """Warm up the configured tools once.""" - if not self._tools_warmed_up: - if self.tools: - warm_up_tools(tools=self.tools) - self._tools_warmed_up = True - def _warm_up_hooks(self) -> None: """Warm up the configured hooks once.""" if not self._hooks_warmed_up: @@ -590,14 +581,14 @@ async def _warm_up_hooks_async(self) -> None: def warm_up(self) -> None: """Warm up the tools, hooks, and the underlying chat generator.""" - self._warm_up_tools() + warm_up_tools(tools=self.tools) self._warm_up_hooks() if hasattr(self.chat_generator, "warm_up"): self.chat_generator.warm_up() async def warm_up_async(self) -> None: """Warm up the tools, hooks, and the underlying chat generator on the serving event loop.""" - self._warm_up_tools() + warm_up_tools(tools=self.tools) await self._warm_up_hooks_async() if hasattr(self.chat_generator, "warm_up_async"): await self.chat_generator.warm_up_async() @@ -791,7 +782,7 @@ def _select_tools(self, tools: ToolsType | list[str] | None = None) -> ToolsType or if any provided tool name is not valid. :raises TypeError: If tools is not a list of Tool objects, a Toolset, or a list of tool names (strings). """ - # Toolsets are spawned into per-run copies (see _spawn_tools / _select_tools_by_name) so concurrent runs + # Toolsets are spawned per run (see _spawn_tools / _select_tools_by_name) so concurrent runs # sharing the same configured Toolset don't corrupt each other's run-scoped state. if tools is None: return _spawn_tools(tools=self.tools) @@ -799,14 +790,8 @@ def _select_tools(self, tools: ToolsType | list[str] | None = None) -> ToolsType if isinstance(tools, list) and all(isinstance(t, str) for t in tools): return _select_tools_by_name(self.tools, cast(list[str], tools)) - if isinstance(tools, Toolset): - # 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=tools) - return _spawn_tools(tools=tools) - - if isinstance(tools, list): - selected = cast(list[Tool | Toolset], tools) # mypy can't narrow the Union type from isinstance check + 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) diff --git a/haystack/components/agents/utils.py b/haystack/components/agents/utils.py index c34743a2cbc..dba07b5de73 100644 --- a/haystack/components/agents/utils.py +++ b/haystack/components/agents/utils.py @@ -90,21 +90,39 @@ def _record_tool_calls(state: State, tool_messages: list[ChatMessage]) -> None: # --------------------------- +def _spawn_selection_copy(item: Tool | Toolset, selected_tool_names: set[str]) -> Toolset | None: + """ + Return the per-run copy carrying the selection, or None if the item does not provide one. + + A Toolset with run-scoped state (e.g. SearchableToolset) overrides `spawn()` to return a copy that + applies `selected_tool_names` itself. A plain Toolset returns itself from `spawn()` (it has nothing + to isolate), and a standalone Tool has no `spawn()`: in both cases the caller applies the selection. + + :param item: A configured Tool or Toolset. + :param selected_tool_names: The tool names selected for this run. + :returns: The selection-carrying per-run copy, or None. + """ + if not isinstance(item, Toolset): + return None + spawned = item.spawn(selected_tool_names=selected_tool_names) + return spawned if spawned is not item else None + + def _select_tools_by_name(configured_tools: ToolsType, names: list[str]) -> list[Tool | Toolset]: """ Select configured tools by name for a single run. - Standalone Tools are kept when their name is requested. A Toolset that exposes a requested name is replaced by a - per-run `spawn()` (an isolated copy) with the requested names registered as its `_selected_tool_names`, so - dynamic toolsets such as SearchableToolset preserve their behavior (search/lazy-loading) over the selected subset - without mutating the shared, configured Toolset. + 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. :param configured_tools: The tools configured on the Agent. :param names: The requested tool names. - :returns: The selected standalone Tools and/or spawned, selection-scoped Toolsets. + :returns: The selected Tools and/or selection-scoped Toolset copies. :raises ValueError: If no tools were configured, or if any requested name is not a valid tool name. """ - if not configured_tools: + if configured_tools is None: raise ValueError("No tools were configured for the Agent at initialization.") requested_names = set(names) @@ -112,14 +130,16 @@ def _select_tools_by_name(configured_tools: ToolsType, names: list[str]) -> list [configured_tools] if isinstance(configured_tools, Toolset) else list(configured_tools) ) - # Resolve selectable names per item. For Toolsets we use get_selectable_tools() so dynamic toolsets - # (e.g. SearchableToolset) offer their full catalog by name, not just the tools exposed by iteration. - selectable_per_item: list[tuple[Tool | Toolset, set[str]]] = [] - valid_tool_names: set[str] = set() + # Resolve the tools each item offers for selection + selectable_per_item: list[tuple[Tool | Toolset, list[Tool]]] = [] for item in items: - item_names = {tool.name for tool in item.get_selectable_tools()} if isinstance(item, Toolset) else {item.name} - selectable_per_item.append((item, item_names)) - valid_tool_names |= item_names + selectable = item.get_selectable_tools() if isinstance(item, Toolset) else [item] + selectable_per_item.append((item, selectable)) + + valid_tool_names = {tool.name for _, selectable in selectable_per_item for tool in selectable} + # A dynamic Toolset may look empty before its catalog is resolved, so emptiness is checked here. + if not valid_tool_names: + raise ValueError("No tools were configured for the Agent at initialization.") invalid_tool_names = requested_names - valid_tool_names if invalid_tool_names: @@ -128,27 +148,28 @@ def _select_tools_by_name(configured_tools: ToolsType, names: list[str]) -> list ) selected: list[Tool | Toolset] = [] - for item, item_names in selectable_per_item: - matched = requested_names & item_names + for item, selectable in selectable_per_item: + matched = requested_names & {tool.name for tool in selectable} if not matched: continue - if isinstance(item, Toolset): - # Apply the selection to a per-run copy so the shared, configured Toolset is never mutated. - spawned = item.spawn() - spawned._selected_tool_names = matched - selected.append(spawned) + run_copy = _spawn_selection_copy(item, matched) + if run_copy is not None: + selected.append(run_copy) else: - selected.append(item) + # Select from `selectable`, the list the names were validated against: iterating a dynamic + # Toolset could silently miss tools. + selected.extend(tool for tool in selectable if tool.name in matched) return selected def _spawn_tools(tools: ToolsType) -> ToolsType: """ - Return per-run copies of `tools`, replacing each Toolset with an isolated `spawn()` (Tools are passed through). + Return per-run copies of `tools`, replacing each Toolset with its `spawn()` (Tools are passed through). This isolates run-scoped Toolset state (e.g. a SearchableToolset's discovered tools and any active name selection) so that concurrent runs sharing the same configured Toolset — such as parallel sub-agent tool calls - or concurrent requests against one Agent — don't corrupt each other. + or concurrent requests against one Agent — don't corrupt each other. A plain Toolset has no run-scoped state + and its `spawn()` returns itself unchanged. """ if isinstance(tools, Toolset): return tools.spawn() diff --git a/haystack/tools/searchable_toolset.py b/haystack/tools/searchable_toolset.py index 58809160c15..50681632336 100644 --- a/haystack/tools/searchable_toolset.py +++ b/haystack/tools/searchable_toolset.py @@ -129,16 +129,15 @@ def __init__( self._bootstrap_tool: Tool | None = None self._document_store: InMemoryDocumentStore | None = None self._passthrough: bool | None = None - self._is_warmed_up = False + + # Optional per-run name filter, set on the copies returned by spawn(). When set, iteration only + # yields tools whose name is in this set, and search is scoped to it. None means no filtering. + self._selected_tool_names: set[str] | None = None # Initialize parent with empty tools list - we manage tools dynamically super().__init__(tools=[]) - def __add__(self, other: Tool | Toolset | list[Tool]) -> "Toolset": - """Concatenation is not supported for SearchableToolset.""" - raise NotImplementedError("SearchableToolset does not support concatenation.") - - def add(self, tool: Tool | Toolset) -> None: + def add(self, tool: Tool) -> None: """Adding new tools after initialization is not supported for SearchableToolset.""" raise NotImplementedError("SearchableToolset does not support adding new tools after initialization.") @@ -153,7 +152,7 @@ def warm_up(self) -> None: :raises ValueError: If the flattened catalog contains tools with duplicate names. """ - if self._is_warmed_up: + if self._passthrough is not None: return # Warm up the catalog first (triggers lazy connections like MCPToolset), then flatten — lazy toolsets will @@ -176,8 +175,6 @@ def warm_up(self) -> None: self._document_store.write_documents(documents, policy=DuplicatePolicy.OVERWRITE) self._bootstrap_tool = self._create_search_tool() - self._is_warmed_up = True - def get_selectable_tools(self) -> list[Tool]: """ Return the full catalog of tools that can be selected by name. @@ -187,8 +184,7 @@ def get_selectable_tools(self) -> list[Tool]: :returns: The flattened catalog of tools. """ - if not self._is_warmed_up: - self.warm_up() + self.warm_up() return list(self._catalog) def clear(self) -> None: @@ -200,21 +196,23 @@ def clear(self) -> None: """ self._discovered_tools.clear() - def spawn(self) -> "SearchableToolset": + def spawn(self, selected_tool_names: set[str] | None = None) -> "SearchableToolset": """ - Return an isolated copy for a single run. + Return an isolated copy for a single run, carrying the given name selection. The copy shares the read-only catalog and BM25 index but gets fresh discovered tools and name selection, - plus a bootstrap search tool bound to the copy. This way concurrent runs sharing the same configured - SearchableToolset don't share discovered tools or collide on the active selection. + plus a bootstrap search tool bound to the copy; the selection scopes both iteration and search. This way + concurrent runs sharing the same configured SearchableToolset don't share discovered tools or collide on + the active selection. + :param selected_tool_names: Optional catalog tool names this run is restricted to. None means no + restriction. :returns: A run-scoped copy of this SearchableToolset. """ - if not self._is_warmed_up: - self.warm_up() + self.warm_up() new = copy.copy(self) new._discovered_tools = {} - new._selected_tool_names = None + new._selected_tool_names = set(selected_tool_names) if selected_tool_names is not None else None # Rebuild the bootstrap tool so its closure is bound to the copy's discovered tools / selection # rather than the original's. The document store and catalog are read-only and stay shared. if not self._passthrough: @@ -307,12 +305,10 @@ def __iter__(self) -> Iterator[Tool]: set, but the bootstrap search tool is always exposed so search keeps working over the selected subset. Automatically calls warm_up() if needed to ensure the bootstrap tool is available. """ - # Unlike base Toolset/MCPToolset, which expose a placeholder tool before warm_up, this toolset materializes - # everything (flattened catalog, bootstrap tool, passthrough decision) in warm_up. + # 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. - if not self._is_warmed_up: - self.warm_up() + self.warm_up() if self._passthrough: yield from (tool for tool in self._catalog if self._is_selected(tool.name)) else: @@ -320,11 +316,6 @@ def __iter__(self) -> Iterator[Tool]: yield self._bootstrap_tool yield from (tool for tool in self._discovered_tools.values() if self._is_selected(tool.name)) - def __len__(self) -> int: - """Return the number of currently available tools.""" - # the number of tools is computed by invoking __iter__ on the toolset - return sum(1 for _ in self) - def __contains__(self, item: str | Tool) -> bool: """ Check if a tool is available by Tool instance or tool name string. @@ -338,16 +329,6 @@ def __contains__(self, item: str | Tool) -> bool: return any(tool == item for tool in self) raise TypeError(f"Invalid item type: {type(item)}. Must be Tool or str.") - def __getitem__(self, index: int) -> Tool: - """ - Get a tool by index. - - :param index: Index of the tool to retrieve. - :returns: The tool at the given index. - :raises IndexError: If the index is out of range. - """ - return list(self)[index] - def to_dict(self) -> dict[str, Any]: """ Serialize the toolset to a dictionary. diff --git a/haystack/tools/skills/skill_toolset.py b/haystack/tools/skills/skill_toolset.py index c7ea90b6a7c..d3b46d4187b 100644 --- a/haystack/tools/skills/skill_toolset.py +++ b/haystack/tools/skills/skill_toolset.py @@ -103,20 +103,13 @@ def warm_up(self) -> None: self._load_skill_tool.description = self._load_skill_description() self._is_warmed_up = True - def add(self, tool: Tool | Toolset) -> None: + def add(self, tool: Tool) -> None: """Adding tools is not supported: a SkillToolset's tools are fixed and defined by its store.""" raise NotImplementedError( "SkillToolset does not support adding tools. To combine it with other tools, pass it to the Agent " "alongside them, e.g. tools=[skill_toolset, other_tool]." ) - def __add__(self, other: Tool | Toolset | list[Tool]) -> "Toolset": - """Concatenation is not supported for SearchableToolset.""" - raise NotImplementedError( - "SkillToolset does not support concatenation. To combine it with other tools, pass it to the Agent " - "alongside them, e.g. tools=[skill_toolset, other_tool]." - ) - def _load_skill_description(self) -> str: """ Build the `load_skill` tool description, including the catalog of discovered skills. diff --git a/haystack/tools/toolset.py b/haystack/tools/toolset.py index 845a24fa11c..6bb61f705db 100644 --- a/haystack/tools/toolset.py +++ b/haystack/tools/toolset.py @@ -2,7 +2,6 @@ # # SPDX-License-Identifier: Apache-2.0 -import copy from collections.abc import Iterator from dataclasses import dataclass, field from typing import Any @@ -51,58 +50,43 @@ def subtract(a: Annotated[int, "first number"], b: Annotated[int, "second number By subclassing Toolset, you can create implementations that dynamically load tools from external sources like OpenAPI URLs, MCP servers, or other resources. + When implementing a custom Toolset subclass for dynamic tool loading: + - Load the tools in `warm_up()` and assign them to `self.tools`. Since `warm_up()` may be called before + every run, make it idempotent by guarding on your own state (e.g. `if self._client is not None: return`). + - Override `to_dict()` and `from_dict()` to serialize the endpoint descriptor (URL, server info) rather than + the dynamically loaded Tool instances. + Example: ```python - from typing import Annotated from haystack.core.serialization import generate_qualified_class_name - from haystack.tools import tool, Toolset - from haystack.components.agents import Agent - from haystack.components.generators.chat import OpenAIChatGenerator - - class CalculatorToolset(Toolset): - '''A toolset for calculator operations.''' + from haystack.tools import Toolset - def __init__(self) -> None: - super().__init__(self._create_tools()) + class RemoteServiceToolset(Toolset): + def __init__(self, endpoint: str) -> None: + self.endpoint = endpoint + self._client = None + super().__init__(tools=[]) # tools are loaded on warm_up() - def _create_tools(self): - # These tools are defined statically for illustration purposes only. - # In a real-world scenario, you would dynamically load tools from an external source here. - @tool - def add(a: Annotated[int, "first number"], b: Annotated[int, "second number"]) -> int: - '''Add two numbers.''' - return a + b - - @tool - def multiply(a: Annotated[int, "first number"], b: Annotated[int, "second number"]) -> int: - '''Multiply two numbers.''' - return a * b - - return [add, multiply] + def warm_up(self) -> None: + if self._client is not None: + return + self._client = connect(self.endpoint) + self.tools = self._client.fetch_tools() def to_dict(self): return { "type": generate_qualified_class_name(type(self)), - "data": {}, # no data to serialize as we define the tools dynamically + "data": {"endpoint": self.endpoint}, } @classmethod def from_dict(cls, data): - return cls() # Recreate the tools dynamically during deserialization - - # Create the dynamic toolset and use it with an Agent - calculator_toolset = CalculatorToolset() - agent = Agent(chat_generator=OpenAIChatGenerator(), tools=calculator_toolset) + return cls(endpoint=data["data"]["endpoint"]) ``` Toolset implements the collection interface (__iter__, __contains__, __len__, __getitem__), making it behave like a list of Tools. This makes it compatible with components that expect iterable tools, such as Agent or Haystack chat generators. - - When implementing a custom Toolset subclass for dynamic tool loading: - - Perform the dynamic loading in the __init__ method - - Override to_dict() and from_dict() methods if your tools are defined dynamically - - Serialize endpoint descriptors rather than tool instances if your tools are loaded from external sources """ # Use field() with default_factory to initialize the list @@ -110,9 +94,7 @@ def from_dict(cls, data): def __post_init__(self) -> None: """ - Validate and set up the toolset after initialization. - - This handles the case when tools are provided during initialization. + Validate the tools provided during initialization. """ # If initialization was done a single Tool, raise an error if isinstance(self.tools, Tool): @@ -121,59 +103,41 @@ def __post_init__(self) -> None: # Check for duplicate tool names in the initial set _check_duplicate_tool_names(self.tools) - # Tracks whether warm_up() has already run so subsequent calls become a no-op. - self._is_warmed_up = False - - # Optional per-run name filter. When set, iteration only yields tools whose name is in this set. - # None means no filtering. Set on a per-run spawn(), so it never leaks across runs. - self._selected_tool_names: set[str] | None = None - def __iter__(self) -> Iterator[Tool]: """ Return an iterator over the Tools in this Toolset. - This allows the Toolset to be used wherever a list of Tools is expected. If a name filter is active, - only the tools whose names are in it are yielded. + This allows the Toolset to be used wherever a list of Tools is expected. :returns: An iterator yielding Tool instances """ - for tool in self.tools: - if self._selected_tool_names is None or tool.name in self._selected_tool_names: - yield tool + return iter(self.tools) def get_selectable_tools(self) -> list[Tool]: """ - Return the full set of tools that can be selected by name, ignoring any active name filter. + Return the tools available for name-based selection (e.g. via `Agent.run(tools=["tool_name"])`). - This differs from iteration, which yields only the tools currently exposed (and respects the name filter). - Override this when a Toolset's iteration does not surface every selectable tool, so name-based selection - can still target the full set. - - Warms up the Toolset first if needed, so lazily loaded tools (those a Toolset fetches in `warm_up()`) - are available for selection. + 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. :returns: The list of tools available for name-based selection. """ - if not self._is_warmed_up: - self.warm_up() + self.warm_up() return list(self.tools) - def spawn(self) -> "Toolset": + def spawn(self, selected_tool_names: set[str] | None = None) -> "Toolset": # noqa: ARG002 """ - Return an isolated copy of this Toolset for a single run. + Return this Toolset, or an isolated copy of it, for a single run. - The copy shares this Toolset's read-only state (its tools and any warmed-up resources) but gets fresh - run-scoped state, so concurrent runs that share the same configured Toolset don't corrupt each other (for - example, one run's name selection leaking into another). Warms up first if needed so the copy shares the - warmed state. Subclasses with additional run-scoped state should override this. + A plain Toolset has no run-scoped state, so the default implementation returns `self` and ignores the + selection (the Agent materializes it). Subclasses with run-scoped state (e.g. SearchableToolset) override + this to return a copy carrying the selection, so concurrent runs sharing the same configured Toolset + don't corrupt each other. - :returns: A run-scoped copy of this Toolset. + :param selected_tool_names: Optional tool names this run is restricted to. None means no restriction. + :returns: This Toolset, or a run-scoped copy of it. """ - if not self._is_warmed_up: - self.warm_up() - new = copy.copy(self) - new._selected_tool_names = None - return new + return self def __contains__(self, item: str | Tool) -> bool: """ @@ -192,6 +156,23 @@ def __contains__(self, item: str | Tool) -> bool: return any(tool is item or tool == item for tool in self) return False + def __len__(self) -> int: + """ + Return the number of Tools in this Toolset. + + :returns: Number of Tools + """ + return sum(1 for _ in self) + + def __getitem__(self, index: int) -> Tool: + """ + Get a Tool by index. + + :param index: Index of the Tool to get + :returns: The Tool at the specified index + """ + return list(self)[index] + def warm_up(self) -> None: """ Prepare the Toolset for use. @@ -201,65 +182,43 @@ def warm_up(self) -> None: - Setting up shared resources (database connections, HTTP sessions) instead of warming individual tools - - Implementing custom initialization logic for dynamically loaded 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 rather than warming up - individual tools: + might override this to initialize a shared connection and load the tools through it: ```python class MCPToolset(Toolset): def warm_up(self) -> None: - # Only warm up the shared MCP connection, not individual tools + 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 is idempotent: it only warms up the tools the first time it is called. - Subclasses overriding it should preserve this contract (for example by guarding on - `self._is_warmed_up`). + 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()`. """ - if self._is_warmed_up: - return for tool in self.tools: if hasattr(tool, "warm_up"): tool.warm_up() - self._is_warmed_up = True - def add(self, tool: "Tool | Toolset") -> None: + def add(self, tool: Tool) -> None: """ - Add a new Tool or merge another Toolset. - - If this Toolset has already been warmed up, the newly added Tool (or the tools of the - added Toolset) are warmed up immediately so they are ready to use without requiring a - second `warm_up()` call on the whole Toolset. + Add a new Tool to this Toolset. - Note: adding a Toolset flattens it into its individual tools, so this is only recommended - for Toolsets that don't manage shared resources in their `warm_up()` (or `__init__`). - For example, combining with an `MCPToolset`, which owns a shared connection, is not - recommended: the connection's lifecycle would no longer be managed by the original - Toolset. In those cases combine Toolsets with `+` (which preserves each Toolset as a - unit via `_ToolsetWrapper`) instead. - - :param tool: A Tool instance or another Toolset to add + :param tool: A Tool instance to add :raises ValueError: If adding the tool would result in duplicate tool names - :raises TypeError: If the provided object is not a Tool or Toolset + :raises TypeError: If the provided object is not a Tool """ - if not isinstance(tool, (Tool, Toolset)): - raise TypeError(f"Expected Tool or Toolset, got {type(tool).__name__}") - - # Warm up the source before flattening so that lazily-loaded toolsets (e.g. MCPToolset) - # expose their tools, and so newly added tools are ready to use right away. - if self._is_warmed_up and hasattr(tool, "warm_up"): - tool.warm_up() - - new_tools = [tool] if isinstance(tool, Tool) else list(tool) + if not isinstance(tool, Tool): + raise TypeError(f"Expected Tool, got {type(tool).__name__}") # Check for duplicates before adding - combined_tools = self.tools + new_tools - _check_duplicate_tool_names(combined_tools) - - self.tools.extend(new_tools) + _check_duplicate_tool_names(self.tools + [tool]) + self.tools.append(tool) def to_dict(self) -> dict[str, Any]: """ @@ -302,131 +261,3 @@ def from_dict(cls, data: dict[str, Any]) -> "Toolset": tools.append(tool_class.from_dict(tool_data)) return cls(tools=tools) - - def __add__(self, other: "Tool | Toolset | list[Tool]") -> "Toolset": - """ - Concatenate this Toolset with another Tool, Toolset, or list of Tools. - - :param other: Another Tool, Toolset, or list of Tools to concatenate - :returns: A new Toolset containing all tools - :raises TypeError: If the other parameter is not a Tool, Toolset, or list of Tools - :raises ValueError: If the combination would result in duplicate tool names - """ - if isinstance(other, Tool): - return Toolset(tools=self.tools + [other]) - if isinstance(other, Toolset): - return _ToolsetWrapper([self, other]) - if isinstance(other, list) and all(isinstance(item, Tool) for item in other): - return Toolset(tools=self.tools + other) - raise TypeError(f"Cannot add {type(other).__name__} to Toolset") - - def __len__(self) -> int: - """ - Return the number of Tools in this Toolset (respecting any active name filter). - - :returns: Number of Tools - """ - return sum(1 for _ in self) - - def __getitem__(self, index: int) -> Tool: - """ - Get a Tool by index (respecting any active name filter). - - :param index: Index of the Tool to get - :returns: The Tool at the specified index - """ - return list(self)[index] - - -class _ToolsetWrapper(Toolset): - """ - A wrapper that holds multiple toolsets and provides a unified interface. - - This is used internally when combining different types of toolsets to preserve - their individual configurations while still being usable with Agent and Haystack chat generators. - """ - - def __init__(self, toolsets: list[Toolset]) -> None: - super().__init__([tool for toolset in toolsets for tool in toolset]) - self.toolsets = toolsets - # Tracks whether warm_up() has already run so subsequent calls become a no-op. - self._is_warmed_up = False - - def __iter__(self) -> Iterator[Tool]: - """Iterate over all tools from all toolsets, honoring any active name filter.""" - for toolset in self.toolsets: - for tool in toolset: - if self._selected_tool_names is None or tool.name in self._selected_tool_names: - yield tool - - def get_selectable_tools(self) -> list[Tool]: - """Return every selectable tool across all wrapped toolsets, ignoring any active filter.""" - return [tool for toolset in self.toolsets for tool in toolset.get_selectable_tools()] - - def spawn(self) -> "_ToolsetWrapper": - """Return an isolated copy with each wrapped toolset spawned.""" - return _ToolsetWrapper([toolset.spawn() for toolset in self.toolsets]) - - def __contains__(self, item: Any) -> bool: - """Check if a tool is in any of the toolsets.""" - return any(item in toolset for toolset in self.toolsets) - - def warm_up(self) -> None: - """ - Warm up all wrapped toolsets. - - This method is idempotent: it only warms up the wrapped toolsets the first time it is - called. The individual toolsets are themselves expected to have idempotent `warm_up()` - methods. - """ - if self._is_warmed_up: - return - for toolset in self.toolsets: - toolset.warm_up() - self._is_warmed_up = True - - def to_dict(self) -> dict[str, Any]: - """ - Serialize the wrapper to a dictionary. - - Each wrapped toolset is serialized via its own `to_dict()`, so any subclass that - overrides serialization (e.g. a toolset that serializes a connection/endpoint - descriptor) is preserved. - - :returns: A dictionary representation of the wrapper. - """ - return { - "type": generate_qualified_class_name(type(self)), - "data": {"toolsets": [toolset.to_dict() for toolset in self.toolsets]}, - } - - @classmethod - def from_dict(cls, data: dict[str, Any]) -> "_ToolsetWrapper": - """ - Deserialize a wrapper from a dictionary. - - :param data: Dictionary representation of the wrapper. - :returns: A new `_ToolsetWrapper` instance. - :raises TypeError: If any serialized entry is not a subclass of Toolset. - """ - inner_data = data["data"] - toolsets_data = inner_data.get("toolsets", []) - - toolsets = [] - for toolset_data in toolsets_data: - toolset_class = import_class_by_name(toolset_data["type"]) - if not issubclass(toolset_class, Toolset): - raise TypeError(f"Class '{toolset_class}' is not a subclass of Toolset") - toolsets.append(toolset_class.from_dict(toolset_data)) - - return cls(toolsets=toolsets) - - def __add__(self, other: Toolset | Tool | list[Tool]) -> "_ToolsetWrapper": - """Add another toolset or tool to this wrapper.""" - if isinstance(other, Toolset): - return _ToolsetWrapper(self.toolsets + [other]) - if isinstance(other, Tool): - return _ToolsetWrapper(self.toolsets + [Toolset([other])]) - if isinstance(other, list) and all(isinstance(item, Tool) for item in other): - return _ToolsetWrapper(self.toolsets + [Toolset(other)]) - raise TypeError(f"Cannot add {type(other).__name__} to _ToolsetWrapper") diff --git a/releasenotes/notes/simplify-toolset-102a9effbe4f188f.yaml b/releasenotes/notes/simplify-toolset-102a9effbe4f188f.yaml new file mode 100644 index 00000000000..25f520eb569 --- /dev/null +++ b/releasenotes/notes/simplify-toolset-102a9effbe4f188f.yaml @@ -0,0 +1,16 @@ +--- +upgrade: + - | + The ``+`` operator for combining Toolsets has been removed. Pass Toolsets as a list wherever tools are + accepted instead: ``Agent(tools=[toolset_a, toolset_b])``. Pipelines serialized with a ``+``-combined + Toolset cannot be loaded anymore (their YAML references the removed internal ``_ToolsetWrapper`` class): + recreate them with the list form and serialize them again. + - | + ``Toolset.add()`` now only accepts a single Tool, not another Toolset. To combine Toolsets, pass them as + a list: ``Agent(tools=[toolset_a, toolset_b])``. + - | + Haystack can call ``warm_up()`` on Tools and Toolsets more than once, for example before every run. + Previously ``Toolset`` absorbed repeated calls with an internal ``_is_warmed_up`` flag; that flag is gone + and every call now reaches your ``warm_up()``. If your custom Tool or Toolset does expensive work there + (connecting to a server, loading a model), or relied on the ``_is_warmed_up`` attribute, guard with your + own state and return early, for example ``if self._client is not None: return``. diff --git a/test/components/agents/test_agent.py b/test/components/agents/test_agent.py index 01616b3217e..d5d119b5532 100644 --- a/test/components/agents/test_agent.py +++ b/test/components/agents/test_agent.py @@ -1772,6 +1772,34 @@ def test_agent_span_has_parent_when_in_pipeline(self, spying_tracer, weather_too class TestAgentToolSelection: + @staticmethod + def _agent_with_duplicate_tool_names() -> Agent: + def make_tool(description: str) -> Tool: + return Tool( + name="same_name", + description=description, + parameters={"type": "object", "properties": {}}, + function=lambda: None, + ) + + agent = Agent( + chat_generator=MockChatGenerator("Hello"), + tools=[Toolset([make_tool("first")]), Toolset([make_tool("second")])], + ) + agent.warm_up() + return agent + + def test_run_raises_on_duplicate_tool_names_across_toolsets(self): + agent = self._agent_with_duplicate_tool_names() + with pytest.raises(ValueError, match="Duplicate tool names"): + agent.run(messages=[ChatMessage.from_user("hi")]) + + @pytest.mark.asyncio + async def test_run_async_raises_on_duplicate_tool_names_across_toolsets(self): + agent = self._agent_with_duplicate_tool_names() + with pytest.raises(ValueError, match="Duplicate tool names"): + await agent.run_async(messages=[ChatMessage.from_user("hi")]) + def test_tool_selection_new_tool(self, weather_tool: Tool, component_tool: Tool): chat_generator = MockChatGenerator("Hello") agent = Agent(chat_generator=chat_generator, tools=[weather_tool], system_prompt="This is a system prompt.") @@ -2202,13 +2230,14 @@ def test_warm_up_toolset(self): agent.warm_up() assert toolset.was_warmed_up - def test_warm_up_mixed_toolsets(self): + def test_warm_up_list_of_toolsets(self): + # Toolsets are combined by passing them as a list; each keeps its own warm_up. tool1 = self._make_tracking_tool("tool1") toolset1 = self._make_tracking_toolset([tool1]) tool2 = self._make_tracking_tool("tool2") toolset2 = self._make_tracking_toolset([tool2]) - agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=toolset1 + toolset2) + agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=[toolset1, toolset2]) assert not toolset1.was_warmed_up assert not toolset2.was_warmed_up @@ -2236,7 +2265,7 @@ def test_warm_up_mixed_list_of_tools_and_toolsets(self): assert toolset1.was_warmed_up assert toolset2.was_warmed_up - def test_warm_up_is_idempotent(self): + def test_warm_up_rewarms_tools_on_every_call(self): call_count = {"n": 0} tool = Tool( name="counting_tool", @@ -2257,16 +2286,28 @@ def counting_warm_up(): agent.warm_up() agent.warm_up() - assert call_count["n"] == 1 + assert call_count["n"] == 3 - def test_warm_up_refreshes_toolset(self): - """Agent.warm_up() must warm up lazy toolsets (e.g. MCPToolset) so the actual tools are available at runtime.""" - placeholder_tool = Tool( - name="mcp_not_connected_placeholder_123", - description="Placeholder tool before connection", - parameters={"type": "object", "properties": {}}, - function=lambda: "placeholder", - ) + @pytest.mark.parametrize( + "initial_tools", + [ + pytest.param([], id="empty"), + pytest.param( + [ + Tool( + name="mcp_not_connected_placeholder_123", + description="Placeholder tool before connection", + parameters={"type": "object", "properties": {}}, + function=lambda: "placeholder", + ) + ], + id="placeholder", + ), + ], + ) + def test_warm_up_loads_lazy_toolset(self, initial_tools): + # A lazy toolset (e.g. MCPToolset) loads its real tools on warm_up(). Before that it may be empty + # (a truthiness check would skip it) or expose a placeholder: warm_up must load the real tools either way. actual_tool = Tool( name="get_time", description="Get the current time in ISO format", @@ -2274,21 +2315,21 @@ def test_warm_up_refreshes_toolset(self): function=lambda: "2024-12-01T12:00:00Z", ) - class MockMCPToolset(Toolset): + class LazyToolset(Toolset): def __init__(self): - super().__init__([placeholder_tool]) self._connected = False + super().__init__(list(initial_tools)) def warm_up(self): if not self._connected: self.tools = [actual_tool] self._connected = True - mcp_toolset = MockMCPToolset() - agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=mcp_toolset) - assert mcp_toolset.tools == [placeholder_tool] + toolset = LazyToolset() + agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=toolset) + assert toolset.tools == initial_tools agent.warm_up() - assert mcp_toolset.tools == [actual_tool] + assert toolset.tools == [actual_tool] def test_run_warms_lazy_toolset_before_tool_selection(self): """ @@ -2400,7 +2441,6 @@ def test_warm_up_delegates_to_chat_generator(self, weather_tool): chat_generator.warm_up.reset_mock() agent.run([ChatMessage.from_user("What is the weather in Berlin?")]) - assert agent._tools_warmed_up is True # warm_up runs twice here: the Agent delegates to the generator, and the generator's own run() self-warms assert chat_generator.warm_up.call_count == 2 diff --git a/test/components/agents/test_utils.py b/test/components/agents/test_utils.py index 4bbf950014e..cde495b64c3 100644 --- a/test/components/agents/test_utils.py +++ b/test/components/agents/test_utils.py @@ -95,27 +95,83 @@ def test_raises_for_invalid_name(self, first_tool: Tool): with pytest.raises(ValueError, match="The following tool names are not valid"): _select_tools_by_name([first_tool], ["unknown"]) - def test_raises_when_no_tools_configured(self, first_tool: Tool): + @pytest.mark.parametrize("configured_tools", [[], Toolset([])], ids=["empty_list", "empty_toolset"]) + def test_raises_when_no_tools_configured(self, configured_tools, first_tool: Tool): with pytest.raises(ValueError, match="No tools were configured for the Agent at initialization."): - _select_tools_by_name([], [first_tool.name]) + _select_tools_by_name(configured_tools, [first_tool.name]) - def test_spawns_toolsets_without_mutating_them(self, first_tool: Tool, second_tool: Tool): + def test_reduces_plain_toolsets_to_matching_tools(self, first_tool: Tool, second_tool: Tool): toolset = Toolset([first_tool, second_tool]) selected = _select_tools_by_name([toolset], [first_tool.name]) - spawned = selected[0] - assert isinstance(spawned, Toolset) - assert spawned is not toolset - assert spawned._selected_tool_names == {first_tool.name} - assert toolset._selected_tool_names is None + assert selected == [first_tool] + # The configured toolset is untouched. + assert list(toolset) == [first_tool, second_tool] def test_selects_standalone_tools_and_toolsets(self, first_tool: Tool, second_tool: Tool): toolset = Toolset([first_tool]) selected = _select_tools_by_name([second_tool, toolset], [first_tool.name, second_tool.name]) - assert second_tool in selected - spawned = next(item for item in selected if isinstance(item, Toolset)) - assert spawned is not toolset - assert spawned._selected_tool_names == {first_tool.name} - assert toolset._selected_tool_names is None + assert selected == [second_tool, first_tool] + + def test_warms_up_lazy_toolsets_to_resolve_names(self, first_tool: Tool, second_tool: Tool): + 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): + class RunScopedToolset(Toolset): + """A Toolset overriding spawn(), signaling run-scoped state.""" + + def __init__(self, tools): + super().__init__(tools) + self.selected: set[str] | None = None + + def spawn(self, selected_tool_names: set[str] | None = None) -> "RunScopedToolset": + new = RunScopedToolset(list(self.tools)) + new.selected = set(selected_tool_names) if selected_tool_names is not None else None + return new + + toolset = RunScopedToolset([first_tool, second_tool]) + selected = _select_tools_by_name([toolset], [first_tool.name]) + + run_copy = selected[0] + assert isinstance(run_copy, RunScopedToolset) + assert run_copy is not toolset + assert run_copy.selected == {first_tool.name} + # The configured toolset is untouched. + assert toolset.selected is None + + def test_selects_tools_not_surfaced_by_iteration(self, first_tool: Tool, second_tool: Tool): + class DiscoveryToolset(Toolset): + """A dynamic Toolset without a spawn() override: iteration yields less than get_selectable_tools().""" + + def __init__(self, catalog): + self._catalog = catalog + super().__init__([]) + + def __iter__(self): + yield from [] # nothing discovered yet + + def get_selectable_tools(self) -> list[Tool]: + return list(self._catalog) + + toolset = DiscoveryToolset([first_tool, second_tool]) + # The requested tool must be selected even though iteration does not surface it. + selected = _select_tools_by_name([toolset], [first_tool.name]) + assert selected == [first_tool] class TestContextTokensFromUsage: diff --git a/test/tools/skills/test_skill_toolset.py b/test/tools/skills/test_skill_toolset.py index e8ec3e1dec8..efd3eab876b 100644 --- a/test/tools/skills/test_skill_toolset.py +++ b/test/tools/skills/test_skill_toolset.py @@ -113,13 +113,6 @@ def test_add_is_not_supported(self, tmp_path): with pytest.raises(NotImplementedError, match="does not support adding tools"): toolset.add(extra) - def test_concat_is_not_supported(self, tmp_path): - _write_skill(tmp_path, "pdf-forms", description="Use to fill PDF forms.") - toolset = SkillToolset(FileSystemSkillStore(tmp_path)) - extra = Tool(name="extra", description="d", parameters={"type": "object", "properties": {}}, function=len) - with pytest.raises(NotImplementedError, match="does not support concatenation"): - _ = toolset + extra - def test_accepts_skill_store_instance(self, tmp_path): _write_skill(tmp_path, "pdf-forms", description="Use to fill PDF forms.") store = FileSystemSkillStore(tmp_path) diff --git a/test/tools/test_searchable_toolset.py b/test/tools/test_searchable_toolset.py index c83854c7a15..fd9a9a2ead9 100644 --- a/test/tools/test_searchable_toolset.py +++ b/test/tools/test_searchable_toolset.py @@ -129,10 +129,6 @@ def test_init_with_invalid_catalog(self): def test_not_implemented_methods(self): toolset = SearchableToolset(catalog=[]) - with pytest.raises(NotImplementedError): - toolset + Tool( - name="test", description="test", parameters={"type": "object", "properties": {}}, function=lambda: None - ) with pytest.raises(NotImplementedError): toolset.add( Tool( @@ -344,10 +340,10 @@ def test_iter_with_discovered_tools(self, large_catalog): def test_iter_automatically_warms_up(self, large_catalog): toolset = SearchableToolset(catalog=large_catalog) - assert not toolset._is_warmed_up + assert toolset._passthrough is None # not warmed up yet list(toolset) - assert toolset._is_warmed_up + assert toolset._passthrough is not None def test_contains_bootstrap_tool(self, large_catalog): """Test __contains__ for bootstrap tool.""" @@ -528,11 +524,11 @@ 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") toolset = SearchableToolset(catalog=large_catalog) - assert toolset._is_warmed_up is False + assert toolset._passthrough is None # not warmed up yet Agent(chat_generator=OpenAIChatGenerator(), tools=toolset) - assert toolset._is_warmed_up is False + assert toolset._passthrough is None # still not warmed up def test_warm_up_idempotent(self, large_catalog): """Test that warm_up can be called multiple times safely.""" @@ -858,7 +854,7 @@ def test_spawns_have_independent_discovered_tools_and_selection(self, large_cata toolset = SearchableToolset(catalog=large_catalog, search_threshold=3) toolset.warm_up() - spawn_a = toolset.spawn() + spawn_a = toolset.spawn(selected_tool_names={"get_weather"}) spawn_b = toolset.spawn() assert spawn_a is not spawn_b @@ -867,7 +863,6 @@ def test_spawns_have_independent_discovered_tools_and_selection(self, large_cata assert spawn_a._bootstrap_tool is not None assert spawn_a._bootstrap_tool is not spawn_b._bootstrap_tool - spawn_a._selected_tool_names = {"get_weather"} spawn_a._bootstrap_tool.invoke(tool_keywords="weather add stock multiply") # Discovery on spawn_a does not leak into spawn_b or the configured toolset. diff --git a/test/tools/test_tools_utils.py b/test/tools/test_tools_utils.py index d0d31fe870e..eddc3677e52 100644 --- a/test/tools/test_tools_utils.py +++ b/test/tools/test_tools_utils.py @@ -371,11 +371,13 @@ class WarmupCountingToolset(Toolset): def __init__(self, tools): super().__init__(tools) self.warm_up_count = 0 + self._loaded = False def warm_up(self): - if self._is_warmed_up: + if self._loaded: return self.warm_up_count += 1 + self._loaded = True super().warm_up() # Also warm up individual tools tool = WarmupCountingTool( diff --git a/test/tools/test_toolset.py b/test/tools/test_toolset.py index 139e13d2a5d..937752ba7f1 100644 --- a/test/tools/test_toolset.py +++ b/test/tools/test_toolset.py @@ -121,20 +121,6 @@ def warm_up(self) -> None: self.warm_up_count += 1 -class WarmUpCountingToolset(Toolset): - """A Toolset that records how many times its own warm_up() did real work.""" - - def __init__(self, tools): - super().__init__(tools) - self.warm_up_count = 0 - - def warm_up(self) -> None: - if self._is_warmed_up: - return - self.warm_up_count += 1 - super().warm_up() - - class TestToolset: def test_toolset_with_multiple_tools(self, add_tool, multiply_tool): """Test that a Toolset with multiple tools works properly.""" @@ -181,42 +167,20 @@ def test_toolset_contains(self, add_tool, multiply_tool): assert "multiply" not in toolset assert "non_existent_tool" not in toolset - def test_toolset_addition(self, add_tool, multiply_tool, subtract_tool): - """Test that the __add__ method combines toolsets with various operand types.""" - base = Toolset([add_tool]) - - # Toolset + Tool - result = base + multiply_tool - assert isinstance(result, Toolset) - assert [t.name for t in result] == ["add", "multiply"] - - # Toolset + Toolset - result = base + Toolset([subtract_tool]) - assert isinstance(result, Toolset) - assert [t.name for t in result] == ["add", "subtract"] - - # Toolset + list[Tool] - result = base + [multiply_tool, subtract_tool] - assert isinstance(result, Toolset) - assert [t.name for t in result] == ["add", "multiply", "subtract"] - - # Unsupported operand types raise TypeError - with pytest.raises(TypeError): - base + "not_a_tool" # type: ignore[operator] - with pytest.raises(TypeError): - base + 123 # type: ignore[operator] + def test_combining_toolsets_via_unpacking(self, add_tool, multiply_tool, subtract_tool): + combined = Toolset([*Toolset([add_tool, subtract_tool]), multiply_tool]) + assert [t.name for t in combined] == ["add", "subtract", "multiply"] # The combined tools remain invocable message = ChatMessage.from_assistant( tool_calls=[ ToolCall(tool_name="add", arguments={"a": 10, "b": 5}), ToolCall(tool_name="multiply", arguments={"a": 10, "b": 5}), - ToolCall(tool_name="subtract", arguments={"a": 10, "b": 5}), ] ) - tool_messages = _run_tool_messages(messages=[message], tools=result) + tool_messages = _run_tool_messages(messages=[message], tools=combined) tool_results = [tcr.result for message in tool_messages for tcr in message.tool_call_results] - assert tool_results == ["15", "50", "5"] + assert tool_results == ["15", "50"] def test_toolset_serialization(self, add_tool): """Test that a Toolset can be serialized and deserialized.""" @@ -244,9 +208,8 @@ def test_toolset_duplicate_tool_names(self, add_tool): with pytest.raises(ValueError, match="Duplicate tool names found"): toolset.add(add_tool) - toolset2 = Toolset([add_tool]) with pytest.raises(ValueError, match="Duplicate tool names found"): - _ = toolset + toolset2 + Toolset([*toolset, *Toolset([add_tool])]) class TestToolsetWithAgent: @@ -312,6 +275,25 @@ def test_agent_serde_with_list_of_toolsets(self, weather_tool, add_tool, monkeyp assert len(deserialized_agent.tools) == 2 assert all(isinstance(ts, Toolset) for ts in deserialized_agent.tools) + def test_agent_serde_with_list_containing_custom_toolset(self, multiply_tool, monkeypatch): + """A custom-serde Toolset inside a list roundtrips through its own to_dict/from_dict.""" + monkeypatch.setenv("OPENAI_API_KEY", "test") + agent = Agent(chat_generator=OpenAIChatGenerator(), tools=[DynamicToolset(), Toolset([multiply_tool])]) + data = agent.to_dict() + + tools_data = data["init_parameters"]["tools"] + assert tools_data[0]["type"] == "test_toolset.DynamicToolset" + assert tools_data[0]["data"] == {} + assert tools_data[1]["type"] == "haystack.tools.toolset.Toolset" + + deserialized_agent = Agent.from_dict(data) + custom_toolset, plain_toolset = deserialized_agent.tools[0], deserialized_agent.tools[1] + assert isinstance(custom_toolset, DynamicToolset) + assert isinstance(plain_toolset, Toolset) + # The custom toolset rebuilt its tools via its own from_dict. + assert [tool.name for tool in custom_toolset] == ["add"] + assert [tool.name for tool in plain_toolset] == ["multiply"] + def test_list_of_toolsets_runtime_override(self, weather_tool, add_tool, multiply_tool): """Test that list of Toolsets can be passed as runtime override to Agent.run().""" toolset2 = Toolset([add_tool]) @@ -364,10 +346,6 @@ def test_pipeline_with_list_of_toolsets(self, add_tool, multiply_tool, monkeypat class TestToolsetWarmUp: """Stress tests for Toolset warm_up behavior.""" - def test_new_toolset_is_not_warmed_up(self): - toolset = Toolset([WarmUpCountingTool("a")]) - assert toolset._is_warmed_up is False - def test_warm_up_warms_all_tools(self): t1, t2 = WarmUpCountingTool("a"), WarmUpCountingTool("b") toolset = Toolset([t1, t2]) @@ -376,170 +354,66 @@ def test_warm_up_warms_all_tools(self): toolset.warm_up() assert t1.warm_up_count == 1 assert t2.warm_up_count == 1 - assert toolset._is_warmed_up is True - def test_warm_up_is_idempotent(self): + 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 == 1 + assert t1.warm_up_count == 3 - def test_add_before_warm_up_does_not_warm_tools(self): + 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) - # Nothing is warmed until warm_up() is called explicitly. - assert existing.warm_up_count == 0 assert new_tool.warm_up_count == 0 toolset.warm_up() - assert existing.warm_up_count == 1 - assert new_tool.warm_up_count == 1 - - def test_add_tool_after_warm_up_warms_only_new_tool(self): - existing = WarmUpCountingTool("a") - toolset = Toolset([existing]) - toolset.warm_up() - assert existing.warm_up_count == 1 - new_tool = WarmUpCountingTool("b") - toolset.add(new_tool) - # The new tool is warmed immediately, the already-warmed tool is not re-warmed. assert new_tool.warm_up_count == 1 - assert existing.warm_up_count == 1 - def test_add_toolset_after_warm_up_warms_added_toolset(self): + def test_add_toolset_raises(self): toolset = Toolset([WarmUpCountingTool("a")]) - toolset.warm_up() - added_tools = [WarmUpCountingTool("b"), WarmUpCountingTool("c")] - added = WarmUpCountingToolset(added_tools) - toolset.add(added) - # The added toolset's own warm_up() is invoked, warming its tools. - assert added.warm_up_count == 1 - assert all(tool.warm_up_count == 1 for tool in added_tools) - - def test_plus_returns_new_unwarmed_toolset(self): - ts1 = Toolset([WarmUpCountingTool("a")]) - ts1.warm_up() - assert ts1._is_warmed_up is True - new_tool = WarmUpCountingTool("b") - ts2 = ts1 + new_tool - # `+` returns a brand new Toolset object that has not been warmed up yet. - assert ts2 is not ts1 - assert ts2._is_warmed_up is False - assert new_tool.warm_up_count == 0 - ts2.warm_up() - assert new_tool.warm_up_count == 1 + not_a_tool: Any = Toolset([WarmUpCountingTool("b")]) + with pytest.raises(TypeError, match="Expected Tool"): + toolset.add(not_a_tool) -class TestToolsetToolSelection: - """Tests for get_selectable_tools(), the name filter, and spawn().""" +class TestToolsetSpawn: + """Tests for spawn(), the run-scoping hook.""" - def test_no_filter_yields_all_tools(self, add_tool, multiply_tool): + def test_spawn_returns_self_for_plain_toolset(self, add_tool, multiply_tool): toolset = Toolset([add_tool, multiply_tool]) - assert toolset._selected_tool_names is None + assert toolset.spawn() is toolset + assert toolset.spawn(selected_tool_names={"add"}) is toolset assert [tool.name for tool in toolset] == ["add", "multiply"] - assert len(toolset) == 2 + + +class TestToolsetToolSelection: + """Tests for get_selectable_tools().""" def test_get_selectable_tools_returns_all_tools(self, add_tool, multiply_tool): toolset = Toolset([add_tool, multiply_tool]) assert toolset.get_selectable_tools() == [add_tool, multiply_tool] - def test_get_selectable_tools_ignores_active_filter(self, add_tool, multiply_tool): - toolset = Toolset([add_tool, multiply_tool]) - toolset._selected_tool_names = {"add"} - # Iteration is filtered, but get_selectable_tools still returns the full set. - assert [tool.name for tool in toolset] == ["add"] - assert {tool.name for tool in toolset.get_selectable_tools()} == {"add", "multiply"} - 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._is_warmed_up: + if self._loaded: return + self._loaded = True self.tools = [add_tool, multiply_tool] - self._is_warmed_up = True toolset = LazyToolset() - assert toolset._is_warmed_up is False assert toolset.tools == [] # not loaded yet selectable = toolset.get_selectable_tools() - assert toolset._is_warmed_up is True # get_selectable_tools triggered warm_up assert [tool.name for tool in selectable] == ["add", "multiply"] - - def test_filter_restricts_iteration(self, add_tool, multiply_tool, subtract_tool): - toolset = Toolset([add_tool, multiply_tool, subtract_tool]) - toolset._selected_tool_names = {"add", "subtract"} - assert [tool.name for tool in toolset] == ["add", "subtract"] - - def test_filter_restricts_len(self, add_tool, multiply_tool, subtract_tool): - toolset = Toolset([add_tool, multiply_tool, subtract_tool]) - toolset._selected_tool_names = {"add"} - assert len(toolset) == 1 - - def test_filter_restricts_getitem(self, add_tool, multiply_tool, subtract_tool): - toolset = Toolset([add_tool, multiply_tool, subtract_tool]) - toolset._selected_tool_names = {"subtract"} - assert toolset[0].name == "subtract" - - def test_filter_restricts_contains(self, add_tool, multiply_tool): - toolset = Toolset([add_tool, multiply_tool]) - toolset._selected_tool_names = {"add"} - assert "add" in toolset - assert "multiply" not in toolset - assert add_tool in toolset - assert multiply_tool not in toolset - - def test_spawn_returns_isolated_copy(self, add_tool, multiply_tool): - toolset = Toolset([add_tool, multiply_tool]) - - spawned = toolset.spawn() - - assert spawned is not toolset - assert spawned._selected_tool_names is None - # The copy shares the same (read-only) tools. - assert list(spawned.tools) == list(toolset.tools) - - def test_spawn_selection_does_not_leak_to_original(self, add_tool, multiply_tool): - """A per-run selection set on a spawn must not affect the configured toolset or other spawns.""" - toolset = Toolset([add_tool, multiply_tool]) - - spawn_a = toolset.spawn() - spawn_b = toolset.spawn() - spawn_a._selected_tool_names = {"add"} - - # Each run sees only its own selection; the configured toolset stays unfiltered. - assert [tool.name for tool in spawn_a] == ["add"] - assert [tool.name for tool in spawn_b] == ["add", "multiply"] - assert [tool.name for tool in toolset] == ["add", "multiply"] - assert toolset._selected_tool_names is None - - def test_spawn_warms_up_lazy_toolset(self, add_tool, multiply_tool): - """spawn() warms up a lazy toolset so the copy shares the warmed state.""" - - class LazyToolset(Toolset): - def __init__(self): - super().__init__([]) - - def warm_up(self): - if self._is_warmed_up: - return - self.tools = [add_tool, multiply_tool] - self._is_warmed_up = True - - toolset = LazyToolset() - assert toolset._is_warmed_up is False - - spawned = toolset.spawn() - - assert toolset._is_warmed_up is True # spawn triggered warm_up - assert spawned._is_warmed_up is True - assert [tool.name for tool in spawned] == ["add", "multiply"] diff --git a/test/tools/test_toolset_wrapper.py b/test/tools/test_toolset_wrapper.py deleted file mode 100644 index cb4ac58b8b5..00000000000 --- a/test/tools/test_toolset_wrapper.py +++ /dev/null @@ -1,253 +0,0 @@ -# SPDX-FileCopyrightText: 2022-present deepset GmbH -# -# SPDX-License-Identifier: Apache-2.0 - -from typing import Annotated - -import pytest - -from haystack.components.agents import Agent -from haystack.components.generators.chat import OpenAIChatGenerator -from haystack.core.serialization import generate_qualified_class_name -from haystack.tools import Tool, Toolset, tool -from haystack.tools.toolset import _ToolsetWrapper - - -@tool -def add(a: Annotated[int, "first number"], b: Annotated[int, "second number"]) -> int: - """Add two numbers.""" - return a + b - - -@tool -def multiply(a: Annotated[int, "first number"], b: Annotated[int, "second number"]) -> int: - """Multiply two numbers.""" - return a * b - - -@tool -def subtract(a: Annotated[int, "first number"], b: Annotated[int, "second number"]) -> int: - """Subtract b from a.""" - return a - b - - -@tool -def rebuilt() -> str: - """A rebuilt tool.""" - return "rebuilt" - - -@pytest.fixture -def add_tool(): - return add - - -@pytest.fixture -def multiply_tool(): - return multiply - - -@pytest.fixture -def subtract_tool(): - return subtract - - -class WarmUpCountingTool(Tool): - """A Tool that records how many times warm_up() was called.""" - - def __init__(self, name: str): - super().__init__( - name=name, - description=f"{name} tool", - parameters={"type": "object", "properties": {}}, - function=lambda: None, - ) - self.warm_up_count = 0 - - def warm_up(self) -> None: - self.warm_up_count += 1 - - -class WarmUpCountingToolset(Toolset): - """A Toolset that records how many times its own warm_up() did real work.""" - - def __init__(self, tools): - super().__init__(tools) - self.warm_up_count = 0 - - def warm_up(self) -> None: - if self._is_warmed_up: - return - self.warm_up_count += 1 - super().warm_up() - - -class RebuildingToolset(Toolset): - """A toolset that rebuilds its tools on from_dict() instead of serializing them (like a dynamic toolset).""" - - def __init__(self): - super().__init__([rebuilt]) - - def to_dict(self): - return {"type": generate_qualified_class_name(type(self)), "data": {}} - - @classmethod - def from_dict(cls, data): - return cls() - - -class TestToolsetWrapper: - """Tests for the _ToolsetWrapper class""" - - def test_toolset_plus_toolset_creates_wrapper(self, add_tool, multiply_tool): - """Test that combining two Toolsets creates a _ToolsetWrapper and works correctly.""" - result = Toolset([add_tool]) + Toolset([multiply_tool]) - assert isinstance(result, _ToolsetWrapper) - assert len(result) == 2 - assert add_tool in result - assert multiply_tool in result - - def test_wrapper_getitem_supports_negative_indices(self, add_tool, multiply_tool): - """Test that _ToolsetWrapper.__getitem__ behaves like a list, including negative indices.""" - toolset1 = Toolset([add_tool]) - toolset2 = Toolset([multiply_tool]) - - result = toolset1 + toolset2 - - assert result[0] is add_tool - assert result[1] is multiply_tool - assert result[-1] is multiply_tool - assert result[-2] is add_tool - - with pytest.raises(IndexError): - _ = result[2] - - def test_wrapper_with_agent(self, add_tool, multiply_tool, monkeypatch): - """Test that _ToolsetWrapper works with Agent.""" - monkeypatch.setenv("OPENAI_API_KEY", "test") - wrapper = Toolset([add_tool]) + Toolset([multiply_tool]) - agent = Agent(chat_generator=OpenAIChatGenerator(), tools=wrapper) - agent.warm_up() - assert len(list(agent.tools)) == 2 - - def test_wrapper_chaining_and_duplicate_detection(self, add_tool, multiply_tool, subtract_tool): - """Test chaining operations and that duplicates are still detected.""" - # Chaining should work - result = Toolset([add_tool]) + Toolset([multiply_tool]) + Toolset([subtract_tool]) - assert len(result) == 3 - # Duplicates should be detected - toolset_with_dup = Toolset([add_tool]) - with pytest.raises(ValueError, match="Duplicate tool names found"): - _ = result + toolset_with_dup - - -class TestToolsetWrapperWarmUp: - """Tests for warm_up behavior of _ToolsetWrapper.""" - - def test_new_wrapper_is_not_warmed_up(self): - wrapper = Toolset([WarmUpCountingTool("a")]) + Toolset([WarmUpCountingTool("b")]) - assert wrapper._is_warmed_up is False - - def test_warm_up_delegates_to_each_toolset(self): - ts1 = WarmUpCountingToolset([WarmUpCountingTool("a")]) - ts2 = WarmUpCountingToolset([WarmUpCountingTool("b")]) - wrapper = ts1 + ts2 - wrapper.warm_up() - assert ts1.warm_up_count == 1 - assert ts2.warm_up_count == 1 - assert wrapper._is_warmed_up is True - - def test_warm_up_is_idempotent(self): - ts1 = WarmUpCountingToolset([WarmUpCountingTool("a")]) - ts2 = WarmUpCountingToolset([WarmUpCountingTool("b")]) - wrapper = ts1 + ts2 - wrapper.warm_up() - wrapper.warm_up() - wrapper.warm_up() - assert ts1.warm_up_count == 1 - assert ts2.warm_up_count == 1 - - -class TestToolsetWrapperSerialization: - """Tests for to_dict/from_dict of _ToolsetWrapper.""" - - def test_to_dict(self, add_tool, multiply_tool): - wrapper = Toolset([add_tool]) + Toolset([multiply_tool]) - data = wrapper.to_dict() - assert data["type"] == "haystack.tools.toolset._ToolsetWrapper" - assert len(data["data"]["toolsets"]) == 2 - assert all(ts["type"] == "haystack.tools.toolset.Toolset" for ts in data["data"]["toolsets"]) - - def test_from_dict_round_trip(self, add_tool, multiply_tool): - wrapper = Toolset([add_tool]) + Toolset([multiply_tool]) - restored = _ToolsetWrapper.from_dict(wrapper.to_dict()) - assert isinstance(restored, _ToolsetWrapper) - assert len(restored) == 2 - assert len(restored.toolsets) == 2 - assert "add" in restored - assert "multiply" in restored - - def test_to_dict_preserves_subclass_serialization(self, add_tool): - # RebuildingToolset has a custom to_dict that serializes no tools (they are rebuilt on from_dict). - wrapper = RebuildingToolset() + Toolset([add_tool]) - data = wrapper.to_dict() - - # Each wrapped toolset is serialized via its own to_dict, so the custom one is preserved. - assert data["data"]["toolsets"][0]["type"].endswith("RebuildingToolset") - assert data["data"]["toolsets"][0]["data"] == {} - - restored = _ToolsetWrapper.from_dict(data) - assert isinstance(restored.toolsets[0], RebuildingToolset) - assert "rebuilt" in restored - assert "add" in restored - - def test_from_dict_rejects_non_toolset(self, add_tool): - data = Toolset([add_tool]).to_dict() - data["data"] = {"toolsets": [{"type": "haystack.tools.tool.Tool", "data": {}}]} - - with pytest.raises(TypeError, match="is not a subclass of Toolset"): - _ToolsetWrapper.from_dict(data) - - -class TestToolsetWrapperToolSelection: - """Tests for get_selectable_tools(), the name filter, and spawn() on _ToolsetWrapper.""" - - def test_get_selectable_tools_aggregates_all_toolsets(self, add_tool, multiply_tool, subtract_tool): - wrapper = Toolset([add_tool]) + Toolset([multiply_tool, subtract_tool]) - assert {tool.name for tool in wrapper.get_selectable_tools()} == {"add", "multiply", "subtract"} - - def test_get_selectable_tools_ignores_active_filter(self, add_tool, multiply_tool): - wrapper = Toolset([add_tool]) + Toolset([multiply_tool]) - wrapper._selected_tool_names = {"add"} - # Iteration is filtered, but get_selectable_tools still returns the full set. - assert [tool.name for tool in wrapper] == ["add"] - assert {tool.name for tool in wrapper.get_selectable_tools()} == {"add", "multiply"} - - def test_filter_restricts_iteration_and_len(self, add_tool, multiply_tool, subtract_tool): - wrapper = Toolset([add_tool, multiply_tool]) + Toolset([subtract_tool]) - wrapper._selected_tool_names = {"add", "subtract"} - assert [tool.name for tool in wrapper] == ["add", "subtract"] - assert len(wrapper) == 2 - - def test_filter_restricts_getitem(self, add_tool, multiply_tool, subtract_tool): - wrapper = Toolset([add_tool, multiply_tool]) + Toolset([subtract_tool]) - wrapper._selected_tool_names = {"add", "subtract"} - assert wrapper[0].name == "add" - assert wrapper[-1].name == "subtract" - - def test_spawn_isolates_own_and_child_state(self, add_tool, multiply_tool): - ts1 = Toolset([add_tool]) - ts2 = Toolset([multiply_tool]) - wrapper = ts1 + ts2 - - spawned = wrapper.spawn() - - # The spawn and its wrapped toolsets are independent copies. - assert spawned is not wrapper - spawned._selected_tool_names = {"add"} - assert {tool.name for tool in spawned} == {"add"} - # The configured wrapper and its children are untouched. - assert wrapper._selected_tool_names is None - assert ts1._selected_tool_names is None - assert ts2._selected_tool_names is None - assert {tool.name for tool in wrapper} == {"add", "multiply"}