mirror of
https://github.com/usestrix/strix.git
synced 2026-10-27 06:12:38 +00:00
feat(agents): let callers remove base tools from scan agents
This commit is contained in:
parent
aa2036f229
commit
168826a335
2 changed files with 20 additions and 0 deletions
|
|
@ -670,6 +670,7 @@ _BASE_TOOLS: tuple[Tool, ...] = (
|
|||
# ``strix.runtime.backends.register_backend``: register before the first
|
||||
# ``build_strix_agent`` call and every agent (root + children) gets them.
|
||||
_EXTRA_TOOLS: list[Tool] = []
|
||||
_REMOVED_TOOL_NAMES: set[str] = set()
|
||||
|
||||
|
||||
def _ensure_unique_tool_names(tools: Sequence[Tool]) -> None:
|
||||
|
|
@ -703,6 +704,11 @@ def register_agent_tools(*tools: Tool) -> None:
|
|||
logger.info("Registered extra agent tool: %s", getattr(tool, "name", tool))
|
||||
|
||||
|
||||
def remove_agent_tools(*names: str) -> None:
|
||||
"""Leave these base tools out of every scan agent built afterwards."""
|
||||
_REMOVED_TOOL_NAMES.update(names)
|
||||
|
||||
|
||||
def registered_agent_tools() -> tuple[Tool, ...]:
|
||||
"""Return the currently registered scan-agent tools."""
|
||||
return tuple(_EXTRA_TOOLS)
|
||||
|
|
@ -760,6 +766,7 @@ def build_strix_agent(
|
|||
tools: list[Tool] = [*_BASE_TOOLS, *agent_tools, finish_tool]
|
||||
else:
|
||||
tools = [*_BASE_TOOLS, *agent_tools, agent_finish]
|
||||
tools = [tool for tool in tools if tool.name not in _REMOVED_TOOL_NAMES]
|
||||
_ensure_unique_tool_names(tools)
|
||||
tools = [
|
||||
_with_bounded_result(_with_strictness(_with_coerced_arguments(tool), strict_tool_schemas))
|
||||
|
|
|
|||
|
|
@ -31,11 +31,15 @@ def _tool(name: str) -> FunctionTool:
|
|||
@pytest.fixture(autouse=True)
|
||||
def _reset_registry() -> object:
|
||||
saved = list(factory._EXTRA_TOOLS)
|
||||
saved_removed = set(factory._REMOVED_TOOL_NAMES)
|
||||
factory._EXTRA_TOOLS.clear()
|
||||
factory._REMOVED_TOOL_NAMES.clear()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
factory._EXTRA_TOOLS[:] = saved
|
||||
factory._REMOVED_TOOL_NAMES.clear()
|
||||
factory._REMOVED_TOOL_NAMES.update(saved_removed)
|
||||
|
||||
|
||||
def test_register_agent_tools_is_deduped() -> None:
|
||||
|
|
@ -93,6 +97,15 @@ def test_per_call_extra_tools_reject_duplicate_registered_names() -> None:
|
|||
factory.build_strix_agent(is_root=True, extra_tools=[_tool("same_name")])
|
||||
|
||||
|
||||
def test_removed_tools_are_left_out_of_root_and_child() -> None:
|
||||
factory.remove_agent_tools("create_vulnerability_report")
|
||||
|
||||
for is_root in (True, False):
|
||||
names = [t.name for t in factory.build_strix_agent(is_root=is_root).tools]
|
||||
assert "create_vulnerability_report" not in names
|
||||
assert "create_dependency_report" in names
|
||||
|
||||
|
||||
def test_instructions_override_is_used_verbatim() -> None:
|
||||
custom = "You are a scan agent. Follow the provided scope."
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue