mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-06 17:38:35 +00:00
perf(tools): cache ToolRegistry.get_definitions() between mutations
get_definitions() sorts tools on every LLM iteration for prompt cache stability. Cache the sorted result and invalidate on register/unregister so the sort only runs when the tool set actually changes. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
7ce8f247a0
commit
1304ff78cc
@ -14,14 +14,17 @@ class ToolRegistry:
|
|||||||
|
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
self._tools: dict[str, Tool] = {}
|
self._tools: dict[str, Tool] = {}
|
||||||
|
self._cached_definitions: list[dict[str, Any]] | None = None
|
||||||
|
|
||||||
def register(self, tool: Tool) -> None:
|
def register(self, tool: Tool) -> None:
|
||||||
"""Register a tool."""
|
"""Register a tool."""
|
||||||
self._tools[tool.name] = tool
|
self._tools[tool.name] = tool
|
||||||
|
self._cached_definitions = None
|
||||||
|
|
||||||
def unregister(self, name: str) -> None:
|
def unregister(self, name: str) -> None:
|
||||||
"""Unregister a tool by name."""
|
"""Unregister a tool by name."""
|
||||||
self._tools.pop(name, None)
|
self._tools.pop(name, None)
|
||||||
|
self._cached_definitions = None
|
||||||
|
|
||||||
def get(self, name: str) -> Tool | None:
|
def get(self, name: str) -> Tool | None:
|
||||||
"""Get a tool by name."""
|
"""Get a tool by name."""
|
||||||
@ -46,8 +49,12 @@ class ToolRegistry:
|
|||||||
"""Get tool definitions with stable ordering for cache-friendly prompts.
|
"""Get tool definitions with stable ordering for cache-friendly prompts.
|
||||||
|
|
||||||
Built-in tools are sorted first as a stable prefix, then MCP tools are
|
Built-in tools are sorted first as a stable prefix, then MCP tools are
|
||||||
sorted and appended.
|
sorted and appended. The result is cached until the next
|
||||||
|
register/unregister call.
|
||||||
"""
|
"""
|
||||||
|
if self._cached_definitions is not None:
|
||||||
|
return self._cached_definitions
|
||||||
|
|
||||||
definitions = [tool.to_schema() for tool in self._tools.values()]
|
definitions = [tool.to_schema() for tool in self._tools.values()]
|
||||||
builtins: list[dict[str, Any]] = []
|
builtins: list[dict[str, Any]] = []
|
||||||
mcp_tools: list[dict[str, Any]] = []
|
mcp_tools: list[dict[str, Any]] = []
|
||||||
@ -60,7 +67,8 @@ class ToolRegistry:
|
|||||||
|
|
||||||
builtins.sort(key=self._schema_name)
|
builtins.sort(key=self._schema_name)
|
||||||
mcp_tools.sort(key=self._schema_name)
|
mcp_tools.sort(key=self._schema_name)
|
||||||
return builtins + mcp_tools
|
self._cached_definitions = builtins + mcp_tools
|
||||||
|
return self._cached_definitions
|
||||||
|
|
||||||
def prepare_call(
|
def prepare_call(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@ -71,3 +71,33 @@ def test_prepare_call_other_tools_keep_generic_object_validation() -> None:
|
|||||||
assert tool is not None
|
assert tool is not None
|
||||||
assert params == ["TODO"]
|
assert params == ["TODO"]
|
||||||
assert error == "Error: Invalid parameters for tool 'grep': parameters must be an object, got list"
|
assert error == "Error: Invalid parameters for tool 'grep': parameters must be an object, got list"
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_definitions_returns_cached_result() -> None:
|
||||||
|
registry = ToolRegistry()
|
||||||
|
registry.register(_FakeTool("read_file"))
|
||||||
|
first = registry.get_definitions()
|
||||||
|
assert registry._cached_definitions is not None
|
||||||
|
second = registry.get_definitions()
|
||||||
|
assert first == second
|
||||||
|
|
||||||
|
|
||||||
|
def test_register_invalidates_cache() -> None:
|
||||||
|
registry = ToolRegistry()
|
||||||
|
registry.register(_FakeTool("read_file"))
|
||||||
|
first = registry.get_definitions()
|
||||||
|
registry.register(_FakeTool("write_file"))
|
||||||
|
second = registry.get_definitions()
|
||||||
|
assert first is not second
|
||||||
|
assert len(second) == 2
|
||||||
|
|
||||||
|
|
||||||
|
def test_unregister_invalidates_cache() -> None:
|
||||||
|
registry = ToolRegistry()
|
||||||
|
registry.register(_FakeTool("read_file"))
|
||||||
|
registry.register(_FakeTool("write_file"))
|
||||||
|
first = registry.get_definitions()
|
||||||
|
registry.unregister("write_file")
|
||||||
|
second = registry.get_definitions()
|
||||||
|
assert first is not second
|
||||||
|
assert len(second) == 1
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user