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