mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-06-06 20:15:17 +02:00
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:
parent
add9e14694
commit
482aefc32a
2 changed files with 140 additions and 3 deletions
|
|
@ -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.llm_config import AgentConfig
|
||||||
from app.agents.shared.prompt_caching import apply_litellm_prompt_caching
|
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.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.db import ChatVisibility
|
||||||
from app.services.connector_service import ConnectorService
|
from app.services.connector_service import ConnectorService
|
||||||
from app.services.user_tool_allowlist import (
|
from app.services.user_tool_allowlist import (
|
||||||
|
|
@ -42,6 +41,7 @@ from ..tools import (
|
||||||
MAIN_AGENT_SURFSENSE_TOOL_NAMES,
|
MAIN_AGENT_SURFSENSE_TOOL_NAMES,
|
||||||
MAIN_AGENT_SURFSENSE_TOOL_NAMES_ORDERED,
|
MAIN_AGENT_SURFSENSE_TOOL_NAMES_ORDERED,
|
||||||
)
|
)
|
||||||
|
from ..tools.registry import build_main_agent_tools
|
||||||
from .agent_cache import build_agent_with_cache
|
from .agent_cache import build_agent_with_cache
|
||||||
|
|
||||||
_perf_log = get_perf_logger()
|
_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)
|
main_agent_enabled_tools = list(MAIN_AGENT_SURFSENSE_TOOL_NAMES_ORDERED)
|
||||||
|
|
||||||
_t0 = time.perf_counter()
|
_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,
|
dependencies=dependencies,
|
||||||
enabled_tools=main_agent_enabled_tools,
|
enabled_tools=main_agent_enabled_tools,
|
||||||
disabled_tools=modified_disabled_tools,
|
disabled_tools=modified_disabled_tools,
|
||||||
additional_tools=list(additional_tools) if additional_tools else None,
|
additional_tools=list(additional_tools) if additional_tools else None,
|
||||||
include_mcp_tools=False,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
_flags: AgentFeatureFlags = get_flags()
|
_flags: AgentFeatureFlags = get_flags()
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
Loading…
Add table
Add a link
Reference in a new issue