# SPDX-FileCopyrightText: 2022-present deepset GmbH # # SPDX-License-Identifier: Apache-2.0 import copy from collections.abc import Iterator from typing import TYPE_CHECKING, Annotated, Any from haystack.core.serialization import generate_qualified_class_name from haystack.dataclasses import Document from haystack.document_stores.in_memory import InMemoryDocumentStore from haystack.document_stores.types import DuplicatePolicy from haystack.tools.from_function import create_tool_from_function from haystack.tools.serde_utils import deserialize_tools_or_toolset_inplace, serialize_tools_or_toolset from haystack.tools.tool import Tool, _check_duplicate_tool_names from haystack.tools.toolset import Toolset from haystack.tools.utils import flatten_tools_or_toolsets, warm_up_tools if TYPE_CHECKING: from haystack.tools import ToolsType class SearchableToolset(Toolset): """ Dynamic tool discovery from large catalogs using BM25 search. This Toolset enables LLMs to discover and use tools from large catalogs through BM25-based search. Instead of exposing all tools at once (which can overwhelm the LLM context), it provides a `search_tools` bootstrap tool that allows the LLM to find and load specific tools as needed. For very small catalogs (below `search_threshold`), acts as a simple passthrough exposing all tools directly without any discovery mechanism. ### Usage Example ```python from typing import Annotated from haystack.components.agents import Agent from haystack.components.generators.chat import OpenAIChatGenerator from haystack.dataclasses import ChatMessage from haystack.tools import SearchableToolset, tool @tool def get_weather(city: Annotated[str, "The city to get the weather for"]) -> str: '''Get the current weather for a city.''' return f"The weather in {city} is 22°C and sunny." @tool def search_web(query: Annotated[str, "The query to search the web for"]) -> str: '''Search the web for a query.''' return f"Top result for '{query}': ..." @tool def convert_currency( amount: Annotated[float, "The amount to convert"], to_currency: Annotated[str, "The currency to convert to, e.g. 'EUR'"], ) -> str: '''Convert an amount in USD to another currency.''' return f"{amount} USD is {amount * 0.9} {to_currency}" # search_threshold=2 means a catalog of 2+ tools activates discovery: the agent only sees the # `search_tools` tool and must search to load the others (set it higher for larger catalogs). toolset = SearchableToolset(catalog=[get_weather, search_web, convert_currency], search_threshold=2) agent = Agent(chat_generator=OpenAIChatGenerator(), tools=toolset) # The agent is initially provided only with the search_tools tool and will use it to find relevant tools. result = agent.run(messages=[ChatMessage.from_user("What's the weather in Milan?")]) print(result["last_message"].text) ``` """ _VALID_SEARCH_TOOL_PARAMS = {"tool_keywords", "k"} def __init__( self, catalog: "ToolsType", *, top_k: int = 3, search_threshold: int = 8, search_tool_name: str = "search_tools", search_tool_description: str | None = None, search_tool_parameters_description: dict[str, str] | None = None, ) -> None: """ Initialize the SearchableToolset. :param catalog: Source of tools - a list of Tools, list of Toolsets, or a single Toolset. :param top_k: Default number of results for search_tools. :param search_threshold: Minimum catalog size to activate search. If catalog has fewer tools, acts as passthrough (all tools visible). Default is 8. :param search_tool_name: Custom name for the bootstrap search tool. Default is "search_tools". :param search_tool_description: Custom description for the bootstrap search tool. If not provided, uses a default description. :param search_tool_parameters_description: Custom descriptions for the bootstrap search tool's parameters. Keys must be a subset of `{"tool_keywords", "k"}`. Example: `{"tool_keywords": "Keywords to find tools, e.g. 'email send'"}` """ valid_catalog = isinstance(catalog, Toolset) or ( isinstance(catalog, list) and all(isinstance(item, (Tool, Toolset)) for item in catalog) ) if not valid_catalog: raise TypeError( f"Invalid catalog type: {type(catalog)}. Expected Tool, Toolset, or list of Tools and/or Toolsets." ) if search_tool_parameters_description is not None: invalid_keys = set(search_tool_parameters_description.keys()) - self._VALID_SEARCH_TOOL_PARAMS if invalid_keys: raise ValueError( f"Invalid search_tool_parameters_description keys: {invalid_keys}. " f"Valid keys are: {self._VALID_SEARCH_TOOL_PARAMS}" ) # Store raw catalog; flattening is deferred to warm_up() so that lazy toolsets # (e.g. MCPToolset with eager_connect=False) can connect first. self._raw_catalog: "ToolsType" = catalog self._catalog: list[Tool] = [] self._top_k = top_k self._search_threshold = search_threshold self._search_tool_name = search_tool_name self._search_tool_description = search_tool_description self._search_tool_parameters_description = search_tool_parameters_description # Runtime state (initialized in warm_up) self._discovered_tools: dict[str, Tool] = {} self._bootstrap_tool: Tool | None = None self._document_store: InMemoryDocumentStore | None = None self._passthrough: bool | None = None self._is_warmed_up = False # 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: """Adding new tools after initialization is not supported for SearchableToolset.""" raise NotImplementedError("SearchableToolset does not support adding new tools after initialization.") def warm_up(self) -> None: """ Prepare the toolset for use. Warms up the catalog (so lazy toolsets like MCPToolset can connect) and flattens it. Above the passthrough threshold, it also indexes the catalog and creates the search_tools bootstrap tool. This method is idempotent: it only warms up the toolset the first time it is called. :raises ValueError: If the flattened catalog contains tools with duplicate names. """ if self._is_warmed_up: return # Warm up the catalog first (triggers lazy connections like MCPToolset), then flatten — lazy toolsets will # have their real tools available. warm_up_tools(self._raw_catalog) self._catalog = flatten_tools_or_toolsets(self._raw_catalog) _check_duplicate_tool_names(self._catalog) self._passthrough = len(self._catalog) < self._search_threshold # Build the BM25 search index only when the catalog is large enough to need discovery. if not self._passthrough: # shared=False keeps the BM25 index instance-local so it is freed with this toolset instead of # accumulating in InMemoryDocumentStore's process-global storage (e.g. when a SearchableToolset is # built per request in a served application). self._document_store = InMemoryDocumentStore(shared=False) documents = [ Document(content=f"{tool.name} {tool.description}", meta={"tool_name": tool.name}) for tool in self._catalog ] 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. Iteration only exposes the search tool plus already-discovered tools, but name-based selection can target any tool in the catalog, so this returns the entire flattened catalog (warming up first if needed). :returns: The flattened catalog of tools. """ if not self._is_warmed_up: self.warm_up() return list(self._catalog) def clear(self) -> None: """ Clear all discovered tools. This method allows resetting the toolset's discovered tools between agent runs when the same toolset instance is reused. This can be useful for long-running applications to control memory usage or to start fresh searches. """ self._discovered_tools.clear() def spawn(self) -> "SearchableToolset": """ Return an isolated copy for a single run. 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. :returns: A run-scoped copy of this SearchableToolset. """ if not self._is_warmed_up: self.warm_up() new = copy.copy(self) new._discovered_tools = {} new._selected_tool_names = 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: new._bootstrap_tool = new._create_search_tool() return new def _create_search_tool(self) -> Tool: """Create the search_tools bootstrap tool.""" tool_by_name = {tool.name: tool for tool in self._catalog} def search_tools( tool_keywords: Annotated[ str, "Space-separated words from tool names/descriptions (e.g. 'route weather search')." " NOT the user's question or task—use vocabulary from the tools you need.", ], k: Annotated[int | None, f"Number of results to return (default: {self._top_k})"] = None, ) -> str: """ ALWAYS use this tool FIRST when you need to invoke some tools but don't have the right one loaded yet. Provide space separated tool keywords likely to appear in tool names/descriptions (e.g. 'route distance weather', 'search email'). Do NOT pass the user's request or task (e.g. 'things to do in X', 'user question'); matching is keyword-based. Returns loaded tool names; they become available immediately. """ num_results = k if k is not None else self._top_k if not tool_keywords.strip(): return ( "No tool keywords provided. Please provide space-separated words likely to appear in tool " "names/descriptions (e.g. 'route weather search')." ) # Scope the search to the selected subset if active so that top_k applies within the selected tools filters = None if self._selected_tool_names is not None: filters = {"field": "meta.tool_name", "operator": "in", "value": list(self._selected_tool_names)} # at this point, the toolset has been warmed up, so self._document_store is not None results = self._document_store.bm25_retrieval( # type: ignore[union-attr] query=tool_keywords, top_k=num_results, filters=filters ) if not results: return "No tools found matching these keywords. Try different keywords." # Add found tools to _discovered_tools. These become available to the LLM on the next agent iteration # when __iter__ is called again - the Agent re-iterates over the toolset each loop, picking up newly # discovered tools. # The return message here just confirms what was found; actual tool availability comes through the dynamic # iteration mechanism. This way we also save tokens by not returning the full tool definitions. # # NOTE: The Agent can run tool calls in a step concurrently (ThreadPoolExecutor), so multiple search_tools # calls can mutate self._discovered_tools from different threads at once. This is currently safe only # because CPython's GIL makes individual dict assignments atomic; on a free-threaded (no-GIL) build these # unguarded writes could corrupt the dict. tool_names = [] for doc in results: tool = tool_by_name[doc.meta["tool_name"]] self._discovered_tools[tool.name] = tool tool_names.append(tool.name) return f"Found and loaded {len(tool_names)} tool(s): {', '.join(tool_names)}. Use them directly as tools." bootstrap_tool = create_tool_from_function( function=search_tools, name=self._search_tool_name, description=self._search_tool_description ) # Override parameter descriptions if custom ones were provided if self._search_tool_parameters_description: for param_name, desc in self._search_tool_parameters_description.items(): if param_name in bootstrap_tool.parameters.get("properties", {}): bootstrap_tool.parameters["properties"][param_name]["description"] = desc return bootstrap_tool def _is_selected(self, name: str) -> bool: """Whether a catalog tool name is allowed by the active `_selected_tool_names` filter (None means all).""" return self._selected_tool_names is None or name in self._selected_tool_names def __iter__(self) -> Iterator[Tool]: """ Iterate over available tools. In passthrough mode, yields all catalog tools. Otherwise, yields the bootstrap search tool plus the already-discovered tools. If `_selected_tool_names` is set, catalog/discovered tools are restricted to that 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. # 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() if self._passthrough: yield from (tool for tool in self._catalog if self._is_selected(tool.name)) else: if self._bootstrap_tool is not None: 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. :param item: Tool instance or tool name string. :returns: True if the tool is available, False otherwise. """ if isinstance(item, str): return any(tool.name == item for tool in self) if isinstance(item, Tool): 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. :returns: Dictionary representation of the toolset. """ data: dict[str, Any] = { "catalog": serialize_tools_or_toolset(self._raw_catalog), "top_k": self._top_k, "search_threshold": self._search_threshold, "search_tool_name": self._search_tool_name, "search_tool_description": self._search_tool_description, "search_tool_parameters_description": self._search_tool_parameters_description, } return {"type": generate_qualified_class_name(type(self)), "data": data} @classmethod def from_dict(cls, data: dict[str, Any]) -> "SearchableToolset": """ Deserialize a toolset from a dictionary. :param data: Dictionary representation of the toolset. :returns: New SearchableToolset instance. :raises TypeError: If a serialized catalog entry is not a subclass of Tool or Toolset. """ inner_data = data["data"] deserialize_tools_or_toolset_inplace(inner_data, key="catalog") optional_keys = ( "top_k", "search_threshold", "search_tool_name", "search_tool_description", "search_tool_parameters_description", ) return cls(catalog=inner_data["catalog"], **{k: inner_data[k] for k in optional_keys if k in inner_data})