diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/action_log/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/action_log/__init__.py new file mode 100644 index 000000000..46fa28009 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/action_log/__init__.py @@ -0,0 +1,10 @@ +"""Action-log middleware: audit row per tool call (impl + builder).""" + +from .builder import build_action_log_mw +from .middleware import ActionLogMiddleware, ToolDefinition + +__all__ = [ + "ActionLogMiddleware", + "ToolDefinition", + "build_action_log_mw", +] diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/action_log.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/action_log/builder.py similarity index 94% rename from surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/action_log.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/action_log/builder.py index 88e01bfd7..8d5220637 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/action_log.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/action_log/builder.py @@ -4,10 +4,10 @@ from __future__ import annotations import logging -from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.shared.middleware import ActionLogMiddleware - from app.agents.multi_agent_chat.shared.middleware.flags import enabled +from app.agents.shared.feature_flags import AgentFeatureFlags + +from .middleware import ActionLogMiddleware def build_action_log_mw( diff --git a/surfsense_backend/app/agents/shared/middleware/action_log.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/action_log/middleware.py similarity index 100% rename from surfsense_backend/app/agents/shared/middleware/action_log.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/action_log/middleware.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/anonymous_document/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/anonymous_document/__init__.py new file mode 100644 index 000000000..5684a592c --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/anonymous_document/__init__.py @@ -0,0 +1,9 @@ +"""Anonymous-document middleware: Redis hydration, cloud only (impl + builder).""" + +from .builder import build_anonymous_doc_mw +from .middleware import AnonymousDocumentMiddleware + +__all__ = [ + "AnonymousDocumentMiddleware", + "build_anonymous_doc_mw", +] diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/anonymous_doc.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/anonymous_document/builder.py similarity index 86% rename from surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/anonymous_doc.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/anonymous_document/builder.py index 40090320e..0958c7a84 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/anonymous_doc.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/anonymous_document/builder.py @@ -3,7 +3,8 @@ from __future__ import annotations from app.agents.shared.filesystem_selection import FilesystemMode -from app.agents.shared.middleware import AnonymousDocumentMiddleware + +from .middleware import AnonymousDocumentMiddleware def build_anonymous_doc_mw( diff --git a/surfsense_backend/app/agents/shared/middleware/anonymous_document.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/anonymous_document/middleware.py similarity index 100% rename from surfsense_backend/app/agents/shared/middleware/anonymous_document.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/anonymous_document/middleware.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/context_editing/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/context_editing/__init__.py new file mode 100644 index 000000000..0c86c8cbd --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/context_editing/__init__.py @@ -0,0 +1,15 @@ +"""Context-editing middleware: spill + clear-tool-uses passes (impl + builder).""" + +from .builder import build_context_editing_mw +from .middleware import ( + ClearToolUsesEdit, + SpillingContextEditingMiddleware, + SpillToBackendEdit, +) + +__all__ = [ + "ClearToolUsesEdit", + "SpillToBackendEdit", + "SpillingContextEditingMiddleware", + "build_context_editing_mw", +] diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/context_editing.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/context_editing/builder.py similarity index 97% rename from surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/context_editing.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/context_editing/builder.py index a907758fe..b5f8c233e 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/context_editing.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/context_editing/builder.py @@ -10,15 +10,15 @@ from langchain_core.tools import BaseTool from app.agents.multi_agent_chat.main_agent.context_prune.prune_tool_names import ( safe_exclude_tools, ) +from app.agents.multi_agent_chat.shared.middleware.flags import enabled from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.shared.middleware import ( + +from .middleware import ( ClearToolUsesEdit, SpillingContextEditingMiddleware, SpillToBackendEdit, ) -from app.agents.multi_agent_chat.shared.middleware.flags import enabled - def build_context_editing_mw( *, diff --git a/surfsense_backend/app/agents/shared/middleware/context_editing.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/context_editing/middleware.py similarity index 100% rename from surfsense_backend/app/agents/shared/middleware/context_editing.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/context_editing/middleware.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/doom_loop/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/doom_loop/__init__.py new file mode 100644 index 000000000..d0a1126a5 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/doom_loop/__init__.py @@ -0,0 +1,9 @@ +"""Doom-loop middleware: detect repeated identical tool calls (impl + builder).""" + +from .builder import build_doom_loop_mw +from .middleware import DoomLoopMiddleware + +__all__ = [ + "DoomLoopMiddleware", + "build_doom_loop_mw", +] diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/doom_loop.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/doom_loop/builder.py similarity index 87% rename from surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/doom_loop.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/doom_loop/builder.py index 6b1c50949..0afb12add 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/doom_loop.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/doom_loop/builder.py @@ -2,10 +2,10 @@ from __future__ import annotations -from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.shared.middleware import DoomLoopMiddleware - from app.agents.multi_agent_chat.shared.middleware.flags import enabled +from app.agents.shared.feature_flags import AgentFeatureFlags + +from .middleware import DoomLoopMiddleware def build_doom_loop_mw(flags: AgentFeatureFlags) -> DoomLoopMiddleware | None: diff --git a/surfsense_backend/app/agents/shared/middleware/doom_loop.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/doom_loop/middleware.py similarity index 100% rename from surfsense_backend/app/agents/shared/middleware/doom_loop.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/doom_loop/middleware.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/knowledge_tree/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/knowledge_tree/__init__.py new file mode 100644 index 000000000..f2d456b34 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/knowledge_tree/__init__.py @@ -0,0 +1,9 @@ +"""Knowledge-tree middleware: injection, cloud only (impl + builder).""" + +from .builder import build_knowledge_tree_mw +from .middleware import KnowledgeTreeMiddleware + +__all__ = [ + "KnowledgeTreeMiddleware", + "build_knowledge_tree_mw", +] diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/knowledge_tree.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/knowledge_tree/builder.py similarity index 90% rename from surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/knowledge_tree.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/knowledge_tree/builder.py index 4f6a19805..d1f77774a 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/knowledge_tree.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/knowledge_tree/builder.py @@ -5,7 +5,8 @@ from __future__ import annotations from langchain_core.language_models import BaseChatModel from app.agents.shared.filesystem_selection import FilesystemMode -from app.agents.shared.middleware import KnowledgeTreeMiddleware + +from .middleware import KnowledgeTreeMiddleware def build_knowledge_tree_mw( diff --git a/surfsense_backend/app/agents/shared/middleware/knowledge_tree.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/knowledge_tree/middleware.py similarity index 100% rename from surfsense_backend/app/agents/shared/middleware/knowledge_tree.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/knowledge_tree/middleware.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/noop_injection/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/noop_injection/__init__.py new file mode 100644 index 000000000..c4c004618 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/noop_injection/__init__.py @@ -0,0 +1,9 @@ +"""Noop-injection middleware: provider-compat _noop tool (impl + builder).""" + +from .builder import build_noop_injection_mw +from .middleware import NoopInjectionMiddleware + +__all__ = [ + "NoopInjectionMiddleware", + "build_noop_injection_mw", +] diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/noop_injection.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/noop_injection/builder.py similarity index 87% rename from surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/noop_injection.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/noop_injection/builder.py index 5388cfac7..5c62a9291 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/noop_injection.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/noop_injection/builder.py @@ -2,10 +2,10 @@ from __future__ import annotations -from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.shared.middleware import NoopInjectionMiddleware - from app.agents.multi_agent_chat.shared.middleware.flags import enabled +from app.agents.shared.feature_flags import AgentFeatureFlags + +from .middleware import NoopInjectionMiddleware def build_noop_injection_mw(flags: AgentFeatureFlags) -> NoopInjectionMiddleware | None: diff --git a/surfsense_backend/app/agents/shared/middleware/noop_injection.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/noop_injection/middleware.py similarity index 100% rename from surfsense_backend/app/agents/shared/middleware/noop_injection.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/noop_injection/middleware.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/otel_span/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/otel_span/__init__.py new file mode 100644 index 000000000..801d08962 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/otel_span/__init__.py @@ -0,0 +1,9 @@ +"""OTel-span middleware: spans on model and tool calls (impl + builder).""" + +from .builder import build_otel_mw +from .middleware import OtelSpanMiddleware + +__all__ = [ + "OtelSpanMiddleware", + "build_otel_mw", +] diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/otel.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/otel_span/builder.py similarity index 85% rename from surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/otel.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/otel_span/builder.py index f4c465ce5..09130fcb8 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/otel.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/otel_span/builder.py @@ -2,10 +2,10 @@ from __future__ import annotations -from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.shared.middleware import OtelSpanMiddleware - from app.agents.multi_agent_chat.shared.middleware.flags import enabled +from app.agents.shared.feature_flags import AgentFeatureFlags + +from .middleware import OtelSpanMiddleware def build_otel_mw(flags: AgentFeatureFlags) -> OtelSpanMiddleware | None: diff --git a/surfsense_backend/app/agents/shared/middleware/otel_span.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/otel_span/middleware.py similarity index 100% rename from surfsense_backend/app/agents/shared/middleware/otel_span.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/otel_span/middleware.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/stack.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/stack.py index 4975fcbd0..2bdf115e2 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/stack.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/stack.py @@ -20,6 +20,24 @@ from langchain_core.language_models import BaseChatModel from langchain_core.tools import BaseTool from langgraph.types import Checkpointer +from app.agents.multi_agent_chat.shared.middleware.anthropic_cache import ( + build_anthropic_cache_mw, +) +from app.agents.multi_agent_chat.shared.middleware.compaction import build_compaction_mw +from app.agents.multi_agent_chat.shared.middleware.kb_context_projection import ( + build_kb_context_projection_mw, +) +from app.agents.multi_agent_chat.shared.middleware.memory import build_memory_mw +from app.agents.multi_agent_chat.shared.middleware.patch_tool_calls import ( + build_patch_tool_calls_mw, +) +from app.agents.multi_agent_chat.shared.middleware.permissions import ( + build_permission_mw, +) +from app.agents.multi_agent_chat.shared.middleware.resilience import ( + build_resilience_middlewares, +) +from app.agents.multi_agent_chat.shared.middleware.todos import build_todos_mw from app.agents.multi_agent_chat.subagents import ( build_subagents, get_subagents_to_exclude, @@ -31,12 +49,15 @@ from app.agents.multi_agent_chat.subagents.builtins.knowledge_base.agent import from app.agents.multi_agent_chat.subagents.builtins.knowledge_base.ask_knowledge_base_tool import ( build_ask_knowledge_base_tool, ) +from app.agents.multi_agent_chat.subagents.shared.middleware.middleware_stack import ( + build_subagent_middleware_stack, +) from app.agents.shared.feature_flags import AgentFeatureFlags from app.agents.shared.filesystem_selection import FilesystemMode from app.db import ChatVisibility from .action_log import build_action_log_mw -from .anonymous_doc import build_anonymous_doc_mw +from .anonymous_document import build_anonymous_doc_mw from .busy_mutex import build_busy_mutex_mw from .checkpointed_subagent_middleware import ( SurfSenseCheckpointedSubAgentMiddleware, @@ -51,21 +72,10 @@ from .kb_persistence import build_kb_persistence_mw from .knowledge_priority import build_knowledge_priority_mw from .knowledge_tree import build_knowledge_tree_mw from .noop_injection import build_noop_injection_mw -from .otel import build_otel_mw +from .otel_span import build_otel_mw from .plugins import build_plugin_middlewares -from .repair import build_repair_mw from .skills import build_skills_mw -from app.agents.multi_agent_chat.shared.middleware.anthropic_cache import build_anthropic_cache_mw -from app.agents.multi_agent_chat.shared.middleware.compaction import build_compaction_mw -from app.agents.multi_agent_chat.shared.middleware.kb_context_projection import build_kb_context_projection_mw -from app.agents.multi_agent_chat.shared.middleware.memory import build_memory_mw -from app.agents.multi_agent_chat.shared.middleware.patch_tool_calls import build_patch_tool_calls_mw -from app.agents.multi_agent_chat.shared.middleware.permissions import build_permission_mw -from app.agents.multi_agent_chat.shared.middleware.resilience import build_resilience_middlewares -from app.agents.multi_agent_chat.shared.middleware.todos import build_todos_mw -from app.agents.multi_agent_chat.subagents.shared.middleware.middleware_stack import ( - build_subagent_middleware_stack, -) +from .tool_call_repair import build_repair_mw def build_main_agent_deepagent_middleware( diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/tool_call_repair/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/tool_call_repair/__init__.py new file mode 100644 index 000000000..1e6d93750 --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/tool_call_repair/__init__.py @@ -0,0 +1,9 @@ +"""Tool-call-repair middleware: fix miscased/unknown tool names (impl + builder).""" + +from .builder import build_repair_mw +from .middleware import ToolCallNameRepairMiddleware + +__all__ = [ + "ToolCallNameRepairMiddleware", + "build_repair_mw", +] diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/repair.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/tool_call_repair/builder.py similarity index 94% rename from surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/repair.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/tool_call_repair/builder.py index 3d1578826..30de80387 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/repair.py +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/tool_call_repair/builder.py @@ -6,10 +6,10 @@ from collections.abc import Sequence from langchain_core.tools import BaseTool -from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.shared.middleware import ToolCallNameRepairMiddleware - from app.agents.multi_agent_chat.shared.middleware.flags import enabled +from app.agents.shared.feature_flags import AgentFeatureFlags + +from .middleware import ToolCallNameRepairMiddleware # deepagents-built-in tool names the repair pass treats as known. _DEEPAGENT_BUILTIN_TOOL_NAMES: frozenset[str] = frozenset( diff --git a/surfsense_backend/app/agents/shared/middleware/tool_call_repair.py b/surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/tool_call_repair/middleware.py similarity index 100% rename from surfsense_backend/app/agents/shared/middleware/tool_call_repair.py rename to surfsense_backend/app/agents/multi_agent_chat/main_agent/middleware/tool_call_repair/middleware.py diff --git a/surfsense_backend/app/agents/shared/middleware/__init__.py b/surfsense_backend/app/agents/shared/middleware/__init__.py index e9652325c..57719f6a6 100644 --- a/surfsense_backend/app/agents/shared/middleware/__init__.py +++ b/surfsense_backend/app/agents/shared/middleware/__init__.py @@ -1,23 +1,10 @@ -"""Middleware components for the SurfSense new chat agent.""" +"""Shared middleware components for the SurfSense chat agents.""" -from app.agents.shared.middleware.action_log import ( - ActionLogMiddleware, - ToolDefinition, -) -from app.agents.shared.middleware.anonymous_document import ( - AnonymousDocumentMiddleware, -) from app.agents.shared.middleware.busy_mutex import BusyMutexMiddleware from app.agents.shared.middleware.compaction import ( SurfSenseCompactionMiddleware, create_surfsense_compaction_middleware, ) -from app.agents.shared.middleware.context_editing import ( - ClearToolUsesEdit, - SpillingContextEditingMiddleware, - SpillToBackendEdit, -) -from app.agents.shared.middleware.doom_loop import DoomLoopMiddleware from app.agents.shared.middleware.kb_persistence import ( KnowledgeBasePersistenceMiddleware, commit_staged_filesystem_state, @@ -25,39 +12,20 @@ from app.agents.shared.middleware.kb_persistence import ( from app.agents.shared.middleware.knowledge_search import ( KnowledgePriorityMiddleware, ) -from app.agents.shared.middleware.knowledge_tree import ( - KnowledgeTreeMiddleware, -) from app.agents.shared.middleware.memory_injection import ( MemoryInjectionMiddleware, ) -from app.agents.shared.middleware.noop_injection import NoopInjectionMiddleware -from app.agents.shared.middleware.otel_span import OtelSpanMiddleware from app.agents.shared.middleware.permission import PermissionMiddleware from app.agents.shared.middleware.retry_after import RetryAfterMiddleware -from app.agents.shared.middleware.tool_call_repair import ( - ToolCallNameRepairMiddleware, -) __all__ = [ - "ActionLogMiddleware", - "AnonymousDocumentMiddleware", "BusyMutexMiddleware", - "ClearToolUsesEdit", - "DoomLoopMiddleware", "KnowledgeBasePersistenceMiddleware", "KnowledgePriorityMiddleware", - "KnowledgeTreeMiddleware", "MemoryInjectionMiddleware", - "NoopInjectionMiddleware", - "OtelSpanMiddleware", "PermissionMiddleware", "RetryAfterMiddleware", - "SpillToBackendEdit", - "SpillingContextEditingMiddleware", "SurfSenseCompactionMiddleware", - "ToolCallNameRepairMiddleware", - "ToolDefinition", "commit_staged_filesystem_state", "create_surfsense_compaction_middleware", ] diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_action_log.py b/surfsense_backend/tests/unit/agents/new_chat/test_action_log.py index 3e8c7ffaf..bed1b5578 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_action_log.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_action_log.py @@ -10,9 +10,11 @@ import pytest from langchain_core.messages import ToolMessage from langchain_core.tools import tool +from app.agents.multi_agent_chat.main_agent.middleware.action_log.middleware import ( + ActionLogMiddleware, + ToolDefinition, +) from app.agents.shared.feature_flags import AgentFeatureFlags -from app.agents.shared.middleware.action_log import ActionLogMiddleware -from app.agents.shared.middleware.action_log import ToolDefinition @dataclass @@ -58,7 +60,7 @@ def _disabled_flags() -> AgentFeatureFlags: def patch_get_flags(): def _patch(flags: AgentFeatureFlags): return patch( - "app.agents.shared.middleware.action_log.get_flags", + "app.agents.multi_agent_chat.main_agent.middleware.action_log.middleware.get_flags", return_value=flags, ) @@ -360,7 +362,7 @@ class TestActionLogDispatch: patch_get_flags(_enabled_flags()), patch("app.db.shielded_async_session", side_effect=lambda: factory()), patch( - "app.agents.shared.middleware.action_log.adispatch_custom_event", + "app.agents.multi_agent_chat.main_agent.middleware.action_log.middleware.adispatch_custom_event", dispatch_mock, ), ): @@ -395,7 +397,7 @@ class TestActionLogDispatch: patch_get_flags(_enabled_flags()), patch("app.db.shielded_async_session", side_effect=_exploding_session), patch( - "app.agents.shared.middleware.action_log.adispatch_custom_event", + "app.agents.multi_agent_chat.main_agent.middleware.action_log.middleware.adispatch_custom_event", dispatch_mock, ), ): diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_context_editing.py b/surfsense_backend/tests/unit/agents/new_chat/test_context_editing.py index fbb32713e..71b91bf18 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_context_editing.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_context_editing.py @@ -7,7 +7,7 @@ from typing import Any import pytest from langchain_core.messages import AIMessage, HumanMessage, ToolMessage -from app.agents.shared.middleware.context_editing import ( +from app.agents.multi_agent_chat.main_agent.middleware.context_editing.middleware import ( SpillToBackendEdit, _build_spill_placeholder, ) diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_doom_loop.py b/surfsense_backend/tests/unit/agents/new_chat/test_doom_loop.py index 0bb338628..6f6c32271 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_doom_loop.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_doom_loop.py @@ -5,7 +5,10 @@ from __future__ import annotations import pytest from langchain_core.messages import AIMessage -from app.agents.shared.middleware.doom_loop import DoomLoopMiddleware, _signature +from app.agents.multi_agent_chat.main_agent.middleware.doom_loop.middleware import ( + DoomLoopMiddleware, + _signature, +) pytestmark = pytest.mark.unit diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_noop_injection.py b/surfsense_backend/tests/unit/agents/new_chat/test_noop_injection.py index 92c4ba1b3..bc52052d2 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_noop_injection.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_noop_injection.py @@ -5,7 +5,7 @@ from __future__ import annotations import pytest from langchain_core.messages import AIMessage, HumanMessage -from app.agents.shared.middleware.noop_injection import ( +from app.agents.multi_agent_chat.main_agent.middleware.noop_injection.middleware import ( NOOP_TOOL_NAME, NoopInjectionMiddleware, _last_ai_has_tool_calls, diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_otel_span.py b/surfsense_backend/tests/unit/agents/new_chat/test_otel_span.py index 0e924210d..8a4ed7f75 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_otel_span.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_otel_span.py @@ -8,7 +8,7 @@ from unittest.mock import MagicMock import pytest from langchain_core.messages import AIMessage, ToolMessage -from app.agents.shared.middleware.otel_span import ( +from app.agents.multi_agent_chat.main_agent.middleware.otel_span.middleware import ( OtelSpanMiddleware, _annotate_model_response, _annotate_tool_result, @@ -206,13 +206,13 @@ class TestMiddlewareIntegration: duration_calls: list[dict[str, Any]] = [] token_calls: list[dict[str, Any]] = [] monkeypatch.setattr( - "app.agents.shared.middleware.otel_span.ot_metrics.record_model_call_duration", + "app.agents.multi_agent_chat.main_agent.middleware.otel_span.middleware.ot_metrics.record_model_call_duration", lambda duration_ms, **attrs: duration_calls.append( {"duration_ms": duration_ms, **attrs} ), ) monkeypatch.setattr( - "app.agents.shared.middleware.otel_span.ot_metrics.record_model_token_usage", + "app.agents.multi_agent_chat.main_agent.middleware.otel_span.middleware.ot_metrics.record_model_token_usage", lambda **attrs: token_calls.append(attrs), ) @@ -257,11 +257,11 @@ class TestMiddlewareIntegration: errors: list[str] = [] monkeypatch.setattr( - "app.agents.shared.middleware.otel_span.ot_metrics.record_tool_call_error", + "app.agents.multi_agent_chat.main_agent.middleware.otel_span.middleware.ot_metrics.record_tool_call_error", lambda *, tool_name: errors.append(tool_name), ) monkeypatch.setattr( - "app.agents.shared.middleware.otel_span.ot_metrics.record_tool_call_duration", + "app.agents.multi_agent_chat.main_agent.middleware.otel_span.middleware.ot_metrics.record_tool_call_duration", lambda *args, **kwargs: None, ) diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_tool_call_repair.py b/surfsense_backend/tests/unit/agents/new_chat/test_tool_call_repair.py index eb2704546..042f044bb 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_tool_call_repair.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_tool_call_repair.py @@ -5,10 +5,10 @@ from __future__ import annotations import pytest 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.multi_agent_chat.main_agent.middleware.tool_call_repair.middleware import ( ToolCallNameRepairMiddleware, ) +from app.agents.multi_agent_chat.main_agent.tools.invalid_tool import INVALID_TOOL_NAME pytestmark = pytest.mark.unit diff --git a/surfsense_backend/tests/unit/middleware/test_knowledge_tree.py b/surfsense_backend/tests/unit/middleware/test_knowledge_tree.py index 87403c14a..1d0f40259 100644 --- a/surfsense_backend/tests/unit/middleware/test_knowledge_tree.py +++ b/surfsense_backend/tests/unit/middleware/test_knowledge_tree.py @@ -9,7 +9,9 @@ contract cannot silently regress. from __future__ import annotations -from app.agents.shared.middleware.knowledge_tree import KnowledgeTreeMiddleware +from app.agents.multi_agent_chat.main_agent.middleware.knowledge_tree.middleware import ( + KnowledgeTreeMiddleware, +) from app.agents.shared.path_resolver import DOCUMENTS_ROOT