diff --git a/raven/agent/loop/main.py b/raven/agent/loop/main.py index 0bd8b9a..3a37731 100644 --- a/raven/agent/loop/main.py +++ b/raven/agent/loop/main.py @@ -648,7 +648,6 @@ def _register_default_tools(self) -> None: from raven.agent.tools.tool_search import ( DEFAULT_ALWAYS_VISIBLE, ToolCallTool, - ToolDescribeTool, ToolSearchController, ToolSearchStrategy, ToolSearchTool, @@ -661,7 +660,6 @@ def _register_default_tools(self) -> None: search_result_limit=cfg.search_result_limit, ) self.tools.register(ToolSearchTool(self.tool_search_controller)) - self.tools.register(ToolDescribeTool(self.tool_search_controller)) self.tools.register(ToolCallTool(self.tool_search_controller)) # ``first=True``: filter the tool list before CacheOptimizer marks # the final tool with ``cache_control`` (else the marked tool may be diff --git a/raven/agent/tools/tool_index.py b/raven/agent/tools/tool_index.py index c7d37c7..3570a04 100644 --- a/raven/agent/tools/tool_index.py +++ b/raven/agent/tools/tool_index.py @@ -1,36 +1,48 @@ """BM25 keyword index over the agent's tool catalog. -Backs ``tool_search``: ranks tools by a query over their name and description, -reusing the shared Okapi BM25 in :mod:`raven.utils.bm25` (CJK-aware, so -Chinese queries match). The name is repeated in the indexed text so a hit on -the tool name outranks an incidental description hit — the same field-weighting -idea as the skill ``LocalPool``. - -The index rebuilds only when the catalog's ``(name, description)`` set changes -(startup + each MCP connect is one burst); steady-state ``search`` is one -query-side tokenize plus a BM25 dot product over precomputed ``doc_freqs``. +Backs ``tool_search``: ranks tools by a query over each tool's name, +description and parameter schema, reusing the shared Okapi BM25 in +:mod:`raven.utils.bm25` (CJK-aware, so Chinese queries match). The name is +repeated in the indexed text so a hit on the tool name outranks an incidental +body hit — the same field-weighting idea as the skill ``LocalPool``. Parameter +names, their descriptions and enum values are folded into the indexed body: +many tools carry the discriminating keywords (a repo owner, a channel id, an +image size) in the schema rather than the one-line description, so indexing the +schema lifts recall without touching the schema the model sees. + +The index rebuilds only when the catalog's ``(name, description, parameters)`` +set changes (startup + each MCP connect is one burst); steady-state ``search`` +is one query-side tokenize plus a BM25 dot product over precomputed +``doc_freqs``. Each ``ensure`` flattens every tool's text to compute the +signature (cheap string work); the expensive BM25 rebuild is skipped on a hit. A single process-level slot caches the most recently built index keyed by that same signature, so the many short-lived agent loops a resident process spins up over one fixed tool set reuse one prebuilt BM25 rather than each rebuilding an identical one. A single slot suffices because only the main loop feeds this index (subagents run a separate minimal tool set), so one signature dominates; -keying on description (not just name) also forces a rebuild when a tool's -description changes under hot-reload. +keying on the full indexed text (not just the name) also forces a rebuild when +a tool's description or parameters change under hot-reload. """ from __future__ import annotations import threading -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from raven.utils.bm25 import BM25Okapi, tokenize if TYPE_CHECKING: from raven.agent.tools.base import Tool -# Catalog signature: the exact inputs to the BM25 build (name + description). +# Catalog signature: one (name, indexed-text) pair per tool. The indexed text +# embeds description + parameter schema, so the signature changes whenever any +# indexed field does — the exact trigger for a rebuild. _Signature = frozenset[tuple[str, str]] + +# Cap schema-walk recursion so a pathologically nested MCP schema can't blow the +# stack or balloon the indexed text; tool schemas are shallow in practice. +_MAX_SCHEMA_DEPTH = 6 # A built index: the names list paired with the BM25 built from it (same # ordering). Read-only for searchers, so safe to share across loops. _BuiltIndex = tuple[list[str], BM25Okapi] @@ -44,22 +56,61 @@ _cached_index: _BuiltIndex | None = None +def _schema_text(parameters: "dict[str, Any] | None") -> str: + """Natural-language keywords drawn from a JSON-Schema parameter block. + + Collects property names, their ``description`` strings and enum values, + recursing through nested objects and array items. Types and structural + keywords are skipped — they carry no query signal. + """ + if not isinstance(parameters, dict): + return "" + parts: list[str] = [] + + def walk(node: Any, depth: int) -> None: + if depth > _MAX_SCHEMA_DEPTH or not isinstance(node, dict): + return + props = node.get("properties") + if isinstance(props, dict): + for key, sub in props.items(): + parts.append(str(key)) + walk(sub, depth + 1) + desc = node.get("description") + if isinstance(desc, str): + parts.append(desc) + enum = node.get("enum") + if isinstance(enum, list): + parts.extend(str(v) for v in enum if isinstance(v, (str, int, float))) + items = node.get("items") + if isinstance(items, dict): + walk(items, depth + 1) + + walk(parameters, 0) + return " ".join(parts) + + def _format_tool_text(tool: "Tool") -> str: - """Indexed text for one tool. name ×3, description ×1 — TF-based field - weighting so a query term on the name dominates a description hit.""" + """Indexed text for one tool: name ×3, then description + parameter-schema + keywords ×1 — TF-based field weighting so a query term on the name + dominates a hit in the body.""" name = tool.name - desc = tool.description or "" - return f"{name} {name} {name} {desc}" + body = f"{tool.description or ''} {_schema_text(tool.parameters)}".strip() + return f"{name} {name} {name} {body}".rstrip() -def _get_or_build(sig: _Signature, tools: "list[Tool]") -> _BuiltIndex: - """Return the shared prebuilt index for ``sig``, building it once on miss.""" +def _get_or_build(sig: _Signature, items: list[tuple[str, str]]) -> _BuiltIndex: + """Return the shared prebuilt index for ``sig``, building it once on miss. + + ``items`` are ``(name, indexed-text)`` pairs already flattened by the + caller, so the text is computed once per ``ensure`` (for the signature) and + reused here rather than recomputed. + """ global _cached_sig, _cached_index with _cache_lock: if sig == _cached_sig and _cached_index is not None: return _cached_index - names = [t.name for t in tools] - corpus = [tokenize(_format_tool_text(t)) for t in tools] + names = [name for name, _ in items] + corpus = [tokenize(text) for _, text in items] built: _BuiltIndex = (names, BM25Okapi(corpus)) with _cache_lock: _cached_sig, _cached_index = sig, built @@ -67,7 +118,7 @@ def _get_or_build(sig: _Signature, tools: "list[Tool]") -> _BuiltIndex: class ToolIndex: - """Prebuilt BM25 over a tool catalog, rebuilt on (name, description) change. + """Prebuilt BM25 over a tool catalog, rebuilt on (name, description, parameters) change. Thread-safety: a lock guards the ``(names, BM25Okapi)`` pair. The expensive build is delegated to the shared ``_get_or_build`` slot; the lock here is @@ -81,12 +132,13 @@ def __init__(self) -> None: self._bm25: BM25Okapi | None = None def ensure(self, tools: "list[Tool]") -> None: - """Adopt the shared index iff the (name, description) set changed.""" - sig: _Signature = frozenset((t.name, t.description or "") for t in tools) + """Adopt the shared index iff the (name, description, parameters) set changed.""" + items = [(t.name, _format_tool_text(t)) for t in tools] + sig: _Signature = frozenset(items) with self._lock: if sig == self._sig and self._bm25 is not None: return - names, bm25 = _get_or_build(sig, tools) + names, bm25 = _get_or_build(sig, items) with self._lock: self._names, self._bm25, self._sig = names, bm25, sig diff --git a/raven/agent/tools/tool_search.py b/raven/agent/tools/tool_search.py index ee5e08a..7e50deb 100644 --- a/raven/agent/tools/tool_search.py +++ b/raven/agent/tools/tool_search.py @@ -2,17 +2,17 @@ When built-ins + plugins + MCP push the tool count past a threshold, injecting every schema into each request burns context that scales with tool count. This -module withholds most schemas and exposes three meta-tools instead: +module withholds most schemas and exposes two meta-tools instead: - - ``tool_search`` — BM25 keyword search over the hidden catalog; compact - hits (name + description, no parameters). - - ``tool_describe`` — one tool's full parameter schema (so the model knows - the argument shape before calling it). - - ``tool_call`` — invoke a cataloged tool by name; forwards through the - registry, which still validates arguments. + - ``tool_search`` — BM25 keyword search over the hidden catalog; each hit + carries name + description + parameter schema, enough to + call the tool without a second lookup. + - ``tool_call`` — invoke a cataloged tool by name; forwards through the + registry, which validates arguments and returns a + correctable error when they don't fit the schema. The tool list sent to the model never changes turn-to-turn (always the core -set + these three meta-tools), so the prompt cache stays stable across the +set + these two meta-tools), so the prompt cache stays stable across the whole session — tools sit ahead of system+messages in the cached prefix, so a changing tool list would invalidate everything after it. The cost is that cataloged tools are invoked through ``tool_call`` rather than native @@ -56,7 +56,7 @@ TOOL_CALL_NAME: str = "tool_call" # The meta-tools: always registered when the feature is on, never cataloged. -META_TOOL_NAMES: frozenset[str] = frozenset({"tool_search", "tool_describe", TOOL_CALL_NAME}) +META_TOOL_NAMES: frozenset[str] = frozenset({"tool_search", TOOL_CALL_NAME}) class ToolSearchController: @@ -98,23 +98,23 @@ def visible_names(self) -> set[str]: return self.always_visible def search(self, query: str, limit: int | None = None) -> list[dict[str, Any]]: - """Compact hits: name + description, never parameters.""" + """Hits carry name + description + parameter schema, so the model can go + straight to tool_call without a separate describe round-trip.""" names = self._index.search(query, limit or self.search_result_limit) hits = [] for name in names: tool = self._registry.get(name) if tool is None: continue - hits.append({"name": name, "description": tool.description}) + hits.append( + { + "name": name, + "description": tool.description, + "parameters": tool.parameters, + } + ) return hits - def describe(self, name: str) -> dict[str, Any] | None: - """Return the tool's full function schema (for argument shape).""" - tool = self._registry.get(name) - if tool is None or name in META_TOOL_NAMES: - return None - return tool.to_schema()["function"] - async def call(self, name: str, arguments: dict[str, Any] | None) -> str: """Invoke a cataloged tool: forward to the registry (validates args). @@ -147,10 +147,9 @@ def name(self) -> str: def description(self) -> str: return ( "Search the catalog of additional tools that are available but not " - "currently loaded. Returns matching tool names with a short " - "description (no parameters). Use tool_describe to see a tool's " - "arguments, then tool_call to invoke it. Query with task keywords, " - "e.g. 'create github issue' or '生成图片'." + "currently loaded. Returns matching tools with their description and " + "parameter schema, ready to invoke with tool_call. Query with task " + "keywords, e.g. 'create github issue' or '生成图片'." ) @property @@ -166,7 +165,7 @@ def parameters(self) -> dict[str, Any]: "type": "integer", "description": "Max number of results.", "minimum": 1, - "maximum": 25, + "maximum": 10, }, }, "required": ["query"], @@ -179,43 +178,6 @@ async def execute(self, query: str, limit: int | None = None) -> str: return json.dumps(hits, ensure_ascii=False) -class ToolDescribeTool(Tool): - """Return a cataloged tool's full parameter schema.""" - - def __init__(self, controller: ToolSearchController) -> None: - self._ctrl = controller - - @property - def name(self) -> str: - return "tool_describe" - - @property - def description(self) -> str: - return ( - "Show the full parameter schema of a tool found via tool_search, so " - "you know its arguments. Invoke the tool afterwards with tool_call." - ) - - @property - def parameters(self) -> dict[str, Any]: - return { - "type": "object", - "properties": { - "name": { - "type": "string", - "description": "Exact tool name from a tool_search result.", - }, - }, - "required": ["name"], - } - - async def execute(self, name: str) -> str: - schema = self._ctrl.describe(name) - if schema is None: - return f"Tool '{name}' not found. Use tool_search to find available tools." - return json.dumps(schema, ensure_ascii=False) - - class ToolCallTool(Tool): """Invoke a cataloged tool by name. Arguments are validated by the registry.""" @@ -230,7 +192,8 @@ def name(self) -> str: def description(self) -> str: return ( "Invoke a tool found via tool_search by name, passing its arguments. " - "Use tool_describe first to learn the argument shape." + "If the arguments don't fit the tool's schema the registry returns a " + "validation error describing the fix; adjust and call again." ) @property diff --git a/raven/config/schema.py b/raven/config/schema.py index 9063a66..9d4791e 100644 --- a/raven/config/schema.py +++ b/raven/config/schema.py @@ -504,11 +504,10 @@ class ToolSearchConfig(Base): When the live tool catalog (built-ins + plugins + MCP) grows past ``compaction_threshold``, most tool schemas are withheld from each request and reached - on demand through the ``tool_search`` / ``tool_describe`` / ``tool_call`` - meta-tools, so context cost stops scaling with tool count and the per-turn - tool list (and thus the prompt cache) stays stable. At or below the - threshold every tool is exposed directly (unchanged behavior) and the - meta-tools are omitted. + on demand through the ``tool_search`` / ``tool_call`` meta-tools, so context + cost stops scaling with tool count and the per-turn tool list (and thus the + prompt cache) stays stable. At or below the threshold every tool is exposed + directly (unchanged behavior) and the meta-tools are omitted. """ enabled: bool = False diff --git a/tests/test_agent_loop_tool_search.py b/tests/test_agent_loop_tool_search.py index ee0fff3..10529df 100644 --- a/tests/test_agent_loop_tool_search.py +++ b/tests/test_agent_loop_tool_search.py @@ -66,7 +66,7 @@ def _make_loop(workspace: Path, cfg, strategies=None) -> AgentLoop: def test_enabled_registers_meta_tools(workspace) -> None: loop = _make_loop(workspace, ToolSearchConfig(enabled=True)) - for name in ("tool_search", "tool_describe", "tool_call"): + for name in ("tool_search", "tool_call"): assert loop.tools.has(name), f"{name} should be registered" assert loop.strategies.get("tool_search") is not None @@ -83,7 +83,7 @@ def test_strategy_registered_first(workspace) -> None: def test_disabled_registers_nothing(workspace) -> None: loop = _make_loop(workspace, ToolSearchConfig(enabled=False)) - for name in ("tool_search", "tool_describe", "tool_call"): + for name in ("tool_search", "tool_call"): assert not loop.tools.has(name) assert loop.strategies.get("tool_search") is None @@ -106,5 +106,5 @@ async def test_enabled_loop_keeps_interaction_primitives_visible(workspace) -> N _, out, _ = await loop.strategies.before_llm_call([], tools, "stub") names = {t["function"]["name"] for t in out} assert {"read_file", "message", "ask_user", "spawn"} <= names, "primitives must stay visible" - assert {"tool_search", "tool_describe", "tool_call"} <= names, "meta-tools must stay visible" + assert {"tool_search", "tool_call"} <= names, "meta-tools must stay visible" assert "web_search" not in names, "cataloged domain tools are withheld above threshold" diff --git a/tests/test_tool_search.py b/tests/test_tool_search.py index 58547c9..dfd81b0 100644 --- a/tests/test_tool_search.py +++ b/tests/test_tool_search.py @@ -7,23 +7,29 @@ from raven.agent.tools.base import Tool from raven.agent.tools.registry import ToolRegistry -from raven.agent.tools.tool_index import ToolIndex +from raven.agent.tools.tool_index import ToolIndex, _schema_text from raven.agent.tools.tool_search import ( DEFAULT_ALWAYS_VISIBLE, META_TOOL_NAMES, TOOL_CALL_NAME, ToolCallTool, - ToolDescribeTool, ToolSearchController, ToolSearchStrategy, ToolSearchTool, ) +from raven.config.schema import ToolSearchConfig class _FakeTool(Tool): - def __init__(self, name: str, description: str) -> None: + def __init__( + self, + name: str, + description: str, + parameters: dict[str, Any] | None = None, + ) -> None: self._name = name self._description = description + self._parameters = parameters or {"type": "object", "properties": {}} @property def name(self) -> str: @@ -35,7 +41,7 @@ def description(self) -> str: @property def parameters(self) -> dict[str, Any]: - return {"type": "object", "properties": {}} + return self._parameters async def execute(self, **kwargs: Any) -> str: return f"ran {self._name}" @@ -80,6 +86,82 @@ def test_index_rebuilds_on_name_or_description_change() -> None: assert idx._bm25 is not first, "should rebuild when the name set changes" +def test_index_matches_parameter_schema_keywords() -> None: + # A discriminating keyword living only in the parameter schema (not the + # one-line description) should still make the tool findable. + idx = ToolIndex() + tools = [ + _FakeTool( + "create_issue", + "open a ticket", + parameters={ + "type": "object", + "properties": { + "repository": {"type": "string", "description": "target github repository"}, + }, + }, + ), + _FakeTool("send_message", "post to a channel"), + ] + idx.ensure(tools) + assert idx.search("github repository", limit=5)[0] == "create_issue" + + +def test_schema_text_extracts_names_descriptions_enums_and_nesting() -> None: + schema = { + "type": "object", + "properties": { + "channel": {"type": "string", "description": "slack channel id"}, + "mode": {"type": "string", "enum": ["fast", "thorough"]}, + "opts": { + "type": "object", + "properties": {"retries": {"type": "integer", "description": "retry count"}}, + }, + "tags": {"type": "array", "items": {"type": "string", "description": "a label"}}, + }, + } + text = _schema_text(schema) + for token in ( + "channel", + "slack channel id", + "mode", + "fast", + "thorough", + "opts", + "retries", + "retry count", + "tags", + "a label", + ): + assert token in text, f"{token!r} missing from schema text" + + +def test_schema_text_handles_non_dict_and_empty() -> None: + assert _schema_text(None) == "" + assert _schema_text([1, 2]) == "" # type: ignore[arg-type] + assert _schema_text({"type": "object"}) == "" # no properties/desc/enum + + +def test_schema_text_respects_depth_cap() -> None: + # Build nesting deeper than the cap; the deepest description must be dropped. + node: dict[str, Any] = {"type": "string", "description": "TOODEEP"} + for _ in range(10): + node = {"type": "object", "properties": {"n": node}} + assert "TOODEEP" not in _schema_text(node) + + +def test_index_rebuilds_on_parameters_change() -> None: + base = {"type": "object", "properties": {"a": {"type": "string", "description": "alpha"}}} + idx = ToolIndex() + idx.ensure([_FakeTool("t", "desc", parameters=base)]) + first = idx._bm25 + idx.ensure([_FakeTool("t", "desc", parameters=dict(base))]) # same schema content + assert idx._bm25 is first, "should not rebuild when parameters are unchanged" + changed = {"type": "object", "properties": {"a": {"type": "string", "description": "beta"}}} + idx.ensure([_FakeTool("t", "desc", parameters=changed)]) # param description changed + assert idx._bm25 is not first, "should rebuild when a parameter description changes" + + def test_index_reused_across_instances_for_same_catalog() -> None: tools = [_FakeTool("x", "xray"), _FakeTool("y", "yankee")] first = ToolIndex() @@ -100,35 +182,34 @@ def _controller(registry: ToolRegistry) -> ToolSearchController: ) -def test_controller_search_compact_has_no_parameters() -> None: +def test_controller_search_includes_parameters() -> None: + schema = { + "type": "object", + "properties": {"repo": {"type": "string", "description": "target repository"}}, + "required": ["repo"], + } reg = ToolRegistry() - reg.register(_FakeTool("create_issue", "open a github issue")) + reg.register(_FakeTool("create_issue", "open a github issue", parameters=schema)) ctrl = _controller(reg) ctrl.refresh() hits = ctrl.search("github issue") assert hits[0]["name"] == "create_issue" - assert "parameters" not in hits[0] - - -def test_controller_describe_returns_full_schema() -> None: - reg = ToolRegistry() - reg.register(_FakeTool("create_issue", "open a github issue")) - ctrl = _controller(reg) - fn = ctrl.describe("create_issue") - assert fn is not None and fn["name"] == "create_issue" and "parameters" in fn + assert hits[0]["parameters"] == schema -def test_controller_describe_unknown_returns_none() -> None: - ctrl = _controller(ToolRegistry()) - assert ctrl.describe("nope") is None +def test_meta_includes_tool_call() -> None: + assert TOOL_CALL_NAME in META_TOOL_NAMES -def test_controller_describe_rejects_meta() -> None: - assert _controller(ToolRegistry()).describe("tool_search") is None +def test_meta_no_longer_includes_describe() -> None: + assert "tool_describe" not in META_TOOL_NAMES + assert META_TOOL_NAMES == {"tool_search", TOOL_CALL_NAME} -def test_meta_includes_tool_call() -> None: - assert TOOL_CALL_NAME in META_TOOL_NAMES +def test_default_search_result_limit_is_ten() -> None: + assert ToolSearchConfig().search_result_limit == 10 + ctrl = ToolSearchController(ToolRegistry(), always_visible=set()) + assert ctrl.search_result_limit == 10 def test_default_always_visible_covers_core_and_interaction_primitives() -> None: @@ -159,7 +240,6 @@ def test_meta_tools_not_self_searchable() -> None: reg = ToolRegistry() ctrl = _controller(reg) reg.register(ToolSearchTool(ctrl)) - reg.register(ToolDescribeTool(ctrl)) reg.register(ToolCallTool(ctrl)) reg.register(_FakeTool("create_issue", "open a github issue")) ctrl.refresh() @@ -171,13 +251,16 @@ def test_meta_tools_not_self_searchable() -> None: @pytest.mark.asyncio -async def test_tool_search_tool_returns_json_hits() -> None: +async def test_tool_search_tool_returns_json_hits_with_parameters() -> None: + schema = {"type": "object", "properties": {"repo": {"type": "string"}}} reg = ToolRegistry() - reg.register(_FakeTool("create_issue", "open a github issue")) + reg.register(_FakeTool("create_issue", "open a github issue", parameters=schema)) ctrl = _controller(reg) ctrl.refresh() out = await ToolSearchTool(ctrl).execute(query="github issue") - assert json.loads(out)[0]["name"] == "create_issue" + hit = json.loads(out)[0] + assert hit["name"] == "create_issue" + assert hit["parameters"] == schema @pytest.mark.asyncio @@ -190,15 +273,6 @@ async def test_tool_search_tool_no_match_message() -> None: assert "No tools matched" in out -@pytest.mark.asyncio -async def test_tool_describe_tool_returns_schema() -> None: - reg = ToolRegistry() - reg.register(_FakeTool("create_issue", "open a github issue")) - ctrl = _controller(reg) - out = await ToolDescribeTool(ctrl).execute(name="create_issue") - assert json.loads(out)["name"] == "create_issue" - - @pytest.mark.asyncio async def test_tool_call_forwards_to_registry() -> None: reg = ToolRegistry() @@ -244,7 +318,6 @@ def _registry_with_n(n: int) -> tuple[ToolRegistry, ToolSearchController]: search_result_limit=10, ) reg.register(ToolSearchTool(ctrl)) - reg.register(ToolDescribeTool(ctrl)) reg.register(ToolCallTool(ctrl)) return reg, ctrl @@ -311,7 +384,6 @@ async def test_strategy_passthrough_when_meta_tools_absent() -> None: # expose everything rather than strand cataloged tools. reg, ctrl = _registry_with_n(40) reg.unregister("tool_search") - reg.unregister("tool_describe") reg.unregister(TOOL_CALL_NAME) strat = ToolSearchStrategy(ctrl, compaction_threshold=25) tools = reg.get_definitions()