refactor(agents): SRP main-agent tool registry, decoupled from BUILTIN_TOOLS

The main agent only exposes 4 SurfSense tools (web_search, scrape_webpage,
update_memory, create_automation) and delegates connectors/MCP/deliverables
to subagents. Yet it built those 4 by importing and iterating the 900-line,
connector-laden shared BUILTIN_TOOLS via build_tools_async.

Introduce app/agents/multi_agent_chat/main_agent/tools/registry.py owning
just those 4 factories, and switch runtime/factory.py to build_main_agent_tools.
Binding order is preserved (scrape_webpage, web_search, create_automation,
update_memory) to match prior behavior exactly.

shared/tools/registry.py BUILTIN_TOOLS is intentionally unchanged: it remains
the app-wide tool *metadata* catalog used by action_log (revert/dedup
resolvers for subagent-executed connector tools) and the /agent/tools
listing endpoint.

Verified: full unit suite green (2431 passed, 1 skipped); import-all guard ok.
This commit is contained in:
CREDO23 2026-06-04 19:01:44 +02:00
parent add9e14694
commit 482aefc32a
2 changed files with 140 additions and 3 deletions

View file

@ -28,7 +28,6 @@ from app.agents.shared.filesystem_selection import FilesystemMode, FilesystemSel
from app.agents.shared.llm_config import AgentConfig
from app.agents.shared.prompt_caching import apply_litellm_prompt_caching
from app.agents.shared.tools.invalid_tool import INVALID_TOOL_NAME, invalid_tool
from app.agents.shared.tools.registry import build_tools_async
from app.db import ChatVisibility
from app.services.connector_service import ConnectorService
from app.services.user_tool_allowlist import (
@ -42,6 +41,7 @@ from ..tools import (
MAIN_AGENT_SURFSENSE_TOOL_NAMES,
MAIN_AGENT_SURFSENSE_TOOL_NAMES_ORDERED,
)
from ..tools.registry import build_main_agent_tools
from .agent_cache import build_agent_with_cache
_perf_log = get_perf_logger()
@ -212,12 +212,14 @@ async def create_multi_agent_chat_deep_agent(
main_agent_enabled_tools = list(MAIN_AGENT_SURFSENSE_TOOL_NAMES_ORDERED)
_t0 = time.perf_counter()
tools = await build_tools_async(
# Main agent builds only its own small SurfSense toolset via the SRP
# main-agent registry; connectors/MCP/deliverables are delegated to
# subagents, so no MCP loading or connector construction happens here.
tools = build_main_agent_tools(
dependencies=dependencies,
enabled_tools=main_agent_enabled_tools,
disabled_tools=modified_disabled_tools,
additional_tools=list(additional_tools) if additional_tools else None,
include_mcp_tools=False,
)
_flags: AgentFeatureFlags = get_flags()

View file

@ -0,0 +1,135 @@
"""SRP main-agent tool registry.
The main agent exposes only a small, fixed set of SurfSense tools to its LLM;
connector integrations, MCP, and deliverables are delegated to ``task``
subagents (see :mod:`app.agents.multi_agent_chat.main_agent.tools.index`).
This module is the *building* counterpart to that name list: it owns the
factories for those few tools and nothing else. It is deliberately decoupled
from :mod:`app.agents.shared.tools.registry` (the app-wide ``BUILTIN_TOOLS``
metadata catalog, which imports every connector) so the main agent's tool
surface stays self-contained and connector-free.
The ``BUILTIN_TOOLS`` catalog still exists and is still used elsewhere for
tool *metadata* the ``/agent/tools`` listing endpoint and the action-log
revert/dedup resolvers (which must cover subagent-executed connector tools).
This registry only governs what the main agent actually builds and binds.
"""
from __future__ import annotations
from collections.abc import Callable
from typing import Any
from langchain_core.tools import BaseTool
from app.agents.shared.tools.scrape_webpage import create_scrape_webpage_tool
from app.agents.shared.tools.update_memory import (
create_update_memory_tool,
create_update_team_memory_tool,
)
from app.agents.shared.tools.web_search import create_web_search_tool
from app.db import ChatVisibility
def _build_scrape_webpage_tool(deps: dict[str, Any]) -> BaseTool:
return create_scrape_webpage_tool(firecrawl_api_key=deps.get("firecrawl_api_key"))
def _build_web_search_tool(deps: dict[str, Any]) -> BaseTool:
return create_web_search_tool(
search_space_id=deps.get("search_space_id"),
available_connectors=deps.get("available_connectors"),
)
def _build_create_automation_tool(deps: dict[str, Any]) -> BaseTool:
# Deferred import: the automation package is a sibling under ``main_agent``
# and is only needed at build time, mirroring the shared registry's
# call-time import to keep module import order robust.
from .automation import create_create_automation_tool
return create_create_automation_tool(
search_space_id=deps["search_space_id"],
user_id=deps["user_id"],
llm=deps["llm"],
)
def _build_update_memory_tool(deps: dict[str, Any]) -> BaseTool:
if deps["thread_visibility"] == ChatVisibility.SEARCH_SPACE:
return create_update_team_memory_tool(
search_space_id=deps["search_space_id"],
db_session=deps["db_session"],
llm=deps.get("llm"),
)
return create_update_memory_tool(
user_id=deps["user_id"],
db_session=deps["db_session"],
llm=deps.get("llm"),
)
# Ordered to match the historical binding order produced by the shared
# ``build_tools`` (which iterated ``BUILTIN_TOOLS`` in declaration order):
# scrape_webpage, web_search, create_automation, update_memory.
# Each entry is ``(factory, required_dependency_names)``.
_MAIN_AGENT_TOOL_FACTORIES: dict[
str, tuple[Callable[[dict[str, Any]], BaseTool], tuple[str, ...]]
] = {
"scrape_webpage": (_build_scrape_webpage_tool, ()),
"web_search": (_build_web_search_tool, ()),
"create_automation": (
_build_create_automation_tool,
("search_space_id", "user_id", "llm"),
),
"update_memory": (
_build_update_memory_tool,
("user_id", "search_space_id", "db_session", "thread_visibility", "llm"),
),
}
def build_main_agent_tools(
dependencies: dict[str, Any],
enabled_tools: list[str] | None = None,
disabled_tools: list[str] | None = None,
additional_tools: list[BaseTool] | None = None,
) -> list[BaseTool]:
"""Build the main agent's tool instances.
Args:
dependencies: Dependency bag passed to each tool factory.
enabled_tools: Explicit allow-list of tool names. When ``None``, all
main-agent tools are enabled. Names not owned by this registry are
ignored.
disabled_tools: Names to drop after the enabled set is resolved.
additional_tools: Extra tools appended verbatim (e.g. custom tools).
Returns:
Tool instances in the registry's declaration order, with any
``additional_tools`` appended.
"""
if enabled_tools is None:
names = list(_MAIN_AGENT_TOOL_FACTORIES)
else:
wanted = set(enabled_tools)
names = [n for n in _MAIN_AGENT_TOOL_FACTORIES if n in wanted]
if disabled_tools:
disabled = set(disabled_tools)
names = [n for n in names if n not in disabled]
tools: list[BaseTool] = []
for name in names:
factory, requires = _MAIN_AGENT_TOOL_FACTORIES[name]
missing = [dep for dep in requires if dep not in dependencies]
if missing:
msg = f"Tool '{name}' requires dependencies: {missing}"
raise ValueError(msg)
tools.append(factory(dependencies))
if additional_tools:
tools.extend(additional_tools)
return tools