From 8579f560ca1b070166b8c9988a6ed54d8b3badc1 Mon Sep 17 00:00:00 2001 From: Yakir Shriker Date: Sun, 31 May 2026 19:07:10 +0300 Subject: [PATCH] feat: client-side tool search (load_tools) to cut tool-schema tokens MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Large MCP deployments pay a 10-60k "MCP tax": every tool's full schema is in the system prompt on every turn, before the user even sends a message. This adds an opt-in HOLMES_TOOL_SEARCH_ENABLED mode (default off) that holds heavy tool schemas out of context and exposes a single `load_tools(query)` function tool. The model calls load_tools to discover the tools it needs; they become callable on the next step via the existing tool-refetch loop. Design: load_tools is a plain function tool that HolmesGPT executes itself (ToolCallingLLM._handle_load_tools -> ToolExecutor.search_deferred_tools). A function tool works identically across the Anthropic API, Bedrock, OpenAI and any LiteLLM-fronted gateway — no beta header, no server tools. - ToolExecutor: get_visible_tools_openai_format(loaded_names) returns core (non-MCP) tools + load_tools + already-loaded MCP tools; search_deferred_tools(query) regex/substring-searches held-back MCP tools by name/description/toolset. - ToolCallingLLM: per-conversation _loaded_tool_names; _get_tools serves the visible set when enabled; load_tools is intercepted and never hits the ToolExecutor. Only MCP-type toolsets are deferred (built-ins like kubernetes/bash stay loaded so common ops don't pay a search hop); the policy is a single DEFERRABLE_TOOLSET_TYPES set, easy to extend. Verified end-to-end against a LiteLLM gateway (Claude Opus): a "hi" turn and a real AWS investigation both work; the investigation goes load_tools -> tool loads -> call the tool -> answer, with no broken tool_use/tool_result and ~24-89% fewer prompt tokens depending on how much of the catalog is in play. Signed-off-by: Yakir Shriker --- holmes/common/env_vars.py | 4 + holmes/core/tool_calling_llm.py | 66 +++++++++++++++- holmes/core/tool_search.py | 54 +++++++++++++ holmes/core/tools_utils/tool_executor.py | 83 +++++++++++++++++++- tests/core/test_tool_executor.py | 77 +++++++++++++++++++ tests/core/test_tool_search.py | 10 +++ tests/test_tool_calling_llm.py | 96 ++++++++++++++++++++++++ 7 files changed, 387 insertions(+), 3 deletions(-) create mode 100644 holmes/core/tool_search.py create mode 100644 tests/core/test_tool_search.py diff --git a/holmes/common/env_vars.py b/holmes/common/env_vars.py index cbb33ed140..a522b86157 100644 --- a/holmes/common/env_vars.py +++ b/holmes/common/env_vars.py @@ -62,6 +62,10 @@ def load_bool(env_var, default: Optional[bool]) -> Optional[bool]: EXTRA_HEADERS = os.environ.get("EXTRA_HEADERS", "") THINKING = os.environ.get("THINKING", "") +# Client-side tool search. When enabled, heavy tool schemas (MCP toolsets) are held +# back from the default tool list and a load_tools meta-tool is exposed so the model +# loads them on demand instead of putting the whole catalog in context every turn. +TOOL_SEARCH_ENABLED = load_bool("HOLMES_TOOL_SEARCH_ENABLED", False) REASONING_EFFORT = os.environ.get("REASONING_EFFORT", "").strip().lower() diff --git a/holmes/core/tool_calling_llm.py b/holmes/core/tool_calling_llm.py index d3e034c198..1f5c2632d7 100644 --- a/holmes/core/tool_calling_llm.py +++ b/holmes/core/tool_calling_llm.py @@ -23,9 +23,11 @@ LOG_LLM_USAGE_RESPONSE, RESET_REPEATED_TOOL_CALL_CHECK_AFTER_COMPACTION, TEMPERATURE, + TOOL_SEARCH_ENABLED, load_bool, ) from holmes.core.llm import LLM +from holmes.core.tool_search import LOAD_TOOLS_NAME from holmes.core.llm_usage import RequestStats from holmes.core.models import ( FrontendToolResult, @@ -205,6 +207,9 @@ def __init__( self.tool_results_dir = tool_results_dir self._skill_in_use: bool = False + # Names of deferred (held-back) tools the model has loaded via load_tools in + # this conversation. Heavy MCP tool schemas stay out of context until loaded. + self._loaded_tool_names: set[str] = set() def with_executor(self, tool_executor: ToolExecutor) -> "ToolCallingLLM": """Return a shallow copy with a different ToolExecutor. @@ -220,15 +225,18 @@ def with_executor(self, tool_executor: ToolExecutor) -> "ToolCallingLLM": tracer=self.tracer, ) # Preserve transient state so resumed turns keep access to - # skill-unlocked (restricted) tools. + # skill-unlocked (restricted) tools and already-loaded (tool-search) tools. clone._skill_in_use = self._skill_in_use + clone._loaded_tool_names = set(self._loaded_tool_names) return clone def reset_interaction_state(self) -> None: """ - For interactive loop, reset skills in use + For interactive loop, reset transient per-conversation state (skills in + use, tool-search loaded tools) so a new conversation starts fresh. """ self._skill_in_use = False + self._loaded_tool_names.clear() def _supports_vision(self) -> bool: """Check if vision/multimodal input is enabled. @@ -561,6 +569,16 @@ def _get_tools(self) -> list: replace _connect placeholders for authenticated users. """ user_id = (self._request_context or {}).get("user_id") if hasattr(self, "_request_context") else None + if TOOL_SEARCH_ENABLED: + # Progressive disclosure: hold heavy (MCP) tool schemas out of context and + # expose a load_tools meta-tool. Tools the model has loaded this conversation + # are included. The agent loop re-fetches tools each step (see call_stream), + # so newly loaded tools appear on the next turn. + return self.tool_executor.get_visible_tools_openai_format( + self._loaded_tool_names, + include_restricted=self._should_include_restricted_tools(), + user_id=user_id, + ) return self.tool_executor.get_all_tools_openai_format( include_restricted=self._should_include_restricted_tools(), user_id=user_id, @@ -867,6 +885,40 @@ def _log_tool_call_result( metadata=metadata, ) + def _handle_load_tools( + self, tool_call_id: str, tool_params: Dict[str, Any] + ) -> ToolCallResult: + """Execute the load_tools meta-tool (tool search): find held-back tools matching + the query and mark them loaded, so the next agent step can call them.""" + query = (tool_params or {}).get("query", "") or "" + matched = self.tool_executor.search_deferred_tools(query) + self._loaded_tool_names.update(matched) + if matched: + lines = [] + for name in matched: + tool = self.tool_executor.get_tool_by_name(name) + desc = "" + if tool and tool.description: + desc = tool.description.strip().splitlines()[0] + lines.append(f"- {name}: {desc}" if desc else f"- {name}") + data = ( + f"Loaded {len(matched)} tool(s) matching '{query}'. They are now " + "available — call them directly on your next step:\n" + "\n".join(lines) + ) + else: + data = ( + f"No tools matched '{query}'. Try broader keywords or an integration " + "name (e.g. aws, github, pagerduty, kafka, tempo, opensearch, grafana)." + ) + return ToolCallResult( + tool_call_id=tool_call_id, + tool_name=LOAD_TOOLS_NAME, + description=f"load_tools(query={query!r})", + result=StructuredToolResult( + status=StructuredToolResultStatus.SUCCESS, data=data + ), + ) + def _invoke_llm_tool_call( self, tool_to_call: ChatCompletionMessageToolCall, @@ -915,6 +967,16 @@ def _invoke_llm_tool_call( f"Failed to parse arguments for tool: {tool_name}. args: {tool_arguments}" ) + # load_tools is a Holmes-owned meta-tool (tool search). Handle it here so it + # never reaches the ToolExecutor (it isn't a toolset tool) and so the loaded + # set is updated before the next _get_tools() re-fetch exposes the new tools. + if tool_name == LOAD_TOOLS_NAME: + tool_call_result = self._handle_load_tools(tool_id, tool_params) + ToolCallingLLM._log_tool_call_result( + tool_span, tool_call_result, enable_tool_approval + ) + return tool_call_result + tool_response = None if not user_approved: tool_response = prevent_overly_repeated_tool_call( diff --git a/holmes/core/tool_search.py b/holmes/core/tool_search.py new file mode 100644 index 0000000000..f3f8c66366 --- /dev/null +++ b/holmes/core/tool_search.py @@ -0,0 +1,54 @@ +"""Client-side tool search (progressive tool disclosure). + +When ``HOLMES_TOOL_SEARCH_ENABLED`` is set, heavy tool schemas (MCP toolsets) are +NOT loaded into the model's context up front. Instead the model is given a single +``load_tools`` function tool; it calls that with a query to discover and load the +tools it needs, and they become available on the next step. + +This is a deliberate alternative to Anthropic's server-side tool-search / +``defer_loading`` beta: that path is mis-handled by LiteLLM (server_tool_use is +converted to a client tool_use and result blocks are dropped — BerriAI/litellm +issues #17737 and #28083), which breaks multi-step tool use through Bedrock-backed +gateways. ``load_tools`` is a plain function tool that HolmesGPT executes itself, so +it works identically across the Anthropic API, Bedrock, OpenAI and any +LiteLLM-fronted gateway — no beta header, no server tools. + +On large MCP deployments this keeps 10-60k tokens of tool schemas out of every turn +until they're actually needed. +""" + +# Name of the meta-tool the model calls to discover/load held-back tools. +LOAD_TOOLS_NAME = "load_tools" + +# OpenAI function-tool definition for load_tools. It is a normal function tool — +# HolmesGPT intercepts and executes it (see ToolCallingLLM._handle_load_tools). +LOAD_TOOLS_TOOL: dict = { + "type": "function", + "function": { + "name": LOAD_TOOLS_NAME, + "description": ( + "Search for and load additional tools that are not loaded by default. " + "To keep the context small, heavy integrations are hidden until needed — " + "for example: AWS, GitHub, PagerDuty, Kafka/MSK, Tempo tracing, " + "OpenSearch/Elasticsearch metrics, and Grafana. Call this with keywords " + "describing the capability you need (e.g. 'aws s3', 'github pull request', " + "'kafka consumer lag', 'pagerduty incidents', 'tempo trace', " + "'opensearch metrics'). The matching tools become available to call on your " + "next step. If nothing matches, retry with broader keywords or an " + "integration name. Tools already loaded in this conversation stay available." + ), + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": ( + "Keywords (or a regular expression) describing the capability or " + "integration you need, e.g. 'aws ec2', 'github', 'kafka lag'." + ), + } + }, + "required": ["query"], + }, + }, +} diff --git a/holmes/core/tools_utils/tool_executor.py b/holmes/core/tools_utils/tool_executor.py index d01cbce0a9..1a80ffc7d9 100644 --- a/holmes/core/tools_utils/tool_executor.py +++ b/holmes/core/tools_utils/tool_executor.py @@ -1,18 +1,27 @@ import logging -from typing import Any, Dict, List, Optional +import re +from typing import List, Optional import sentry_sdk from holmes.core.init_event import EventCallback, StatusEvent, StatusEventKind +from holmes.core.tool_search import LOAD_TOOLS_TOOL from holmes.core.tools import ( Tool, Toolset, ToolsetStatusEnum, + ToolsetType, ) from holmes.core.tools_utils.oauth_tool_connector import OAuthToolConnector display_logger = logging.getLogger("holmes.display.tool_executor") +# Toolset types whose tools are "deferred" (held back from the default tool list and +# loaded on demand via load_tools) when tool search is enabled. MCP servers dominate +# tool-schema bloat; built-in toolsets (kubernetes, bash, etc.) stay loaded so common +# operations don't pay a search hop. +DEFERRABLE_TOOLSET_TYPES = frozenset({ToolsetType.MCP}) + class ToolExecutor: def __init__(self, toolsets: List[Toolset], on_event: EventCallback = None): @@ -153,3 +162,75 @@ def _get_base_tools(self, include_restricted: bool = True) -> list: continue tools.append(tool.get_openai_format()) return tools + + # ── Tool search (progressive disclosure) ─────────────────────────── + + def _is_deferrable(self, tool_name: str) -> bool: + """Whether a tool is held back from the default list until loaded on demand.""" + toolset = self._tool_to_toolset.get(tool_name) + return toolset is not None and toolset.type in DEFERRABLE_TOOLSET_TYPES + + def get_visible_tools_openai_format( + self, + loaded_tool_names: Optional[set] = None, + include_restricted: bool = True, + user_id: Optional[str] = None, + ) -> list: + """Tool-search mode: return core (non-deferrable) tools + the ``load_tools`` + meta-tool + any deferrable tools that have already been loaded. Heavy MCP tool + schemas stay out of context until the model loads them via ``load_tools``. + """ + loaded = loaded_tool_names or set() + tools = [] + for tool in self.tools_by_name.values(): + if not include_restricted and tool._is_restricted(): + continue + if self._is_deferrable(tool.name) and tool.name not in loaded: + continue + tools.append(tool.get_openai_format()) + tools.append(LOAD_TOOLS_TOOL) + return self.oauth_connector.apply_user_tools(tools, user_id, self._tool_to_toolset) + + def search_deferred_tools(self, query: str, limit: int = 8) -> List[str]: + """Return names of held-back (deferrable) tools matching ``query``. + + ``query`` is treated as a regular expression (Python ``re.search``); if it isn't + valid regex it falls back to a case-insensitive substring match. Each tool's + name, description and owning toolset's name/description are searched, so a query + like ``"kafka"`` matches every tool in the kafka toolset. + """ + query = (query or "").strip() + if not query: + return [] + try: + pattern = re.compile(query, re.IGNORECASE) + + def matches(text: str) -> bool: + return bool(pattern.search(text)) + except re.error: + needle = query.lower() + + def matches(text: str) -> bool: + return needle in text.lower() + + results: List[str] = [] + for name, tool in self.tools_by_name.items(): + if not self._is_deferrable(name): + continue + ts = self._tool_to_toolset.get(name) + haystack = " ".join( + filter( + None, + [ + name, + tool.description or "", + ts.name if ts else "", + (ts.description or "") if ts else "", + ], + ) + ) + if matches(haystack): + results.append(name) + if len(results) >= limit: + break + return results diff --git a/tests/core/test_tool_executor.py b/tests/core/test_tool_executor.py index 24b6b18559..69faebb93a 100644 --- a/tests/core/test_tool_executor.py +++ b/tests/core/test_tool_executor.py @@ -106,3 +106,80 @@ def test_ensure_toolset_initialized_failure_blocks_subsequent_calls(): # The callable should only have been invoked once (lazy init is not retried) mock_callable.assert_called_once() + + +# ── Tool search (progressive disclosure) ───────────────────────────────────── + +from holmes.core.tool_search import LOAD_TOOLS_NAME # noqa: E402 +from holmes.core.tools import ( # noqa: E402 + StructuredToolResult, + StructuredToolResultStatus, + Tool, + ToolInvokeContext, + Toolset, + ToolsetType, +) + + +class _SearchTool(Tool): + def _invoke(self, params, context: ToolInvokeContext) -> StructuredToolResult: + return StructuredToolResult(status=StructuredToolResultStatus.SUCCESS) + + def get_parameterized_one_liner(self, params) -> str: + return "" + + +class _SearchToolset(Toolset): + def __init__(self, name, toolset_type, tool_specs, **kw): + super().__init__(name=name, description=f"{name} toolset", type=toolset_type, **kw) + self.tools = [_SearchTool(name=n, description=d) for n, d in tool_specs] + + +def _executor_with_search() -> ToolExecutor: + core = _SearchToolset( + "kubernetes/core", ToolsetType.BUILTIN, [("kubectl_get", "get k8s resources")] + ) + aws = _SearchToolset( + "aws_api", ToolsetType.MCP, [("call_aws", "run an aws cli command")] + ) + kafka = _SearchToolset( + "kafka", ToolsetType.MCP, [("list_clusters", "list MSK clusters")] + ) + for ts in (core, aws, kafka): + ts.status = ToolsetStatusEnum.ENABLED + return ToolExecutor(toolsets=[core, aws, kafka]) + + +def test_visible_tools_hide_mcp_and_include_load_tools(): + ex = _executor_with_search() + names = {t["function"]["name"] for t in ex.get_visible_tools_openai_format(set())} + assert "kubectl_get" in names # built-in core stays loaded + assert LOAD_TOOLS_NAME in names # meta-tool exposed + assert "call_aws" not in names # MCP tools held back + assert "list_clusters" not in names + + +def test_visible_tools_include_loaded_mcp_tool(): + ex = _executor_with_search() + names = { + t["function"]["name"] + for t in ex.get_visible_tools_openai_format({"call_aws"}) + } + assert "call_aws" in names # loaded on demand → now visible + assert "list_clusters" not in names # other MCP tool still hidden + + +def test_search_deferred_tools_matches_mcp_only(): + ex = _executor_with_search() + assert ex.search_deferred_tools("aws") == ["call_aws"] + assert ex.search_deferred_tools("kafka") == ["list_clusters"] # toolset-name match + assert ex.search_deferred_tools("kubectl") == [] # core tool isn't deferrable + assert ex.search_deferred_tools("no-such-thing") == [] + + +def test_search_deferred_tools_falls_back_for_invalid_regex(): + # A malformed regex (e.g. the model emits "[" or "(") must not raise — it falls + # back to a case-insensitive substring match (no match here → empty list). + ex = _executor_with_search() + assert ex.search_deferred_tools("[") == [] + assert ex.search_deferred_tools("(") == [] diff --git a/tests/core/test_tool_search.py b/tests/core/test_tool_search.py new file mode 100644 index 0000000000..f127952307 --- /dev/null +++ b/tests/core/test_tool_search.py @@ -0,0 +1,10 @@ +from holmes.core.tool_search import LOAD_TOOLS_NAME, LOAD_TOOLS_TOOL + + +def test_load_tools_tool_is_a_function_tool(): + # Must be a normal OpenAI function tool so HolmesGPT can execute it and so it + # works across every provider (no Anthropic server-tool / beta dependency). + assert LOAD_TOOLS_TOOL["type"] == "function" + assert LOAD_TOOLS_TOOL["function"]["name"] == LOAD_TOOLS_NAME + assert "query" in LOAD_TOOLS_TOOL["function"]["parameters"]["properties"] + assert LOAD_TOOLS_TOOL["function"]["parameters"]["required"] == ["query"] diff --git a/tests/test_tool_calling_llm.py b/tests/test_tool_calling_llm.py index 2118cdb225..b0bad78cfc 100644 --- a/tests/test_tool_calling_llm.py +++ b/tests/test_tool_calling_llm.py @@ -1702,3 +1702,99 @@ def test_noop_tool_included_in_tools_list(self, _mock_limit, make_ai, mock_llm, tool_names = [t["function"]["name"] for t in tools_sent] assert "kubectl_get" in tool_names, "Backend tool should be included" assert "navigate_to_page" in tool_names, "Noop tool should be included" + + +# ── Tool search: load_tools (_get_tools routing + interception) ─────────────── + + +def test_get_tools_uses_visible_listing_when_search_enabled( + make_ai, mock_llm, mock_tool_executor +): + ai = make_ai() + with patch("holmes.core.tool_calling_llm.TOOL_SEARCH_ENABLED", True): + ai._get_tools() + mock_tool_executor.get_visible_tools_openai_format.assert_called_once() + + +def test_get_tools_uses_full_listing_when_search_disabled( + make_ai, mock_llm, mock_tool_executor +): + ai = make_ai() + with patch("holmes.core.tool_calling_llm.TOOL_SEARCH_ENABLED", False): + ai._get_tools() + mock_tool_executor.get_all_tools_openai_format.assert_called_once() + mock_tool_executor.get_visible_tools_openai_format.assert_not_called() + + +def test_handle_load_tools_loads_matches_and_returns_success( + make_ai, mock_tool_executor +): + mock_tool_executor.search_deferred_tools.return_value = ["call_aws", "list_buckets"] + mock_tool_executor.get_tool_by_name.return_value = None + ai = make_ai() + + result = ai._handle_load_tools("tc1", {"query": "aws"}) + + assert ai._loaded_tool_names == {"call_aws", "list_buckets"} + assert result.tool_call_id == "tc1" + assert result.result.status == StructuredToolResultStatus.SUCCESS + assert "call_aws" in result.result.data + + +def test_handle_load_tools_no_match(make_ai, mock_tool_executor): + mock_tool_executor.search_deferred_tools.return_value = [] + ai = make_ai() + + result = ai._handle_load_tools("tc2", {"query": "zzz"}) + + assert ai._loaded_tool_names == set() + assert "No tools matched" in result.result.data + + +def test_reset_interaction_state_clears_loaded_tools(make_ai): + ai = make_ai() + ai._loaded_tool_names = {"call_aws", "list_buckets"} + ai.reset_interaction_state() + assert ai._loaded_tool_names == set() + + +@patch(LIMIT_PATCH, side_effect=_make_context_limiter_passthrough) +def test_load_tools_intercepted_during_invocation( + _mock_limit, make_ai, mock_llm, mock_tool_executor +): + """A model tool call named load_tools is intercepted in _invoke_llm_tool_call, + runs the search, updates _loaded_tool_names, and the loop continues to a final + answer (never routed to the ToolExecutor as a normal tool).""" + from holmes.core.tool_search import LOAD_TOOLS_NAME + + load_tools_tc = _make_mock_tool_call( + tool_call_id="tc_load", tool_name=LOAD_TOOLS_NAME, arguments={"query": "aws"} + ) + mock_llm.completion.side_effect = [ + _make_llm_response(content="searching", tool_calls=[load_tools_tc]), + _make_llm_response(content="Found tools", tool_calls=None), + ] + mock_tool_executor.search_deferred_tools.return_value = ["call_aws"] + mock_tool_executor.get_tool_by_name.return_value = None + + ai = make_ai() + result = ai.call([{"role": "user", "content": "find aws tools"}]) + + # Intercepted: the search ran and updated state (the meta-tool is handled here, + # not routed to the ToolExecutor as a normal toolset tool). + mock_tool_executor.search_deferred_tools.assert_called_once() + assert ai._loaded_tool_names == {"call_aws"} + assert result.result == "Found tools" + + +def test_loaded_tool_names_preserved_across_clone(make_ai): + """with_executor() carries _loaded_tool_names as an independent copy.""" + ai = make_ai() + ai._loaded_tool_names = {"call_aws", "list_buckets"} + + cloned_ai = ai.with_executor(MagicMock(spec=ToolExecutor)) + + assert cloned_ai._loaded_tool_names == {"call_aws", "list_buckets"} + # A copy, not a shared reference. + cloned_ai._loaded_tool_names.add("new_tool") + assert "new_tool" not in ai._loaded_tool_names