From 2db4ad479e97b3bc2838da35c3fd80297227726f Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Fri, 5 Jun 2026 10:28:56 +0200 Subject: [PATCH] refactor(agents): colocate KB-search tool with its sole consumer; fix report ImportError shared/tools/knowledge_base.py had exactly one production consumer: the report deliverable, which imported it via `from .knowledge_base import ...` -- a sibling path that did not exist, so the report KB-search path would raise ImportError at runtime. Move the module next to report.py (subagents/builtins/deliverables/tools/) which makes that relative import valid, and move its only dependency (shared/utils.py date helpers) to multi_agent_chat/shared/date_filters.py, shared between the KB tool and the knowledge_search middleware. Drop the now-unused knowledge-base re-exports from the shared/tools barrel and repoint the integration tests. import-all + error-contract stay green. --- .../shared/date_filters.py} | 0 .../deliverables}/tools/knowledge_base.py | 2 +- .../agents/shared/middleware/knowledge_search.py | 5 ++++- .../app/agents/shared/tools/__init__.py | 16 +++------------- .../integration/google_unification/conftest.py | 2 +- .../test_browse_includes_legacy_docs.py | 4 +++- 6 files changed, 12 insertions(+), 17 deletions(-) rename surfsense_backend/app/agents/{shared/utils.py => multi_agent_chat/shared/date_filters.py} (100%) rename surfsense_backend/app/agents/{shared => multi_agent_chat/subagents/builtins/deliverables}/tools/knowledge_base.py (99%) diff --git a/surfsense_backend/app/agents/shared/utils.py b/surfsense_backend/app/agents/multi_agent_chat/shared/date_filters.py similarity index 100% rename from surfsense_backend/app/agents/shared/utils.py rename to surfsense_backend/app/agents/multi_agent_chat/shared/date_filters.py diff --git a/surfsense_backend/app/agents/shared/tools/knowledge_base.py b/surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/deliverables/tools/knowledge_base.py similarity index 99% rename from surfsense_backend/app/agents/shared/tools/knowledge_base.py rename to surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/deliverables/tools/knowledge_base.py index 702b6086e..6b4b9b5a9 100644 --- a/surfsense_backend/app/agents/shared/tools/knowledge_base.py +++ b/surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/deliverables/tools/knowledge_base.py @@ -692,7 +692,7 @@ async def search_knowledge_base_raw_async( # Preserve the public signature for compatibility even if values are unused. _ = (db_session, connector_service) - from app.agents.shared.utils import resolve_date_range + from app.agents.multi_agent_chat.shared.date_filters import resolve_date_range resolved_start_date, resolved_end_date = resolve_date_range( start_date=start_date, diff --git a/surfsense_backend/app/agents/shared/middleware/knowledge_search.py b/surfsense_backend/app/agents/shared/middleware/knowledge_search.py index 9fbfc2a3c..8d20dc1a6 100644 --- a/surfsense_backend/app/agents/shared/middleware/knowledge_search.py +++ b/surfsense_backend/app/agents/shared/middleware/knowledge_search.py @@ -41,6 +41,10 @@ from litellm import token_counter from pydantic import BaseModel, Field, ValidationError from sqlalchemy import select +from app.agents.multi_agent_chat.shared.date_filters import ( + parse_date_or_datetime, + resolve_date_range, +) from app.agents.shared.feature_flags import get_flags from app.agents.shared.filesystem_selection import FilesystemMode from app.agents.shared.filesystem_state import SurfSenseFilesystemState @@ -49,7 +53,6 @@ from app.agents.shared.path_resolver import ( build_path_index, doc_to_virtual_path, ) -from app.agents.shared.utils import parse_date_or_datetime, resolve_date_range from app.db import ( NATIVE_TO_LEGACY_DOCTYPE, Chunk, diff --git a/surfsense_backend/app/agents/shared/tools/__init__.py b/surfsense_backend/app/agents/shared/tools/__init__.py index e4689c25a..21552ad98 100644 --- a/surfsense_backend/app/agents/shared/tools/__init__.py +++ b/surfsense_backend/app/agents/shared/tools/__init__.py @@ -1,24 +1,14 @@ """Cross-agent shared tools and tool metadata. Tool *implementations* live with the agents that own them (e.g. deliverable -generators under ``subagents/builtins/deliverables/tools``). This package -holds only the genuinely shared pieces: the display-metadata catalog and the -knowledge-base helpers used across agents. +generators and their knowledge-base search helper under +``subagents/builtins/deliverables/tools``). This package holds only the +genuinely shared piece: the display-metadata catalog. """ from .catalog import TOOL_CATALOG, ToolMetadata -from .knowledge_base import ( - CONNECTOR_DESCRIPTIONS, - format_documents_for_context, - search_knowledge_base_async, -) __all__ = [ - # Tool catalog (display metadata) "TOOL_CATALOG", "ToolMetadata", - # Knowledge base utilities - "CONNECTOR_DESCRIPTIONS", - "format_documents_for_context", - "search_knowledge_base_async", ] diff --git a/surfsense_backend/tests/integration/google_unification/conftest.py b/surfsense_backend/tests/integration/google_unification/conftest.py index d189afad2..df07c4841 100644 --- a/surfsense_backend/tests/integration/google_unification/conftest.py +++ b/surfsense_backend/tests/integration/google_unification/conftest.py @@ -239,7 +239,7 @@ def patched_shielded_session(async_engine, monkeypatch): yield session monkeypatch.setattr( - "app.agents.shared.tools.knowledge_base.shielded_async_session", + "app.agents.multi_agent_chat.subagents.builtins.deliverables.tools.knowledge_base.shielded_async_session", _test_shielded, ) diff --git a/surfsense_backend/tests/integration/google_unification/test_browse_includes_legacy_docs.py b/surfsense_backend/tests/integration/google_unification/test_browse_includes_legacy_docs.py index 96bf371d6..1afa72d11 100644 --- a/surfsense_backend/tests/integration/google_unification/test_browse_includes_legacy_docs.py +++ b/surfsense_backend/tests/integration/google_unification/test_browse_includes_legacy_docs.py @@ -17,7 +17,9 @@ async def test_browse_recent_documents_with_list_type_returns_both( committed_google_data, patched_shielded_session ): """_browse_recent_documents returns docs of all types when given a list.""" - from app.agents.shared.tools.knowledge_base import _browse_recent_documents + from app.agents.multi_agent_chat.subagents.builtins.deliverables.tools.knowledge_base import ( + _browse_recent_documents, + ) space_id = committed_google_data["search_space_id"]