diff --git a/surfsense_backend/app/agents/multi_agent_chat/core/bindings/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/core/bindings/__init__.py index d6a826113..c15375e47 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/core/bindings/__init__.py +++ b/surfsense_backend/app/agents/multi_agent_chat/core/bindings/__init__.py @@ -1,4 +1,4 @@ -"""Search-space / DB kwargs shared by ``new_chat`` tool factories (distinct from ``expert_agent.connectors`` integrations).""" +"""Search-space / DB kwargs shared by main-chat tool factories (distinct from ``expert_agent.connectors`` integrations).""" from app.agents.multi_agent_chat.core.bindings.binding import connector_binding diff --git a/surfsense_backend/app/agents/multi_agent_chat/core/bindings/binding.py b/surfsense_backend/app/agents/multi_agent_chat/core/bindings/binding.py index 25e6a03fd..da82e3b3c 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/core/bindings/binding.py +++ b/surfsense_backend/app/agents/multi_agent_chat/core/bindings/binding.py @@ -1,4 +1,4 @@ -"""Shared kwargs dict for ``new_chat`` tool factories (DB session + search space + user).""" +"""Shared kwargs dict for main-chat tool factories (DB session + search space + user).""" from __future__ import annotations diff --git a/surfsense_backend/app/agents/multi_agent_chat/core/mcp_partition.py b/surfsense_backend/app/agents/multi_agent_chat/core/mcp_partition.py index 608d16988..a1ee6fdb6 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/core/mcp_partition.py +++ b/surfsense_backend/app/agents/multi_agent_chat/core/mcp_partition.py @@ -1,4 +1,4 @@ -"""Partition MCP tools onto multi-agent expert routes without modifying ``new_chat``. +"""Partition MCP tools onto multi-agent expert routes (read-only; does not change the MCP loader). Uses the same connector discovery shape as ``load_mcp_tools`` (copied query below). Tools come from ``app.agents.new_chat.tools.mcp_tool.load_mcp_tools``; routing uses metadata already set there: @@ -61,7 +61,7 @@ async def fetch_mcp_connector_metadata_maps( ) -> tuple[dict[int, str], dict[str, str]]: """Read-only copy of connector discovery used alongside ``load_mcp_tools``. - Same filter as ``new_chat.tools.mcp_tool.load_mcp_tools`` (connectors with ``server_config``). + Same filter as :func:`app.agents.new_chat.tools.mcp_tool.load_mcp_tools` (connectors with ``server_config``). """ result = await session.execute( select(SearchSourceConnector).filter( @@ -90,7 +90,7 @@ def partition_mcp_tools_by_expert_route( ) -> dict[str, list[BaseTool]]: """Bucket MCP tools by expert route key. Supervisor never receives raw MCP tools. - Same inclusion rule as ``new_chat.tools.registry.build_tools_async``: all tools returned by + Same inclusion rule as :func:`app.agents.new_chat.tools.registry.build_tools_async`: all tools returned by ``load_mcp_tools`` are partitioned — connector availability for **registry** builtins is handled via ``get_connector_gated_tools`` / routing gates; MCP tools are not pre-filtered by inventory here. """ diff --git a/surfsense_backend/app/agents/multi_agent_chat/core/registry/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/core/registry/__init__.py index 0655115c0..cfd8a5d62 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/core/registry/__init__.py +++ b/surfsense_backend/app/agents/multi_agent_chat/core/registry/__init__.py @@ -1,4 +1,4 @@ -"""``new_chat`` tool registry grouping + dependency bundles for domain slices.""" +"""Main chat tool registry grouping + dependency bundles for domain slices.""" from app.agents.multi_agent_chat.core.registry.categories import ( REGISTRY_ROUTING_CATEGORY_KEYS, diff --git a/surfsense_backend/app/agents/multi_agent_chat/core/registry/dependencies.py b/surfsense_backend/app/agents/multi_agent_chat/core/registry/dependencies.py index 68125c208..24fa6b19c 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/core/registry/dependencies.py +++ b/surfsense_backend/app/agents/multi_agent_chat/core/registry/dependencies.py @@ -1,4 +1,4 @@ -"""Dependency dict for :func:`app.agents.new_chat.tools.registry.build_tools` in multi-agent graphs.""" +"""Dependency dict for :func:`app.agents.new_chat.tools.registry.build_tools` on expert subgraphs.""" from __future__ import annotations diff --git a/surfsense_backend/app/agents/multi_agent_chat/core/registry/subset.py b/surfsense_backend/app/agents/multi_agent_chat/core/registry/subset.py index 027a8af8f..95db1b64c 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/core/registry/subset.py +++ b/surfsense_backend/app/agents/multi_agent_chat/core/registry/subset.py @@ -1,4 +1,4 @@ -"""Build :mod:`new_chat` registry tool subsets for multi-agent domain slices.""" +"""Build registry tool subsets (``app.agents.new_chat.tools.registry``) for multi-agent domain slices.""" from __future__ import annotations diff --git a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/builtins/deliverables/slice_tools.py b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/builtins/deliverables/slice_tools.py index 2c8e80a55..42241bda5 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/builtins/deliverables/slice_tools.py +++ b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/builtins/deliverables/slice_tools.py @@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c def build_deliverables_tools(dependencies: dict[str, Any]) -> list[BaseTool]: - """Tools from ``new_chat`` registry: ``deliverables`` category.""" + """Registry-backed tools for the ``deliverables`` category.""" return build_registry_tools_for_category(dependencies, "deliverables") diff --git a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/builtins/memory/slice_tools.py b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/builtins/memory/slice_tools.py index e2a482ff0..7f4d2d29a 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/builtins/memory/slice_tools.py +++ b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/builtins/memory/slice_tools.py @@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c def build_memory_tools(dependencies: dict[str, Any]) -> list[BaseTool]: - """Tools from ``new_chat`` registry: ``memory`` category.""" + """Registry-backed tools for the ``memory`` category.""" return build_registry_tools_for_category(dependencies, "memory") diff --git a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/builtins/research/slice_tools.py b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/builtins/research/slice_tools.py index 4018c5a18..85a2a9dd9 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/builtins/research/slice_tools.py +++ b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/builtins/research/slice_tools.py @@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c def build_research_tools(dependencies: dict[str, Any]) -> list[BaseTool]: - """Tools from ``new_chat`` registry: ``research`` category.""" + """Registry-backed tools for the ``research`` category.""" return build_registry_tools_for_category(dependencies, "research") diff --git a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/calendar/slice_tools.py b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/calendar/slice_tools.py index 49e316c01..e2f2b404a 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/calendar/slice_tools.py +++ b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/calendar/slice_tools.py @@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c def build_calendar_tools(dependencies: dict[str, Any]) -> list[BaseTool]: - """Tools from ``new_chat`` registry: ``calendar`` category.""" + """Registry-backed tools for the ``calendar`` category.""" return build_registry_tools_for_category(dependencies, "calendar") diff --git a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/confluence/slice_tools.py b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/confluence/slice_tools.py index 2889e8a3a..3f4f2d45c 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/confluence/slice_tools.py +++ b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/confluence/slice_tools.py @@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c def build_confluence_tools(dependencies: dict[str, Any]) -> list[BaseTool]: - """Tools from ``new_chat`` registry: ``confluence`` category.""" + """Registry-backed tools for the ``confluence`` category.""" return build_registry_tools_for_category(dependencies, "confluence") diff --git a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/discord/slice_tools.py b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/discord/slice_tools.py index 3511054ab..79eea4f3f 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/discord/slice_tools.py +++ b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/discord/slice_tools.py @@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c def build_discord_tools(dependencies: dict[str, Any]) -> list[BaseTool]: - """Tools from ``new_chat`` registry: ``discord`` category.""" + """Registry-backed tools for the ``discord`` category.""" return build_registry_tools_for_category(dependencies, "discord") diff --git a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/dropbox/slice_tools.py b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/dropbox/slice_tools.py index 3adc4a480..ff28a5b71 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/dropbox/slice_tools.py +++ b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/dropbox/slice_tools.py @@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c def build_dropbox_tools(dependencies: dict[str, Any]) -> list[BaseTool]: - """Tools from ``new_chat`` registry: ``dropbox`` category.""" + """Registry-backed tools for the ``dropbox`` category.""" return build_registry_tools_for_category(dependencies, "dropbox") diff --git a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/gmail/slice_tools.py b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/gmail/slice_tools.py index 50c070075..87876804e 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/gmail/slice_tools.py +++ b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/gmail/slice_tools.py @@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c def build_gmail_tools(dependencies: dict[str, Any]) -> list[BaseTool]: - """Tools from ``new_chat`` registry: ``gmail`` category.""" + """Registry-backed tools for the ``gmail`` category.""" return build_registry_tools_for_category(dependencies, "gmail") diff --git a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/google_drive/slice_tools.py b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/google_drive/slice_tools.py index 7f63f6eb3..ee6defe4b 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/google_drive/slice_tools.py +++ b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/google_drive/slice_tools.py @@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c def build_google_drive_tools(dependencies: dict[str, Any]) -> list[BaseTool]: - """Tools from ``new_chat`` registry: ``google_drive`` category.""" + """Registry-backed tools for the ``google_drive`` category.""" return build_registry_tools_for_category(dependencies, "google_drive") diff --git a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/luma/slice_tools.py b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/luma/slice_tools.py index 4e8350f2e..bf4efde00 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/luma/slice_tools.py +++ b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/luma/slice_tools.py @@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c def build_luma_tools(dependencies: dict[str, Any]) -> list[BaseTool]: - """Tools from ``new_chat`` registry: ``luma`` category.""" + """Registry-backed tools for the ``luma`` category.""" return build_registry_tools_for_category(dependencies, "luma") diff --git a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/notion/slice_tools.py b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/notion/slice_tools.py index 0229b5b82..4fecd13a4 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/notion/slice_tools.py +++ b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/notion/slice_tools.py @@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c def build_notion_tools(dependencies: dict[str, Any]) -> list[BaseTool]: - """Tools from ``new_chat`` registry: ``notion`` category.""" + """Registry-backed tools for the ``notion`` category.""" return build_registry_tools_for_category(dependencies, "notion") diff --git a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/onedrive/slice_tools.py b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/onedrive/slice_tools.py index 2f7c82dad..572cc6e36 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/onedrive/slice_tools.py +++ b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/onedrive/slice_tools.py @@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c def build_onedrive_tools(dependencies: dict[str, Any]) -> list[BaseTool]: - """Tools from ``new_chat`` registry: ``onedrive`` category.""" + """Registry-backed tools for the ``onedrive`` category.""" return build_registry_tools_for_category(dependencies, "onedrive") diff --git a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/teams/slice_tools.py b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/teams/slice_tools.py index b88f29843..e66ed3295 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/teams/slice_tools.py +++ b/surfsense_backend/app/agents/multi_agent_chat/expert_agent/connectors/teams/slice_tools.py @@ -10,5 +10,5 @@ from app.agents.multi_agent_chat.core.registry import build_registry_tools_for_c def build_teams_tools(dependencies: dict[str, Any]) -> list[BaseTool]: - """Tools from ``new_chat`` registry: ``teams`` category.""" + """Registry-backed tools for the ``teams`` category.""" return build_registry_tools_for_category(dependencies, "teams") diff --git a/surfsense_backend/app/agents/multi_agent_chat/integration/create_multi_agent_chat.py b/surfsense_backend/app/agents/multi_agent_chat/integration/create_multi_agent_chat.py index 06c022ec3..36c731735 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/integration/create_multi_agent_chat.py +++ b/surfsense_backend/app/agents/multi_agent_chat/integration/create_multi_agent_chat.py @@ -1,4 +1,4 @@ -"""Single entry: SurfSense connectors + multi-agent stack → compiled supervisor graph.""" +"""Build the multi-agent supervisor graph: MCP partition, registry, routing tools, optional SurfSense middleware.""" from __future__ import annotations @@ -11,14 +11,6 @@ from langchain_core.tools import BaseTool from langgraph.types import Checkpointer from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.chat_deepagent import _map_connectors_to_searchable_types -from app.agents.new_chat.context import SurfSenseContextSchema -from app.agents.new_chat.feature_flags import get_flags -from app.agents.new_chat.filesystem_backends import build_backend_resolver -from app.agents.new_chat.filesystem_selection import FilesystemSelection -from app.agents.new_chat.tools.mcp_tool import load_mcp_tools -from app.db import ChatVisibility - from app.agents.multi_agent_chat.core.mcp_partition import ( fetch_mcp_connector_metadata_maps, partition_mcp_tools_by_expert_route, @@ -27,14 +19,95 @@ from app.agents.multi_agent_chat.core.registry.dependencies import ( build_registry_dependencies, coerce_thread_id_for_registry, ) -from app.agents.multi_agent_chat.middleware.supervisor_stack import build_supervisor_middleware_stack -from app.agents.multi_agent_chat.routing.supervisor_routing import build_supervisor_routing_tools +from app.agents.multi_agent_chat.middleware.supervisor_stack import ( + build_supervisor_middleware_stack, +) +from app.agents.multi_agent_chat.routing.supervisor_routing import ( + build_supervisor_routing_tools, +) from app.agents.multi_agent_chat.supervisor import build_supervisor_agent +from app.agents.new_chat.chat_deepagent import _map_connectors_to_searchable_types +from app.agents.new_chat.context import SurfSenseContextSchema +from app.agents.new_chat.feature_flags import get_flags +from app.agents.new_chat.filesystem_backends import build_backend_resolver +from app.agents.new_chat.filesystem_selection import FilesystemSelection +from app.agents.new_chat.tools.mcp_tool import load_mcp_tools +from app.db import ChatVisibility logger = logging.getLogger(__name__) -def _compile_supervisor_chat_blocking( +async def _discover_connectors_and_doc_types( + *, + connector_service: Any | None, + search_space_id: int, + available_connectors: list[str] | None, + available_document_types: list[str] | None, +) -> tuple[list[str] | None, list[str] | None]: + """Fill connector / document-type lists from ``connector_service`` when callers omit them.""" + connectors = available_connectors + doc_types = available_document_types + if connector_service is None: + return connectors, doc_types + try: + if connectors is None: + raw = await connector_service.get_available_connectors(search_space_id) + if raw: + connectors = _map_connectors_to_searchable_types(raw) + if doc_types is None: + doc_types = await connector_service.get_available_document_types(search_space_id) + except Exception as exc: + logger.warning("Failed to discover available connectors/document types: %s", exc) + return connectors, doc_types + + +async def _mcp_tools_by_expert_route( + *, + db_session: AsyncSession, + search_space_id: int, +) -> dict[str, list[BaseTool]] | None: + mcp_flat = await load_mcp_tools(db_session, search_space_id) + id_map, name_map = await fetch_mcp_connector_metadata_maps(db_session, search_space_id) + return partition_mcp_tools_by_expert_route(mcp_flat, id_map, name_map) + + +def _make_supervisor_routing_tools( + llm: BaseChatModel, + *, + db_session: AsyncSession, + search_space_id: int, + user_id: str, + thread_id: str | int | None, + firecrawl_api_key: str | None, + connector_service: Any | None, + available_connectors: list[str] | None, + available_document_types: list[str] | None, + thread_visibility: ChatVisibility, + mcp_tools_by_route: dict[str, list[BaseTool]] | None, +) -> list[BaseTool]: + registry_dependencies = build_registry_dependencies( + db_session=db_session, + search_space_id=search_space_id, + user_id=user_id, + thread_id=thread_id, + llm=llm, + firecrawl_api_key=firecrawl_api_key, + connector_service=connector_service, + available_connectors=available_connectors, + available_document_types=available_document_types, + thread_visibility=thread_visibility, + ) + return build_supervisor_routing_tools( + llm, + registry_dependencies=registry_dependencies, + include_deliverables=coerce_thread_id_for_registry(thread_id) is not None, + mcp_tools_by_route=mcp_tools_by_route, + available_connectors=available_connectors, + thread_visibility=thread_visibility, + ) + + +def _compile_supervisor_agent_sync( *, llm: BaseChatModel, routing_tools: list[BaseTool], @@ -43,16 +116,16 @@ def _compile_supervisor_chat_blocking( filesystem_mode: Any, search_space_id: int, user_id: str, - thread_id: str | None, + thread_id: str | int | None, thread_visibility: ChatVisibility, anon_session_id: str | None, available_connectors: list[str] | None, available_document_types: list[str] | None, mentioned_document_ids: list[int] | None, max_input_tokens: int | None, + citations_enabled: bool, ) -> Any: - """CPU-heavy: middleware assembly + ``create_agent`` (runs in a worker thread).""" - flags = get_flags() + """CPU-heavy: middleware stack + ``create_agent`` (intended for ``asyncio.to_thread``).""" middleware = build_supervisor_middleware_stack( llm=llm, tools=routing_tools, @@ -67,7 +140,7 @@ def _compile_supervisor_chat_blocking( available_document_types=available_document_types, mentioned_document_ids=mentioned_document_ids, max_input_tokens=max_input_tokens, - flags=flags, + flags=get_flags(), ) return build_supervisor_agent( llm, @@ -76,6 +149,7 @@ def _compile_supervisor_chat_blocking( thread_visibility=thread_visibility, middleware=middleware, context_schema=SurfSenseContextSchema, + citations_enabled=citations_enabled, ) @@ -98,16 +172,17 @@ async def create_multi_agent_chat( mentioned_document_ids: list[int] | None = None, max_input_tokens: int | None = None, surfsense_stack: bool = True, + citations_enabled: bool | None = None, ): - """Build the full multi-agent chat graph (supervisor + domain subgraphs via routing tools). + """Build the full multi-agent chat graph (supervisor + expert subgraphs via routing tools). **Builtins** (:mod:`expert_agent.builtins`): registry-grouped **categories** (research, memory, deliverables). **Connectors** (:mod:`expert_agent.connectors`): **vendor integrations** — one subgraph per route in ``TOOL_NAMES_BY_CATEGORY`` (e.g. calendar, confluence, discord, dropbox, gmail, google_drive, luma, notion, onedrive, teams). - MCP tools from ``new_chat`` (``load_mcp_tools``) are partitioned inside this package and attached only + MCP tools (via ``load_mcp_tools``) are partitioned inside this package and attached only to the matching expert subgraphs — not to the supervisor tool list as raw MCP calls. Inclusion matches - ``new_chat.tools.registry.build_tools_async``: all tools returned by ``load_mcp_tools`` are merged + ``app.agents.new_chat.tools.registry.build_tools_async``: all tools returned by ``load_mcp_tools`` are merged after partitioning (no extra inventory filter on MCP). Connector routing uses ``available_connectors``: pass explicitly, or provide ``connector_service`` so lists are resolved like ``create_surfsense_deep_agent`` (``get_available_connectors`` → searchable types). @@ -115,57 +190,38 @@ async def create_multi_agent_chat( Deliverables (thread-scoped reports, podcasts, etc.) are registered only when ``thread_id`` is set. When ``surfsense_stack`` is true (default), the supervisor uses the same SurfSense middleware shell as - ``new_chat`` (KB priority/tree, filesystem, compaction, permissions, etc.) except ``SubAgentMiddleware`` / - ``task``, since experts are separate graphs behind routing tools. Graph compilation runs in - ``asyncio.to_thread`` so heavy CPU work does not block the event loop. + the main single-agent chat (KB priority/tree, filesystem, compaction, permissions, etc.) except + ``SubAgentMiddleware`` / ``task``, since experts are separate graphs behind routing tools. Graph + compilation runs in ``asyncio.to_thread`` so heavy CPU work does not block the event loop. + + ``citations_enabled``: when ``None``, defaults to ``True`` (same default as ``AgentConfig`` / main chat). """ - resolved_connectors = available_connectors - resolved_doc_types = available_document_types - if connector_service is not None: - try: - if resolved_connectors is None: - connector_types = await connector_service.get_available_connectors( - search_space_id - ) - if connector_types: - resolved_connectors = _map_connectors_to_searchable_types( - connector_types - ) - if resolved_doc_types is None: - resolved_doc_types = ( - await connector_service.get_available_document_types(search_space_id) - ) - except Exception as exc: - logger.warning( - "Failed to discover available connectors/document types: %s", - exc, - ) + citations = True if citations_enabled is None else citations_enabled + connectors, doc_types = await _discover_connectors_and_doc_types( + connector_service=connector_service, + search_space_id=search_space_id, + available_connectors=available_connectors, + available_document_types=available_document_types, + ) - mcp_tools_by_route: dict[str, list[BaseTool]] | None = None + mcp_by_route: dict[str, list[BaseTool]] | None = None if include_mcp_tools: - mcp_flat = await load_mcp_tools(db_session, search_space_id) - id_map, name_map = await fetch_mcp_connector_metadata_maps(db_session, search_space_id) - mcp_tools_by_route = partition_mcp_tools_by_expert_route(mcp_flat, id_map, name_map) + mcp_by_route = await _mcp_tools_by_expert_route( + db_session=db_session, search_space_id=search_space_id + ) - registry_dependencies = build_registry_dependencies( + routing_tools = _make_supervisor_routing_tools( + llm, db_session=db_session, search_space_id=search_space_id, user_id=user_id, thread_id=thread_id, - llm=llm, firecrawl_api_key=firecrawl_api_key, connector_service=connector_service, - available_connectors=resolved_connectors, - available_document_types=resolved_doc_types, - thread_visibility=thread_visibility, - ) - routing_tools = build_supervisor_routing_tools( - llm, - registry_dependencies=registry_dependencies, - include_deliverables=coerce_thread_id_for_registry(thread_id) is not None, - mcp_tools_by_route=mcp_tools_by_route, - available_connectors=resolved_connectors, + available_connectors=connectors, + available_document_types=doc_types, thread_visibility=thread_visibility, + mcp_tools_by_route=mcp_by_route, ) fs_sel = filesystem_selection or FilesystemSelection() @@ -177,10 +233,11 @@ async def create_multi_agent_chat( tools=routing_tools, checkpointer=checkpointer, thread_visibility=thread_visibility, + citations_enabled=citations, ) return await asyncio.to_thread( - _compile_supervisor_chat_blocking, + _compile_supervisor_agent_sync, llm=llm, routing_tools=routing_tools, checkpointer=checkpointer, @@ -191,8 +248,9 @@ async def create_multi_agent_chat( thread_id=thread_id, thread_visibility=thread_visibility, anon_session_id=anon_session_id, - available_connectors=resolved_connectors, - available_document_types=resolved_doc_types, + available_connectors=connectors, + available_document_types=doc_types, mentioned_document_ids=mentioned_document_ids, max_input_tokens=max_input_tokens, + citations_enabled=citations, ) diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/__init__.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/__init__.py index 130b0508f..058cf705a 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/__init__.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/__init__.py @@ -1,4 +1,4 @@ -"""SurfSense supervisor middleware (parity with ``new_chat`` main agent, minus subagents).""" +"""SurfSense supervisor middleware (parity with the main single-agent chat, minus subagents).""" from app.agents.multi_agent_chat.middleware.supervisor_stack import ( build_supervisor_middleware_stack, diff --git a/surfsense_backend/app/agents/multi_agent_chat/middleware/supervisor_stack.py b/surfsense_backend/app/agents/multi_agent_chat/middleware/supervisor_stack.py index 40b377cbf..0cd390949 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/middleware/supervisor_stack.py +++ b/surfsense_backend/app/agents/multi_agent_chat/middleware/supervisor_stack.py @@ -1,4 +1,4 @@ -"""Supervisor middleware stack matching ``new_chat`` main agent (no ``SubAgentMiddleware`` / ``task``).""" +"""Supervisor middleware stack matching the main single-agent chat (no ``SubAgentMiddleware`` / ``task``).""" from __future__ import annotations diff --git a/surfsense_backend/app/agents/multi_agent_chat/routing/route_connector_gate.py b/surfsense_backend/app/agents/multi_agent_chat/routing/route_connector_gate.py index b66b5eb4a..84e2359e7 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/routing/route_connector_gate.py +++ b/surfsense_backend/app/agents/multi_agent_chat/routing/route_connector_gate.py @@ -1,4 +1,4 @@ -"""Gate supervisor routing tools by connected searchable connector types (aligned with ``new_chat`` KB). +"""Gate supervisor routing tools by connected searchable connector types (aligned with main chat KB). When ``available_connectors`` is ``None``, all routes are emitted (caller did not pass an inventory). diff --git a/surfsense_backend/app/agents/multi_agent_chat/routing/supervisor_routing.py b/surfsense_backend/app/agents/multi_agent_chat/routing/supervisor_routing.py index ef496ab17..63f4da744 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/routing/supervisor_routing.py +++ b/surfsense_backend/app/agents/multi_agent_chat/routing/supervisor_routing.py @@ -116,7 +116,7 @@ def build_supervisor_routing_tools( ``mcp_tools_by_route`` maps route keys to MCP tools merged into the matching expert subgraph. - When ``available_connectors`` is set (searchable connector strings, same shape as ``new_chat``), + When ``available_connectors`` is set (searchable connector strings, same shape as the main chat agent), a connector-backed route is registered only if its required searchable connector type is available. """ if registry_dependencies is None: diff --git a/surfsense_backend/app/agents/multi_agent_chat/supervisor/graph.py b/surfsense_backend/app/agents/multi_agent_chat/supervisor/graph.py index abb3bee8d..7823a0380 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/supervisor/graph.py +++ b/surfsense_backend/app/agents/multi_agent_chat/supervisor/graph.py @@ -1,80 +1,18 @@ -"""Compile the supervisor agent graph (supervisor prompt + caller-supplied routing tools).""" +"""Compile the supervisor agent graph (LangChain ``create_agent`` + caller routing tools).""" from __future__ import annotations from collections.abc import Sequence from typing import Any -import app.agents.multi_agent_chat.supervisor as supervisor_pkg - from langchain.agents import create_agent from langchain_core.language_models import BaseChatModel from langchain_core.tools import BaseTool from langgraph.types import Checkpointer -from app.agents.multi_agent_chat.core.prompts import read_prompt_md - -_BUILTIN_SPECIALISTS: frozenset[str] = frozenset({"research", "memory", "deliverables"}) -_SPECIALIST_CAPABILITIES: dict[str, str] = { - "research": "external research: web lookup, source gathering, and SurfSense documentation help.", - "memory": "save durable long-lived memory items.", - "deliverables": "deliverables and shareable artifacts: reports, podcasts, video presentations, resumes, and images.", - "gmail": "email inbox actions: search/read emails, draft updates, send messages, and trash emails.", - "calendar": "scheduling actions: check availability, inspect events, create events, and update events.", - "google_drive": "Drive file/document actions: locate files, inspect content, and manage files/folders.", - "notion": "Notion page actions: create pages, update content, and delete pages.", - "confluence": "Confluence page actions: find/read pages and create/update pages.", - "dropbox": "Dropbox file storage actions: browse folders, read files, and manage file content.", - "onedrive": "OneDrive file storage actions: browse folders, read files, and manage file content.", - "discord": "Discord communication actions: read channels/threads and post replies.", - "teams": "Microsoft Teams communication actions: read channels/threads and post replies.", - "luma": "Luma event actions: list events, inspect event details, and create events.", - "linear": "Linear workflow actions: search/update issues and inspect projects/cycles.", - "jira": "Jira workflow actions: search/update issues and manage workflow transitions.", - "clickup": "ClickUp workflow actions: find/update tasks and lists.", - "airtable": "Airtable data actions: locate bases/tables and create/read/update records.", - "slack": "Slack communication actions: read channel/thread history and post replies.", - # generic_mcp specialist intentionally disabled for now. - # "generic_mcp": "handle tasks through user-defined custom app integration tools not covered above.", -} -_SPECIALIST_ORDER: tuple[str, ...] = tuple(_SPECIALIST_CAPABILITIES.keys()) - - -def _memory_capability_for_visibility(thread_visibility: Any | None) -> str: - vis = str(getattr(thread_visibility, "value", thread_visibility)).upper() - if vis == "SEARCH_SPACE": - return "team memory actions: save shared team preferences, conventions, and long-lived team facts." - return "user memory actions: save personal preferences, instructions, and long-lived user facts." - - -def _render_available_specialists_list( - tools: Sequence[BaseTool], - *, - thread_visibility: Any | None, -) -> str: - available_names = { - tool.name for tool in tools if isinstance(getattr(tool, "name", None), str) - } - capabilities = dict(_SPECIALIST_CAPABILITIES) - capabilities["memory"] = _memory_capability_for_visibility(thread_visibility) - lines: list[str] = [] - for name in _SPECIALIST_ORDER: - if name in _BUILTIN_SPECIALISTS or name in available_names: - capability = capabilities[name] - lines.append(f"- {name}: {capability}") - return "\n".join(lines) - - -def _render_supervisor_prompt( - template: str, - tools: Sequence[BaseTool], - *, - thread_visibility: Any | None, -) -> str: - specialist_list = _render_available_specialists_list( - tools, thread_visibility=thread_visibility - ) - return template.replace("{{AVAILABLE_SPECIALISTS_LIST}}", specialist_list) +from app.agents.multi_agent_chat.supervisor.prompt_assembly import ( + build_supervisor_system_prompt, +) def build_supervisor_agent( @@ -85,13 +23,13 @@ def build_supervisor_agent( thread_visibility: Any | None = None, middleware: Sequence[Any] | None = None, context_schema: Any | None = None, + citations_enabled: bool = True, ): """Compile the supervisor **agent** (graph). ``tools`` = output of ``build_supervisor_routing_tools``.""" - template = read_prompt_md(supervisor_pkg.__name__, "supervisor_prompt") - system_prompt = _render_supervisor_prompt( - template, + system_prompt = build_supervisor_system_prompt( tools, thread_visibility=thread_visibility, + citations_enabled=citations_enabled, ) kwargs: dict[str, Any] = { "system_prompt": system_prompt, diff --git a/surfsense_backend/app/agents/multi_agent_chat/supervisor/prompt_assembly.py b/surfsense_backend/app/agents/multi_agent_chat/supervisor/prompt_assembly.py new file mode 100644 index 000000000..ac7140c7d --- /dev/null +++ b/surfsense_backend/app/agents/multi_agent_chat/supervisor/prompt_assembly.py @@ -0,0 +1,128 @@ +"""Supervisor system prompt: template load, shared agent-identity injection, specialist list.""" + +from __future__ import annotations + +from collections.abc import Sequence +from datetime import UTC, datetime +from typing import Any + +from langchain_core.tools import BaseTool + +import app.agents.multi_agent_chat.supervisor as supervisor_pkg +from app.agents.multi_agent_chat.core.prompts import read_prompt_md +from app.agents.new_chat.prompts.composer import _build_citation_block, _read_fragment +from app.db import ChatVisibility + +_MEMORY_SPECIALIST_PHRASE = "invoke the **memory** specialist" + +_BUILTIN_SPECIALISTS: frozenset[str] = frozenset({"research", "memory", "deliverables"}) +_SPECIALIST_CAPABILITIES: dict[str, str] = { + "research": "external research: web lookup, source gathering, and SurfSense documentation help.", + "memory": "save durable long-lived memory items.", + "deliverables": "deliverables and shareable artifacts: reports, podcasts, video presentations, resumes, and images.", + "gmail": "email inbox actions: search/read emails, draft updates, send messages, and trash emails.", + "calendar": "scheduling actions: check availability, inspect events, create events, and update events.", + "google_drive": "Drive file/document actions: locate files, inspect content, and manage files/folders.", + "notion": "Notion page actions: create pages, update content, and delete pages.", + "confluence": "Confluence page actions: find/read pages and create/update pages.", + "dropbox": "Dropbox file storage actions: browse folders, read files, and manage file content.", + "onedrive": "OneDrive file storage actions: browse folders, read files, and manage file content.", + "discord": "Discord communication actions: read channels/threads and post replies.", + "teams": "Microsoft Teams communication actions: read channels/threads and post replies.", + "luma": "Luma event actions: list events, inspect event details, and create events.", + "linear": "Linear workflow actions: search/update issues and inspect projects/cycles.", + "jira": "Jira workflow actions: search/update issues and manage workflow transitions.", + "clickup": "ClickUp workflow actions: find/update tasks and lists.", + "airtable": "Airtable data actions: locate bases/tables and create/read/update records.", + "slack": "Slack communication actions: read channel/thread history and post replies.", + # generic_mcp specialist intentionally disabled for now. + # "generic_mcp": "handle tasks through user-defined custom app integration tools not covered above.", +} +_SPECIALIST_ORDER: tuple[str, ...] = tuple(_SPECIALIST_CAPABILITIES.keys()) + + +def _normalize_chat_visibility(thread_visibility: Any | None) -> ChatVisibility: + if thread_visibility is None: + return ChatVisibility.PRIVATE + if thread_visibility == ChatVisibility.SEARCH_SPACE: + return ChatVisibility.SEARCH_SPACE + raw = getattr(thread_visibility, "value", thread_visibility) + if str(raw).upper() == "SEARCH_SPACE": + return ChatVisibility.SEARCH_SPACE + return ChatVisibility.PRIVATE + + +def _identity_fragment_key(thread_visibility: Any | None) -> str: + """``private`` / ``team`` suffix for ``agent_*`` and ``memory_protocol_*`` fragments.""" + return ( + "team" + if _normalize_chat_visibility(thread_visibility) == ChatVisibility.SEARCH_SPACE + else "private" + ) + + +def _compose_identity_memory_citations( + *, + thread_visibility: Any | None, + citations_enabled: bool, +) -> str: + """Main-chat identity, memory protocol, and citation fragments (supervisor slice only).""" + key = _identity_fragment_key(thread_visibility) + today = datetime.now(UTC).date().isoformat() + + intro = _read_fragment(f"base/agent_{key}.md") + if intro: + intro = intro.format(resolved_today=today) + + memory = _read_fragment(f"base/memory_protocol_{key}.md").replace( + "call update_memory", + _MEMORY_SPECIALIST_PHRASE, + ) + tail = ( + f"\n{memory}\n\n\n" + + _build_citation_block(citations_enabled) + ) + return "\n\n".join(part for part in (intro.strip(), tail.strip()) if part) + + +def _memory_specialist_capability(thread_visibility: Any | None) -> str: + vis = str(getattr(thread_visibility, "value", thread_visibility)).upper() + if vis == "SEARCH_SPACE": + return "team memory actions: save shared team preferences, conventions, and long-lived team facts." + return "user memory actions: save personal preferences, instructions, and long-lived user facts." + + +def _specialists_markdown( + tools: Sequence[BaseTool], + *, + thread_visibility: Any | None, +) -> str: + available_names = { + tool.name for tool in tools if isinstance(getattr(tool, "name", None), str) + } + capabilities = dict(_SPECIALIST_CAPABILITIES) + capabilities["memory"] = _memory_specialist_capability(thread_visibility) + lines: list[str] = [] + for name in _SPECIALIST_ORDER: + if name in _BUILTIN_SPECIALISTS or name in available_names: + lines.append(f"- {name}: {capabilities[name]}") + return "\n".join(lines) + + +def build_supervisor_system_prompt( + tools: Sequence[BaseTool], + *, + thread_visibility: Any | None, + citations_enabled: bool, +) -> str: + """Load ``supervisor_prompt.md`` and fill placeholders.""" + template = read_prompt_md(supervisor_pkg.__name__, "supervisor_prompt") + specialists = _specialists_markdown(tools, thread_visibility=thread_visibility) + injected = _compose_identity_memory_citations( + thread_visibility=thread_visibility, + citations_enabled=citations_enabled, + ) + return template.replace("{{AVAILABLE_SPECIALISTS_LIST}}", specialists).replace( + "{{SUPERVISOR_BASE_INJECTION}}", + injected, + ) diff --git a/surfsense_backend/app/agents/multi_agent_chat/supervisor/supervisor_prompt.md b/surfsense_backend/app/agents/multi_agent_chat/supervisor/supervisor_prompt.md index 790e98753..632c888c9 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/supervisor/supervisor_prompt.md +++ b/surfsense_backend/app/agents/multi_agent_chat/supervisor/supervisor_prompt.md @@ -1,9 +1,8 @@ -You are SurfSense's multi-agent supervisor. +{{SUPERVISOR_BASE_INJECTION}} - -Your job is to decide whether to answer directly or delegate to one or more specialists. -You optimize for correctness, low confusion, and minimal unnecessary delegation. - + +In this **multi-agent** session you also **coordinate specialists** (listed below): call a specialist only when their domain matches the need; give each call a compact, outcome-focused task; merge structured results into one clear user-facing reply. When you can satisfy the turn with your own tools and reasoning, do so without delegating. + Use only the specialists listed below. diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py index eec0bcfdc..5d1f4b572 100644 --- a/surfsense_backend/app/tasks/chat/stream_new_chat.py +++ b/surfsense_backend/app/tasks/chat/stream_new_chat.py @@ -1581,6 +1581,7 @@ async def stream_new_chat( thread_visibility=visibility, filesystem_selection=filesystem_selection, mentioned_document_ids=mentioned_document_ids, + citations_enabled=agent_config.citations_enabled, ) else: agent = await create_surfsense_deep_agent( @@ -2305,6 +2306,7 @@ async def stream_resume_chat( connector_service=connector_service, thread_visibility=visibility, filesystem_selection=filesystem_selection, + citations_enabled=agent_config.citations_enabled, ) else: agent = await create_surfsense_deep_agent(