Coverage for haystack/tools/searchable_toolset.py: 100%
112 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import copy
6from collections.abc import Iterator
7from typing import TYPE_CHECKING, Annotated, Any
9from haystack.core.serialization import generate_qualified_class_name
10from haystack.dataclasses import Document
11from haystack.document_stores.in_memory import InMemoryDocumentStore
12from haystack.document_stores.types import DuplicatePolicy
13from haystack.tools.from_function import create_tool_from_function
14from haystack.tools.serde_utils import deserialize_tools_or_toolset_inplace, serialize_tools_or_toolset
15from haystack.tools.tool import Tool, _check_duplicate_tool_names
16from haystack.tools.toolset import Toolset
17from haystack.tools.utils import flatten_tools_or_toolsets, warm_up_tools
19if TYPE_CHECKING:
20 from haystack.tools import ToolsType
23class SearchableToolset(Toolset):
24 """
25 Dynamic tool discovery from large catalogs using BM25 search.
27 This Toolset enables LLMs to discover and use tools from large catalogs through BM25-based search.
28 Instead of exposing all tools at once (which can overwhelm the LLM context), it provides a `search_tools` bootstrap
29 tool that allows the LLM to find and load specific tools as needed.
31 For very small catalogs (below `search_threshold`), acts as a simple passthrough exposing all tools directly
32 without any discovery mechanism.
34 ### Usage Example
36 ```python
37 from typing import Annotated
39 from haystack.components.agents import Agent
40 from haystack.components.generators.chat import OpenAIChatGenerator
41 from haystack.dataclasses import ChatMessage
42 from haystack.tools import SearchableToolset, tool
44 @tool
45 def get_weather(city: Annotated[str, "The city to get the weather for"]) -> str:
46 '''Get the current weather for a city.'''
47 return f"The weather in {city} is 22°C and sunny."
49 @tool
50 def search_web(query: Annotated[str, "The query to search the web for"]) -> str:
51 '''Search the web for a query.'''
52 return f"Top result for '{query}': ..."
54 @tool
55 def convert_currency(
56 amount: Annotated[float, "The amount to convert"],
57 to_currency: Annotated[str, "The currency to convert to, e.g. 'EUR'"],
58 ) -> str:
59 '''Convert an amount in USD to another currency.'''
60 return f"{amount} USD is {amount * 0.9} {to_currency}"
62 # search_threshold=2 means a catalog of 2+ tools activates discovery: the agent only sees the
63 # `search_tools` tool and must search to load the others (set it higher for larger catalogs).
64 toolset = SearchableToolset(catalog=[get_weather, search_web, convert_currency], search_threshold=2)
66 agent = Agent(chat_generator=OpenAIChatGenerator(), tools=toolset)
68 # The agent is initially provided only with the search_tools tool and will use it to find relevant tools.
69 result = agent.run(messages=[ChatMessage.from_user("What's the weather in Milan?")])
70 print(result["last_message"].text)
71 ```
72 """
74 _VALID_SEARCH_TOOL_PARAMS = {"tool_keywords", "k"}
76 def __init__(
77 self,
78 catalog: "ToolsType",
79 *,
80 top_k: int = 3,
81 search_threshold: int = 8,
82 search_tool_name: str = "search_tools",
83 search_tool_description: str | None = None,
84 search_tool_parameters_description: dict[str, str] | None = None,
85 ) -> None:
86 """
87 Initialize the SearchableToolset.
89 :param catalog: Source of tools - a list of Tools, list of Toolsets, or a single Toolset.
90 :param top_k: Default number of results for search_tools.
91 :param search_threshold: Minimum catalog size to activate search. If catalog has fewer tools, acts as
92 passthrough (all tools visible). Default is 8.
93 :param search_tool_name: Custom name for the bootstrap search tool. Default is "search_tools".
94 :param search_tool_description: Custom description for the bootstrap search tool. If not provided, uses a
95 default description.
96 :param search_tool_parameters_description: Custom descriptions for the bootstrap search tool's parameters.
97 Keys must be a subset of `{"tool_keywords", "k"}`.
98 Example: `{"tool_keywords": "Keywords to find tools, e.g. 'email send'"}`
99 """
100 valid_catalog = isinstance(catalog, Toolset) or (
101 isinstance(catalog, list) and all(isinstance(item, (Tool, Toolset)) for item in catalog)
102 )
103 if not valid_catalog:
104 raise TypeError(
105 f"Invalid catalog type: {type(catalog)}. Expected Tool, Toolset, or list of Tools and/or Toolsets."
106 )
108 if search_tool_parameters_description is not None:
109 invalid_keys = set(search_tool_parameters_description.keys()) - self._VALID_SEARCH_TOOL_PARAMS
110 if invalid_keys:
111 raise ValueError(
112 f"Invalid search_tool_parameters_description keys: {invalid_keys}. "
113 f"Valid keys are: {self._VALID_SEARCH_TOOL_PARAMS}"
114 )
116 # Store raw catalog; flattening is deferred to warm_up() so that lazy toolsets
117 # (e.g. MCPToolset with eager_connect=False) can connect first.
118 self._raw_catalog: "ToolsType" = catalog
119 self._catalog: list[Tool] = []
121 self._top_k = top_k
122 self._search_threshold = search_threshold
123 self._search_tool_name = search_tool_name
124 self._search_tool_description = search_tool_description
125 self._search_tool_parameters_description = search_tool_parameters_description
127 # Runtime state (initialized in warm_up)
128 self._discovered_tools: dict[str, Tool] = {}
129 self._bootstrap_tool: Tool | None = None
130 self._document_store: InMemoryDocumentStore | None = None
131 self._passthrough: bool | None = None
133 # Optional per-run name filter, set on the copies returned by spawn(). When set, iteration only
134 # yields tools whose name is in this set, and search is scoped to it. None means no filtering.
135 self._selected_tool_names: set[str] | None = None
137 # Initialize parent with empty tools list - we manage tools dynamically
138 super().__init__(tools=[])
140 def __add__(self, other: Tool | Toolset | list[Tool]) -> "Toolset":
141 """Concatenation is not supported for SearchableToolset."""
142 raise NotImplementedError("SearchableToolset does not support concatenation.")
144 def add(self, tool: Tool | Toolset) -> None:
145 """Adding new tools after initialization is not supported for SearchableToolset."""
146 raise NotImplementedError("SearchableToolset does not support adding new tools after initialization.")
148 def warm_up(self) -> None:
149 """
150 Prepare the toolset for use.
152 Warms up the catalog (so lazy toolsets like MCPToolset can connect) and flattens it. Above the passthrough
153 threshold, it also indexes the catalog and creates the search_tools bootstrap tool.
155 This method is idempotent: it only warms up the toolset the first time it is called.
157 :raises ValueError: If the flattened catalog contains tools with duplicate names.
158 """
159 if self._passthrough is not None:
160 return
162 # Warm up the catalog first (triggers lazy connections like MCPToolset), then flatten — lazy toolsets will
163 # have their real tools available.
164 warm_up_tools(self._raw_catalog)
165 self._catalog = flatten_tools_or_toolsets(self._raw_catalog)
166 _check_duplicate_tool_names(self._catalog)
167 self._passthrough = len(self._catalog) < self._search_threshold
169 # Build the BM25 search index only when the catalog is large enough to need discovery.
170 if not self._passthrough:
171 # shared=False keeps the BM25 index instance-local so it is freed with this toolset instead of
172 # accumulating in InMemoryDocumentStore's process-global storage (e.g. when a SearchableToolset is
173 # built per request in a served application).
174 self._document_store = InMemoryDocumentStore(shared=False)
175 documents = [
176 Document(content=f"{tool.name} {tool.description}", meta={"tool_name": tool.name})
177 for tool in self._catalog
178 ]
179 self._document_store.write_documents(documents, policy=DuplicatePolicy.OVERWRITE)
180 self._bootstrap_tool = self._create_search_tool()
182 def get_selectable_tools(self) -> list[Tool]:
183 """
184 Return the full catalog of tools that can be selected by name.
186 Iteration only exposes the search tool plus already-discovered tools, but name-based selection can target
187 any tool in the catalog, so this returns the entire flattened catalog (warming up first if needed).
189 :returns: The flattened catalog of tools.
190 """
191 self.warm_up()
192 return list(self._catalog)
194 def clear(self) -> None:
195 """
196 Clear all discovered tools.
198 This method allows resetting the toolset's discovered tools between agent runs when the same toolset instance
199 is reused. This can be useful for long-running applications to control memory usage or to start fresh searches.
200 """
201 self._discovered_tools.clear()
203 def spawn(self, selected_tool_names: set[str] | None = None) -> "SearchableToolset":
204 """
205 Return an isolated copy for a single run, carrying the given name selection.
207 The copy shares the read-only catalog and BM25 index but gets fresh discovered tools and name selection,
208 plus a bootstrap search tool bound to the copy; the selection scopes both iteration and search. This way
209 concurrent runs sharing the same configured SearchableToolset don't share discovered tools or collide on
210 the active selection.
212 :param selected_tool_names: Optional catalog tool names this run is restricted to. None means no
213 restriction.
214 :returns: A run-scoped copy of this SearchableToolset.
215 """
216 self.warm_up()
217 new = copy.copy(self)
218 new._discovered_tools = {}
219 new._selected_tool_names = set(selected_tool_names) if selected_tool_names is not None else None
220 # Rebuild the bootstrap tool so its closure is bound to the copy's discovered tools / selection
221 # rather than the original's. The document store and catalog are read-only and stay shared.
222 if not self._passthrough:
223 new._bootstrap_tool = new._create_search_tool()
224 return new
226 def _create_search_tool(self) -> Tool:
227 """Create the search_tools bootstrap tool."""
229 tool_by_name = {tool.name: tool for tool in self._catalog}
231 def search_tools(
232 tool_keywords: Annotated[
233 str,
234 "Space-separated words from tool names/descriptions (e.g. 'route weather search')."
235 " NOT the user's question or task—use vocabulary from the tools you need.",
236 ],
237 k: Annotated[int | None, f"Number of results to return (default: {self._top_k})"] = None,
238 ) -> str:
239 """
240 ALWAYS use this tool FIRST when you need to invoke some tools but don't have the right one loaded yet.
242 Provide space separated tool keywords likely to appear in tool names/descriptions
243 (e.g. 'route distance weather', 'search email').
244 Do NOT pass the user's request or task (e.g. 'things to do in X', 'user question'); matching is
245 keyword-based.
246 Returns loaded tool names; they become available immediately.
247 """
248 num_results = k if k is not None else self._top_k
250 if not tool_keywords.strip():
251 return (
252 "No tool keywords provided. Please provide space-separated words likely to appear in tool "
253 "names/descriptions (e.g. 'route weather search')."
254 )
256 # Scope the search to the selected subset if active so that top_k applies within the selected tools
257 filters = None
258 if self._selected_tool_names is not None:
259 filters = {"field": "meta.tool_name", "operator": "in", "value": list(self._selected_tool_names)}
261 # at this point, the toolset has been warmed up, so self._document_store is not None
262 results = self._document_store.bm25_retrieval( # type: ignore[union-attr]
263 query=tool_keywords, top_k=num_results, filters=filters
264 )
266 if not results:
267 return "No tools found matching these keywords. Try different keywords."
269 # Add found tools to _discovered_tools. These become available to the LLM on the next agent iteration
270 # when __iter__ is called again - the Agent re-iterates over the toolset each loop, picking up newly
271 # discovered tools.
272 # The return message here just confirms what was found; actual tool availability comes through the dynamic
273 # iteration mechanism. This way we also save tokens by not returning the full tool definitions.
274 #
275 # NOTE: The Agent can run tool calls in a step concurrently (ThreadPoolExecutor), so multiple search_tools
276 # calls can mutate self._discovered_tools from different threads at once. This is currently safe only
277 # because CPython's GIL makes individual dict assignments atomic; on a free-threaded (no-GIL) build these
278 # unguarded writes could corrupt the dict.
279 tool_names = []
280 for doc in results:
281 tool = tool_by_name[doc.meta["tool_name"]]
282 self._discovered_tools[tool.name] = tool
283 tool_names.append(tool.name)
285 return f"Found and loaded {len(tool_names)} tool(s): {', '.join(tool_names)}. Use them directly as tools."
287 bootstrap_tool = create_tool_from_function(
288 function=search_tools, name=self._search_tool_name, description=self._search_tool_description
289 )
291 # Override parameter descriptions if custom ones were provided
292 if self._search_tool_parameters_description:
293 for param_name, desc in self._search_tool_parameters_description.items():
294 if param_name in bootstrap_tool.parameters.get("properties", {}):
295 bootstrap_tool.parameters["properties"][param_name]["description"] = desc
297 return bootstrap_tool
299 def _is_selected(self, name: str) -> bool:
300 """Whether a catalog tool name is allowed by the active `_selected_tool_names` filter (None means all)."""
301 return self._selected_tool_names is None or name in self._selected_tool_names
303 def __iter__(self) -> Iterator[Tool]:
304 """
305 Iterate over available tools.
307 In passthrough mode, yields all catalog tools. Otherwise, yields the bootstrap search tool plus the
308 already-discovered tools. If `_selected_tool_names` is set, catalog/discovered tools are restricted to that
309 set, but the bootstrap search tool is always exposed so search keeps working over the selected subset.
310 Automatically calls warm_up() if needed to ensure the bootstrap tool is available.
311 """
312 # This toolset materializes everything (flattened catalog, bootstrap tool, passthrough decision) in warm_up.
313 # Without warming here, iterating before warm_up would yield nothing, so we warm up to make the toolset usable
314 # at all.
315 self.warm_up()
316 if self._passthrough:
317 yield from (tool for tool in self._catalog if self._is_selected(tool.name))
318 else:
319 if self._bootstrap_tool is not None:
320 yield self._bootstrap_tool
321 yield from (tool for tool in self._discovered_tools.values() if self._is_selected(tool.name))
323 def __contains__(self, item: str | Tool) -> bool:
324 """
325 Check if a tool is available by Tool instance or tool name string.
327 :param item: Tool instance or tool name string.
328 :returns: True if the tool is available, False otherwise.
329 """
330 if isinstance(item, str):
331 return any(tool.name == item for tool in self)
332 if isinstance(item, Tool):
333 return any(tool == item for tool in self)
334 raise TypeError(f"Invalid item type: {type(item)}. Must be Tool or str.")
336 def to_dict(self) -> dict[str, Any]:
337 """
338 Serialize the toolset to a dictionary.
340 :returns: Dictionary representation of the toolset.
341 """
342 data: dict[str, Any] = {
343 "catalog": serialize_tools_or_toolset(self._raw_catalog),
344 "top_k": self._top_k,
345 "search_threshold": self._search_threshold,
346 "search_tool_name": self._search_tool_name,
347 "search_tool_description": self._search_tool_description,
348 "search_tool_parameters_description": self._search_tool_parameters_description,
349 }
351 return {"type": generate_qualified_class_name(type(self)), "data": data}
353 @classmethod
354 def from_dict(cls, data: dict[str, Any]) -> "SearchableToolset":
355 """
356 Deserialize a toolset from a dictionary.
358 :param data: Dictionary representation of the toolset.
359 :returns: New SearchableToolset instance.
360 :raises TypeError: If a serialized catalog entry is not a subclass of Tool or Toolset.
361 """
362 inner_data = data["data"]
363 deserialize_tools_or_toolset_inplace(inner_data, key="catalog")
364 optional_keys = (
365 "top_k",
366 "search_threshold",
367 "search_tool_name",
368 "search_tool_description",
369 "search_tool_parameters_description",
370 )
371 return cls(catalog=inner_data["catalog"], **{k: inner_data[k] for k in optional_keys if k in inner_data})