From 8ae190a11d0f5a353129435c38a67edcf046cae5 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Fri, 5 Jun 2026 12:04:31 +0200 Subject: [PATCH] refactor(agents): move MAC middleware impls out of shared kernel knowledge_search, memory_injection and scoped_model_fallback no longer belong in the cross-agent kernel (app/agents/shared/middleware): they are consumed only inside multi_agent_chat. Relocate each impl next to the builder that uses it: - knowledge_search.py -> multi_agent_chat/shared/middleware/ (genuinely shared: its _render_priority_message feeds kb_context_projection, used by both the main agent and the KB subagent) - memory_injection.py -> multi_agent_chat/shared/middleware/ (beside its memory.py builder) - scoped_model_fallback.py -> multi_agent_chat/shared/middleware/resilience/ (beside fallback.py/bundle.py) Impls moved verbatim (git rename). Builders/consumers now import the local sibling; main_agent knowledge_priority imports the new shared path; shared middleware barrel trimmed. Tests: repoint imports; convert the knowledge_search monkeypatch targets from brittle dotted-string form to object-based patching (monkeypatch.setattr on the imported module), which is robust to import ordering. No behavior change. --- .../middleware/knowledge_priority.py | 4 ++- .../middleware/kb_context_projection.py | 3 +- .../shared/middleware/knowledge_search.py | 0 .../shared/middleware/memory.py | 3 +- .../shared/middleware/memory_injection.py | 0 .../shared/middleware/resilience/bundle.py | 6 ++-- .../shared/middleware/resilience/fallback.py | 6 ++-- .../resilience}/scoped_model_fallback.py | 0 .../app/agents/shared/middleware/__init__.py | 8 ----- .../test_knowledge_search_date_filters.py | 13 ++++----- .../middleware/test_scoped_model_fallback.py | 2 +- .../unit/middleware/test_knowledge_search.py | 29 ++++++++++--------- 12 files changed, 35 insertions(+), 39 deletions(-) rename surfsense_backend/app/agents/{ => multi_agent_chat}/shared/middleware/knowledge_search.py (100%) rename surfsense_backend/app/agents/{ => multi_agent_chat}/shared/middleware/memory_injection.py (100%) rename surfsense_backend/app/agents/{shared/middleware => multi_agent_chat/shared/middleware/resilience}/scoped_model_fallback.py (100%) diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/knowledge_priority.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/knowledge_priority.py index 3d95a7d20..edfb5376b 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/knowledge_priority.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/knowledge_priority.py @@ -4,8 +4,10 @@ from __future__ import annotations from langchain_core.language_models import BaseChatModel +from app.agents.multi_agent_chat.shared.middleware.knowledge_search import ( + KnowledgePriorityMiddleware, +) from app.agents.shared.filesystem_selection import FilesystemMode -from app.agents.shared.middleware import KnowledgePriorityMiddleware from app.services.llm_service import get_planner_llm diff --git a/surfsense_backend/app/agents/multi_agent_chat/shared/middleware/kb_context_projection.py b/surfsense_backend/app/agents/multi_agent_chat/shared/middleware/kb_context_projection.py index 8b0fa1357..2502e38ca 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/shared/middleware/kb_context_projection.py +++ b/surfsense_backend/app/agents/multi_agent_chat/shared/middleware/kb_context_projection.py @@ -12,9 +12,10 @@ from langgraph.runtime import Runtime from app.agents.multi_agent_chat.shared.state.filesystem_state import ( SurfSenseFilesystemState, ) -from app.agents.shared.middleware.knowledge_search import _render_priority_message from app.utils.perf import get_perf_logger +from .knowledge_search import _render_priority_message + _perf_log = get_perf_logger() diff --git a/surfsense_backend/app/agents/shared/middleware/knowledge_search.py b/surfsense_backend/app/agents/multi_agent_chat/shared/middleware/knowledge_search.py similarity index 100% rename from surfsense_backend/app/agents/shared/middleware/knowledge_search.py rename to surfsense_backend/app/agents/multi_agent_chat/shared/middleware/knowledge_search.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/shared/middleware/memory.py b/surfsense_backend/app/agents/multi_agent_chat/shared/middleware/memory.py index 3e1df1783..1ed621ab5 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/shared/middleware/memory.py +++ b/surfsense_backend/app/agents/multi_agent_chat/shared/middleware/memory.py @@ -2,9 +2,10 @@ from __future__ import annotations -from app.agents.shared.middleware import MemoryInjectionMiddleware from app.db import ChatVisibility +from .memory_injection import MemoryInjectionMiddleware + def build_memory_mw( *, diff --git a/surfsense_backend/app/agents/shared/middleware/memory_injection.py b/surfsense_backend/app/agents/multi_agent_chat/shared/middleware/memory_injection.py similarity index 100% rename from surfsense_backend/app/agents/shared/middleware/memory_injection.py rename to surfsense_backend/app/agents/multi_agent_chat/shared/middleware/memory_injection.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/shared/middleware/resilience/bundle.py b/surfsense_backend/app/agents/multi_agent_chat/shared/middleware/resilience/bundle.py index 13354aff6..6c454a8b0 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/shared/middleware/resilience/bundle.py +++ b/surfsense_backend/app/agents/multi_agent_chat/shared/middleware/resilience/bundle.py @@ -12,13 +12,13 @@ from langchain.agents.middleware import ( from app.agents.shared.feature_flags import AgentFeatureFlags from app.agents.shared.middleware import RetryAfterMiddleware -from app.agents.shared.middleware.scoped_model_fallback import ( - ScopedModelFallbackMiddleware, -) from .fallback import build_fallback_mw from .model_call_limit import build_model_call_limit_mw from .retry import build_retry_mw +from .scoped_model_fallback import ( + ScopedModelFallbackMiddleware, +) from .tool_call_limit import build_tool_call_limit_mw diff --git a/surfsense_backend/app/agents/multi_agent_chat/shared/middleware/resilience/fallback.py b/surfsense_backend/app/agents/multi_agent_chat/shared/middleware/resilience/fallback.py index 3bd8987f5..af0649e18 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/shared/middleware/resilience/fallback.py +++ b/surfsense_backend/app/agents/multi_agent_chat/shared/middleware/resilience/fallback.py @@ -5,11 +5,11 @@ from __future__ import annotations import logging from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.shared.middleware.scoped_model_fallback import ( - ScopedModelFallbackMiddleware, -) from ..flags import enabled +from .scoped_model_fallback import ( + ScopedModelFallbackMiddleware, +) def build_fallback_mw( diff --git a/surfsense_backend/app/agents/shared/middleware/scoped_model_fallback.py b/surfsense_backend/app/agents/multi_agent_chat/shared/middleware/resilience/scoped_model_fallback.py similarity index 100% rename from surfsense_backend/app/agents/shared/middleware/scoped_model_fallback.py rename to surfsense_backend/app/agents/multi_agent_chat/shared/middleware/resilience/scoped_model_fallback.py diff --git a/surfsense_backend/app/agents/shared/middleware/__init__.py b/surfsense_backend/app/agents/shared/middleware/__init__.py index 57719f6a6..9acc5292a 100644 --- a/surfsense_backend/app/agents/shared/middleware/__init__.py +++ b/surfsense_backend/app/agents/shared/middleware/__init__.py @@ -9,20 +9,12 @@ from app.agents.shared.middleware.kb_persistence import ( KnowledgeBasePersistenceMiddleware, commit_staged_filesystem_state, ) -from app.agents.shared.middleware.knowledge_search import ( - KnowledgePriorityMiddleware, -) -from app.agents.shared.middleware.memory_injection import ( - MemoryInjectionMiddleware, -) from app.agents.shared.middleware.permission import PermissionMiddleware from app.agents.shared.middleware.retry_after import RetryAfterMiddleware __all__ = [ "BusyMutexMiddleware", "KnowledgeBasePersistenceMiddleware", - "KnowledgePriorityMiddleware", - "MemoryInjectionMiddleware", "PermissionMiddleware", "RetryAfterMiddleware", "SurfSenseCompactionMiddleware", diff --git a/surfsense_backend/tests/integration/retriever/test_knowledge_search_date_filters.py b/surfsense_backend/tests/integration/retriever/test_knowledge_search_date_filters.py index 9a911cc44..5937af196 100644 --- a/surfsense_backend/tests/integration/retriever/test_knowledge_search_date_filters.py +++ b/surfsense_backend/tests/integration/retriever/test_knowledge_search_date_filters.py @@ -8,7 +8,10 @@ from datetime import UTC, datetime, timedelta import numpy as np import pytest -from app.agents.shared.middleware.knowledge_search import search_knowledge_base +from app.agents.multi_agent_chat.shared.middleware import knowledge_search as ks +from app.agents.multi_agent_chat.shared.middleware.knowledge_search import ( + search_knowledge_base, +) from .conftest import DUMMY_EMBEDDING @@ -26,13 +29,9 @@ async def test_search_knowledge_base_applies_date_filters( async def fake_shielded_async_session(): yield db_session + monkeypatch.setattr(ks, "shielded_async_session", fake_shielded_async_session) monkeypatch.setattr( - "app.agents.shared.middleware.knowledge_search.shielded_async_session", - fake_shielded_async_session, - ) - monkeypatch.setattr( - "app.agents.shared.middleware.knowledge_search.embed_texts", - lambda texts: [np.array(DUMMY_EMBEDDING) for _ in texts], + ks, "embed_texts", lambda texts: [np.array(DUMMY_EMBEDDING) for _ in texts] ) space_id = seed_date_filtered_docs["search_space"].id diff --git a/surfsense_backend/tests/unit/agents/new_chat/middleware/test_scoped_model_fallback.py b/surfsense_backend/tests/unit/agents/new_chat/middleware/test_scoped_model_fallback.py index 0bc614fb4..8ba4d7c91 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/middleware/test_scoped_model_fallback.py +++ b/surfsense_backend/tests/unit/agents/new_chat/middleware/test_scoped_model_fallback.py @@ -87,7 +87,7 @@ class RateLimitError(Exception): def _build_agent(primary: BaseChatModel, fallback: BaseChatModel): from langchain.agents import create_agent - from app.agents.shared.middleware.scoped_model_fallback import ( + from app.agents.multi_agent_chat.shared.middleware.resilience.scoped_model_fallback import ( ScopedModelFallbackMiddleware, ) diff --git a/surfsense_backend/tests/unit/middleware/test_knowledge_search.py b/surfsense_backend/tests/unit/middleware/test_knowledge_search.py index 00304794b..79e2d9050 100644 --- a/surfsense_backend/tests/unit/middleware/test_knowledge_search.py +++ b/surfsense_backend/tests/unit/middleware/test_knowledge_search.py @@ -5,10 +5,11 @@ import json import pytest from langchain_core.messages import AIMessage, HumanMessage +from app.agents.multi_agent_chat.shared.middleware import knowledge_search as ks from app.agents.multi_agent_chat.shared.middleware.filesystem.backends.document_xml import ( build_document_xml as _build_document_xml, ) -from app.agents.shared.middleware.knowledge_search import ( +from app.agents.multi_agent_chat.shared.middleware.knowledge_search import ( KBSearchPlan, KnowledgePriorityMiddleware, _normalize_optional_date_range, @@ -260,7 +261,7 @@ class TestKnowledgePriorityMiddlewarePlanner: return [] monkeypatch.setattr( - "app.agents.shared.middleware.knowledge_search.search_knowledge_base", + ks, "search_knowledge_base", fake_search_knowledge_base, ) @@ -303,7 +304,7 @@ class TestKnowledgePriorityMiddlewarePlanner: return [] monkeypatch.setattr( - "app.agents.shared.middleware.knowledge_search.search_knowledge_base", + ks, "search_knowledge_base", fake_search_knowledge_base, ) @@ -332,7 +333,7 @@ class TestKnowledgePriorityMiddlewarePlanner: return [] monkeypatch.setattr( - "app.agents.shared.middleware.knowledge_search.search_knowledge_base", + ks, "search_knowledge_base", fake_search_knowledge_base, ) @@ -377,11 +378,11 @@ class TestKnowledgePriorityMiddlewarePlanner: return [] monkeypatch.setattr( - "app.agents.shared.middleware.knowledge_search.browse_recent_documents", + ks, "browse_recent_documents", fake_browse_recent_documents, ) monkeypatch.setattr( - "app.agents.shared.middleware.knowledge_search.search_knowledge_base", + ks, "search_knowledge_base", fake_search_knowledge_base, ) @@ -424,11 +425,11 @@ class TestKnowledgePriorityMiddlewarePlanner: return [] monkeypatch.setattr( - "app.agents.shared.middleware.knowledge_search.browse_recent_documents", + ks, "browse_recent_documents", fake_browse_recent_documents, ) monkeypatch.setattr( - "app.agents.shared.middleware.knowledge_search.search_knowledge_base", + ks, "search_knowledge_base", fake_search_knowledge_base, ) @@ -551,11 +552,11 @@ class TestKnowledgePriorityMentionDrain: return [] monkeypatch.setattr( - "app.agents.shared.middleware.knowledge_search.fetch_mentioned_documents", + ks, "fetch_mentioned_documents", fake_fetch_mentioned_documents, ) monkeypatch.setattr( - "app.agents.shared.middleware.knowledge_search.search_knowledge_base", + ks, "search_knowledge_base", fake_search_knowledge_base, ) @@ -599,11 +600,11 @@ class TestKnowledgePriorityMentionDrain: return [] monkeypatch.setattr( - "app.agents.shared.middleware.knowledge_search.fetch_mentioned_documents", + ks, "fetch_mentioned_documents", fake_fetch_mentioned_documents, ) monkeypatch.setattr( - "app.agents.shared.middleware.knowledge_search.search_knowledge_base", + ks, "search_knowledge_base", fake_search_knowledge_base, ) @@ -644,11 +645,11 @@ class TestKnowledgePriorityMentionDrain: return [] monkeypatch.setattr( - "app.agents.shared.middleware.knowledge_search.fetch_mentioned_documents", + ks, "fetch_mentioned_documents", fake_fetch_mentioned_documents, ) monkeypatch.setattr( - "app.agents.shared.middleware.knowledge_search.search_knowledge_base", + ks, "search_knowledge_base", fake_search_knowledge_base, )