mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-06-06 20:15:17 +02:00
refactor(agents): colocate main_agent-only kernel into main_agent/
Move modules out of agents/shared/ that are consumed by a single package (main_agent), placing each next to its only consumer instead of in a "shared" grab-bag: - agent_cache.py -> main_agent/runtime/agent_cache_store.py - connector_searchable_types.py -> main_agent/runtime/ - plugin_loader.py + plugins/ -> main_agent/plugins/ - skills/ + skills_backends.py -> main_agent/skills/ - tools/invalid_tool.py -> main_agent/tools/ Drop the skills_backends re-export from the shared middleware barrel and repoint all consumers + tests. No behavior change; import-all, error-contract, and the moved tests stay green.
This commit is contained in:
parent
c51aca6ccc
commit
a7d7155039
24 changed files with 33 additions and 46 deletions
|
|
@ -7,15 +7,15 @@ from typing import Any
|
||||||
|
|
||||||
from langchain_core.language_models import BaseChatModel
|
from langchain_core.language_models import BaseChatModel
|
||||||
|
|
||||||
|
from app.agents.multi_agent_chat.shared.middleware.flags import enabled
|
||||||
from app.agents.shared.feature_flags import AgentFeatureFlags
|
from app.agents.shared.feature_flags import AgentFeatureFlags
|
||||||
from app.agents.shared.plugin_loader import (
|
from app.db import ChatVisibility
|
||||||
|
|
||||||
|
from ..plugins.loader import (
|
||||||
PluginContext,
|
PluginContext,
|
||||||
load_allowed_plugin_names_from_env,
|
load_allowed_plugin_names_from_env,
|
||||||
load_plugin_middlewares,
|
load_plugin_middlewares,
|
||||||
)
|
)
|
||||||
from app.db import ChatVisibility
|
|
||||||
|
|
||||||
from app.agents.multi_agent_chat.shared.middleware.flags import enabled
|
|
||||||
|
|
||||||
|
|
||||||
def build_plugin_middlewares(
|
def build_plugin_middlewares(
|
||||||
|
|
|
||||||
|
|
@ -6,14 +6,11 @@ import logging
|
||||||
|
|
||||||
from deepagents.middleware.skills import SkillsMiddleware
|
from deepagents.middleware.skills import SkillsMiddleware
|
||||||
|
|
||||||
|
from app.agents.multi_agent_chat.shared.middleware.flags import enabled
|
||||||
from app.agents.shared.feature_flags import AgentFeatureFlags
|
from app.agents.shared.feature_flags import AgentFeatureFlags
|
||||||
from app.agents.shared.filesystem_selection import FilesystemMode
|
from app.agents.shared.filesystem_selection import FilesystemMode
|
||||||
from app.agents.shared.middleware import (
|
|
||||||
build_skills_backend_factory,
|
|
||||||
default_skills_sources,
|
|
||||||
)
|
|
||||||
|
|
||||||
from app.agents.multi_agent_chat.shared.middleware.flags import enabled
|
from ..skills.backends import build_skills_backend_factory, default_skills_sources
|
||||||
|
|
||||||
|
|
||||||
def build_skills_mw(
|
def build_skills_mw(
|
||||||
|
|
|
||||||
|
|
@ -17,7 +17,7 @@ Wire-up in ``pyproject.toml`` (illustrative; the in-repo plugin doesn't
|
||||||
need this -- it's already on the import path)::
|
need this -- it's already on the import path)::
|
||||||
|
|
||||||
[project.entry-points."surfsense.plugins"]
|
[project.entry-points."surfsense.plugins"]
|
||||||
year_substituter = "app.agents.shared.plugins.year_substituter:make_middleware"
|
year_substituter = "app.agents.multi_agent_chat.main_agent.plugins.year_substituter:make_middleware"
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
@ -34,7 +34,7 @@ if TYPE_CHECKING: # pragma: no cover - type-only
|
||||||
from langchain_core.messages import ToolMessage
|
from langchain_core.messages import ToolMessage
|
||||||
from langgraph.types import Command
|
from langgraph.types import Command
|
||||||
|
|
||||||
from app.agents.shared.plugin_loader import PluginContext
|
from .loader import PluginContext
|
||||||
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
@ -10,18 +10,18 @@ from langchain_core.language_models import BaseChatModel
|
||||||
from langchain_core.tools import BaseTool
|
from langchain_core.tools import BaseTool
|
||||||
from langgraph.types import Checkpointer
|
from langgraph.types import Checkpointer
|
||||||
|
|
||||||
from app.agents.shared.agent_cache import (
|
from app.agents.shared.feature_flags import AgentFeatureFlags
|
||||||
|
from app.agents.shared.filesystem_selection import FilesystemMode
|
||||||
|
from app.db import ChatVisibility
|
||||||
|
|
||||||
|
from ..graph.compile_graph_sync import build_compiled_agent_graph_sync
|
||||||
|
from .agent_cache_store import (
|
||||||
flags_signature,
|
flags_signature,
|
||||||
get_cache,
|
get_cache,
|
||||||
stable_hash,
|
stable_hash,
|
||||||
system_prompt_hash,
|
system_prompt_hash,
|
||||||
tools_signature,
|
tools_signature,
|
||||||
)
|
)
|
||||||
from app.agents.shared.feature_flags import AgentFeatureFlags
|
|
||||||
from app.agents.shared.filesystem_selection import FilesystemMode
|
|
||||||
from app.db import ChatVisibility
|
|
||||||
|
|
||||||
from ..graph.compile_graph_sync import build_compiled_agent_graph_sync
|
|
||||||
|
|
||||||
|
|
||||||
def mcp_signature(mcp_tools_by_agent: dict[str, list[BaseTool]]) -> str:
|
def mcp_signature(mcp_tools_by_agent: dict[str, list[BaseTool]]) -> str:
|
||||||
|
|
|
||||||
|
|
@ -19,15 +19,11 @@ from app.agents.multi_agent_chat.subagents import (
|
||||||
from app.agents.multi_agent_chat.subagents.mcp_tools.index import (
|
from app.agents.multi_agent_chat.subagents.mcp_tools.index import (
|
||||||
load_mcp_tools_by_connector,
|
load_mcp_tools_by_connector,
|
||||||
)
|
)
|
||||||
from app.agents.shared.connector_searchable_types import (
|
|
||||||
map_connectors_to_searchable_types,
|
|
||||||
)
|
|
||||||
from app.agents.shared.feature_flags import AgentFeatureFlags, get_flags
|
from app.agents.shared.feature_flags import AgentFeatureFlags, get_flags
|
||||||
from app.agents.shared.filesystem_backends import build_backend_resolver
|
from app.agents.shared.filesystem_backends import build_backend_resolver
|
||||||
from app.agents.shared.filesystem_selection import FilesystemMode, FilesystemSelection
|
from app.agents.shared.filesystem_selection import FilesystemMode, FilesystemSelection
|
||||||
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.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 (
|
||||||
|
|
@ -41,8 +37,10 @@ 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.invalid_tool import INVALID_TOOL_NAME, invalid_tool
|
||||||
from ..tools.registry import build_main_agent_tools
|
from ..tools.registry import build_main_agent_tools
|
||||||
from .agent_cache import build_agent_with_cache
|
from .agent_cache import build_agent_with_cache
|
||||||
|
from .connector_searchable_types import map_connectors_to_searchable_types
|
||||||
|
|
||||||
_perf_log = get_perf_logger()
|
_perf_log = get_perf_logger()
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -59,10 +59,10 @@ _MAX_SKILL_FILE_SIZE = 10 * 1024 * 1024
|
||||||
def _default_builtin_root() -> Path:
|
def _default_builtin_root() -> Path:
|
||||||
"""Return the absolute path to the bundled builtin skills directory.
|
"""Return the absolute path to the bundled builtin skills directory.
|
||||||
|
|
||||||
Located at ``app/agents/shared/skills/builtin/`` relative to this module
|
Located at ``builtin/`` next to this module (this module lives at
|
||||||
(this module lives at ``app/agents/shared/middleware/skills_backends.py``).
|
``app/agents/multi_agent_chat/main_agent/skills/backends.py``).
|
||||||
"""
|
"""
|
||||||
return (Path(__file__).resolve().parent.parent / "skills" / "builtin").resolve()
|
return (Path(__file__).resolve().parent / "builtin").resolve()
|
||||||
|
|
||||||
|
|
||||||
class BuiltinSkillsBackend(BackendProtocol):
|
class BuiltinSkillsBackend(BackendProtocol):
|
||||||
|
|
@ -45,12 +45,6 @@ from app.agents.shared.middleware.noop_injection import NoopInjectionMiddleware
|
||||||
from app.agents.shared.middleware.otel_span import OtelSpanMiddleware
|
from app.agents.shared.middleware.otel_span import OtelSpanMiddleware
|
||||||
from app.agents.shared.middleware.permission import PermissionMiddleware
|
from app.agents.shared.middleware.permission import PermissionMiddleware
|
||||||
from app.agents.shared.middleware.retry_after import RetryAfterMiddleware
|
from app.agents.shared.middleware.retry_after import RetryAfterMiddleware
|
||||||
from app.agents.shared.middleware.skills_backends import (
|
|
||||||
BuiltinSkillsBackend,
|
|
||||||
SearchSpaceSkillsBackend,
|
|
||||||
build_skills_backend_factory,
|
|
||||||
default_skills_sources,
|
|
||||||
)
|
|
||||||
from app.agents.shared.middleware.tool_call_repair import (
|
from app.agents.shared.middleware.tool_call_repair import (
|
||||||
ToolCallNameRepairMiddleware,
|
ToolCallNameRepairMiddleware,
|
||||||
)
|
)
|
||||||
|
|
@ -58,7 +52,6 @@ from app.agents.shared.middleware.tool_call_repair import (
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"ActionLogMiddleware",
|
"ActionLogMiddleware",
|
||||||
"AnonymousDocumentMiddleware",
|
"AnonymousDocumentMiddleware",
|
||||||
"BuiltinSkillsBackend",
|
|
||||||
"BusyMutexMiddleware",
|
"BusyMutexMiddleware",
|
||||||
"ClearToolUsesEdit",
|
"ClearToolUsesEdit",
|
||||||
"DedupHITLToolCallsMiddleware",
|
"DedupHITLToolCallsMiddleware",
|
||||||
|
|
@ -74,14 +67,11 @@ __all__ = [
|
||||||
"OtelSpanMiddleware",
|
"OtelSpanMiddleware",
|
||||||
"PermissionMiddleware",
|
"PermissionMiddleware",
|
||||||
"RetryAfterMiddleware",
|
"RetryAfterMiddleware",
|
||||||
"SearchSpaceSkillsBackend",
|
|
||||||
"SpillToBackendEdit",
|
"SpillToBackendEdit",
|
||||||
"SpillingContextEditingMiddleware",
|
"SpillingContextEditingMiddleware",
|
||||||
"SurfSenseCompactionMiddleware",
|
"SurfSenseCompactionMiddleware",
|
||||||
"ToolCallNameRepairMiddleware",
|
"ToolCallNameRepairMiddleware",
|
||||||
"ToolDefinition",
|
"ToolDefinition",
|
||||||
"build_skills_backend_factory",
|
|
||||||
"commit_staged_filesystem_state",
|
"commit_staged_filesystem_state",
|
||||||
"create_surfsense_compaction_middleware",
|
"create_surfsense_compaction_middleware",
|
||||||
"default_skills_sources",
|
|
||||||
]
|
]
|
||||||
|
|
|
||||||
|
|
@ -120,7 +120,9 @@ class ToolCallNameRepairMiddleware(
|
||||||
# Stage 2 — invalid fallback
|
# Stage 2 — invalid fallback
|
||||||
# Local import keeps the middleware module import-light and avoids any
|
# Local import keeps the middleware module import-light and avoids any
|
||||||
# tools <-> middleware import-order coupling at module scope.
|
# tools <-> middleware import-order coupling at module scope.
|
||||||
from app.agents.shared.tools.invalid_tool import INVALID_TOOL_NAME
|
from app.agents.multi_agent_chat.main_agent.tools.invalid_tool import (
|
||||||
|
INVALID_TOOL_NAME,
|
||||||
|
)
|
||||||
|
|
||||||
if INVALID_TOOL_NAME in registered:
|
if INVALID_TOOL_NAME in registered:
|
||||||
original_args = call.get("args") or {}
|
original_args = call.get("args") or {}
|
||||||
|
|
|
||||||
|
|
@ -16,7 +16,7 @@ from dataclasses import dataclass
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from app.agents.shared.agent_cache import (
|
from app.agents.multi_agent_chat.main_agent.runtime.agent_cache_store import (
|
||||||
flags_signature,
|
flags_signature,
|
||||||
reload_for_tests,
|
reload_for_tests,
|
||||||
stable_hash,
|
stable_hash,
|
||||||
|
|
|
||||||
|
|
@ -6,13 +6,13 @@ from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
from langchain.agents.middleware import AgentMiddleware
|
from langchain.agents.middleware import AgentMiddleware
|
||||||
|
|
||||||
from app.agents.shared.plugin_loader import (
|
from app.agents.multi_agent_chat.main_agent.plugins.loader import (
|
||||||
PLUGIN_ENTRY_POINT_GROUP,
|
PLUGIN_ENTRY_POINT_GROUP,
|
||||||
PluginContext,
|
PluginContext,
|
||||||
load_allowed_plugin_names_from_env,
|
load_allowed_plugin_names_from_env,
|
||||||
load_plugin_middlewares,
|
load_plugin_middlewares,
|
||||||
)
|
)
|
||||||
from app.agents.shared.plugins.year_substituter import (
|
from app.agents.multi_agent_chat.main_agent.plugins.year_substituter import (
|
||||||
_YearSubstituterMiddleware,
|
_YearSubstituterMiddleware,
|
||||||
make_middleware as year_substituter_factory,
|
make_middleware as year_substituter_factory,
|
||||||
)
|
)
|
||||||
|
|
@ -66,7 +66,7 @@ class TestPluginLoaderBasics:
|
||||||
|
|
||||||
ep = _FakeEntryPoint("dangerous_plugin", factory)
|
ep = _FakeEntryPoint("dangerous_plugin", factory)
|
||||||
with patch(
|
with patch(
|
||||||
"app.agents.shared.plugin_loader.entry_points",
|
"app.agents.multi_agent_chat.main_agent.plugins.loader.entry_points",
|
||||||
return_value=[ep],
|
return_value=[ep],
|
||||||
):
|
):
|
||||||
result = load_plugin_middlewares(
|
result = load_plugin_middlewares(
|
||||||
|
|
@ -78,7 +78,7 @@ class TestPluginLoaderBasics:
|
||||||
def test_loads_allowlisted_plugin(self) -> None:
|
def test_loads_allowlisted_plugin(self) -> None:
|
||||||
ep = _FakeEntryPoint("year_substituter", year_substituter_factory)
|
ep = _FakeEntryPoint("year_substituter", year_substituter_factory)
|
||||||
with patch(
|
with patch(
|
||||||
"app.agents.shared.plugin_loader.entry_points",
|
"app.agents.multi_agent_chat.main_agent.plugins.loader.entry_points",
|
||||||
return_value=[ep],
|
return_value=[ep],
|
||||||
):
|
):
|
||||||
result = load_plugin_middlewares(
|
result = load_plugin_middlewares(
|
||||||
|
|
@ -95,7 +95,7 @@ class TestPluginLoaderIsolation:
|
||||||
|
|
||||||
ep = _FakeEntryPoint("buggy", crashing_factory)
|
ep = _FakeEntryPoint("buggy", crashing_factory)
|
||||||
with patch(
|
with patch(
|
||||||
"app.agents.shared.plugin_loader.entry_points",
|
"app.agents.multi_agent_chat.main_agent.plugins.loader.entry_points",
|
||||||
return_value=[ep],
|
return_value=[ep],
|
||||||
):
|
):
|
||||||
result = load_plugin_middlewares(_ctx(), allowed_plugin_names={"buggy"})
|
result = load_plugin_middlewares(_ctx(), allowed_plugin_names={"buggy"})
|
||||||
|
|
@ -107,7 +107,7 @@ class TestPluginLoaderIsolation:
|
||||||
|
|
||||||
ep = _FakeEntryPoint("liar", bad_factory)
|
ep = _FakeEntryPoint("liar", bad_factory)
|
||||||
with patch(
|
with patch(
|
||||||
"app.agents.shared.plugin_loader.entry_points",
|
"app.agents.multi_agent_chat.main_agent.plugins.loader.entry_points",
|
||||||
return_value=[ep],
|
return_value=[ep],
|
||||||
):
|
):
|
||||||
result = load_plugin_middlewares(_ctx(), allowed_plugin_names={"liar"})
|
result = load_plugin_middlewares(_ctx(), allowed_plugin_names={"liar"})
|
||||||
|
|
@ -121,7 +121,7 @@ class TestPluginLoaderIsolation:
|
||||||
raise ImportError("cannot import")
|
raise ImportError("cannot import")
|
||||||
|
|
||||||
with patch(
|
with patch(
|
||||||
"app.agents.shared.plugin_loader.entry_points",
|
"app.agents.multi_agent_chat.main_agent.plugins.loader.entry_points",
|
||||||
return_value=[_BrokenEP()],
|
return_value=[_BrokenEP()],
|
||||||
):
|
):
|
||||||
result = load_plugin_middlewares(_ctx(), allowed_plugin_names={"broken"})
|
result = load_plugin_middlewares(_ctx(), allowed_plugin_names={"broken"})
|
||||||
|
|
@ -137,7 +137,7 @@ class TestPluginLoaderIsolation:
|
||||||
_FakeEntryPoint("crashing", crashing_factory),
|
_FakeEntryPoint("crashing", crashing_factory),
|
||||||
_FakeEntryPoint("ok", year_substituter_factory),
|
_FakeEntryPoint("ok", year_substituter_factory),
|
||||||
]
|
]
|
||||||
with patch("app.agents.shared.plugin_loader.entry_points", return_value=eps):
|
with patch("app.agents.multi_agent_chat.main_agent.plugins.loader.entry_points", return_value=eps):
|
||||||
result = load_plugin_middlewares(
|
result = load_plugin_middlewares(
|
||||||
_ctx(), allowed_plugin_names={"crashing", "ok"}
|
_ctx(), allowed_plugin_names={"crashing", "ok"}
|
||||||
)
|
)
|
||||||
|
|
|
||||||
|
|
@ -7,7 +7,7 @@ from pathlib import Path
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from app.agents.shared.middleware.skills_backends import (
|
from app.agents.multi_agent_chat.main_agent.skills.backends import (
|
||||||
SKILLS_BUILTIN_PREFIX,
|
SKILLS_BUILTIN_PREFIX,
|
||||||
SKILLS_SPACE_PREFIX,
|
SKILLS_SPACE_PREFIX,
|
||||||
BuiltinSkillsBackend,
|
BuiltinSkillsBackend,
|
||||||
|
|
|
||||||
|
|
@ -5,10 +5,10 @@ from __future__ import annotations
|
||||||
import pytest
|
import pytest
|
||||||
from langchain_core.messages import AIMessage
|
from langchain_core.messages import AIMessage
|
||||||
|
|
||||||
|
from app.agents.multi_agent_chat.main_agent.tools.invalid_tool import INVALID_TOOL_NAME
|
||||||
from app.agents.shared.middleware.tool_call_repair import (
|
from app.agents.shared.middleware.tool_call_repair import (
|
||||||
ToolCallNameRepairMiddleware,
|
ToolCallNameRepairMiddleware,
|
||||||
)
|
)
|
||||||
from app.agents.shared.tools.invalid_tool import INVALID_TOOL_NAME
|
|
||||||
|
|
||||||
pytestmark = pytest.mark.unit
|
pytestmark = pytest.mark.unit
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue