mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-06-06 20:15:17 +02:00
refactor(agents): relocate remaining MAC-only kernel (permissions, deliverable_wait)
permissions.py (authorization Rule/Ruleset model) is consumed across all MAC subagents + the permissions middleware, with a single external consumer (user_tool_allowlist service) -> move to multi_agent_chat/shared/permissions.py and repoint all 42 sites. deliverable_wait.py (wait_for_deliverable) is used only by the podcast and video_presentation deliverable tools -> colocate into subagents/builtins/deliverables/. No behavior change; import-all + permission/allowlist/deliverable unit tests stay green.
This commit is contained in:
CREDO23
parent
714c5ffea9
commit
f615d6b530
47 changed files with 61 additions and 53 deletions
|
|
@ -0,0 +1,123 @@
|
|||
"""Shared poll-until-terminal helper for Celery-backed deliverables.
|
||||
|
||||
Lives in ``app.agents.shared`` (neutral kernel package, no dependency on
|
||||
``multi_agent_chat``) so both the shared tools under ``app/agents/shared/tools/``
|
||||
and the multi-agent subagent tools under
|
||||
``app/agents/multi_agent_chat/subagents/builtins/deliverables/tools/`` can import
|
||||
it without creating a circular dependency.
|
||||
|
||||
Background
|
||||
----------
|
||||
Tools like ``generate_podcast`` and ``generate_video_presentation`` enqueue
|
||||
the heavy work to Celery and historically returned immediately with a
|
||||
"pending" status. That works for very-long deliverables but hurts UX for
|
||||
the common case (most podcasts finish in 10-30 seconds): the agent sends
|
||||
a "kicked off, check back in a minute" reply *before* the worker is done,
|
||||
so the user never gets a "ready" confirmation.
|
||||
|
||||
This helper bridges that gap. The tool dispatches the Celery task as
|
||||
before, then polls the artefact row's ``status`` column **until it
|
||||
reaches a terminal value** (READY / FAILED). The tool then returns a
|
||||
real terminal outcome — never a pending one.
|
||||
|
||||
No wall-clock budget here on purpose
|
||||
------------------------------------
|
||||
Layering a second budget on top of the existing per-invocation safety
|
||||
nets just confused the UX. The real ceilings are:
|
||||
|
||||
* **Multi-agent mode** — ``SURFSENSE_SUBAGENT_INVOKE_TIMEOUT_SECONDS``
|
||||
(default ``300.0``, ``0`` to disable) caps how long any single
|
||||
``task(subagent, ...)`` invocation can run. If a deliverable needs
|
||||
longer than this, the subagent invocation is cancelled and the
|
||||
orchestrator surfaces a "subagent timed out" ToolMessage. Operators
|
||||
who routinely generate long videos should raise that ceiling (or set
|
||||
it to ``0`` for true unbounded waits).
|
||||
* **Single-agent mode** — the chat's HTTP stream / process lifetime is
|
||||
the only ceiling. Truly indefinite waits work here, but a dead Celery
|
||||
worker will leave the row in PENDING/GENERATING forever; treat that
|
||||
as an operational concern, not a UX concern.
|
||||
|
||||
Configuration
|
||||
-------------
|
||||
None. The poll cadence is hardcoded at 1.5s — small enough to feel
|
||||
responsive (~6 polls per typical 10s podcast), large enough to avoid
|
||||
hammering the DB under burst traffic. Override at the call site if a
|
||||
specific tool needs a different cadence.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import InstrumentedAttribute
|
||||
|
||||
from app.db import shielded_async_session
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
_DEFAULT_POLL_INTERVAL_SECONDS: float = 1.5
|
||||
|
||||
|
||||
async def wait_for_deliverable(
|
||||
*,
|
||||
model: type,
|
||||
row_id: int,
|
||||
columns: list[InstrumentedAttribute[Any]],
|
||||
terminal_statuses: set[Enum],
|
||||
poll_interval_s: float = _DEFAULT_POLL_INTERVAL_SECONDS,
|
||||
) -> tuple[Enum, tuple[Any, ...], float]:
|
||||
"""Poll ``model`` row ``row_id`` until ``columns[0]`` reaches a terminal status.
|
||||
|
||||
Blocks until the row's status column matches one of
|
||||
``terminal_statuses``. There is no internal wall-clock budget; cancel
|
||||
from the outside (subagent timeout, HTTP disconnect, task
|
||||
cancellation) if you need a ceiling. See module docstring.
|
||||
|
||||
The first entry of ``columns`` must be the status column; additional
|
||||
columns (e.g. ``Podcast.file_location``) are returned alongside the
|
||||
final status so callers can build their payload without a second
|
||||
roundtrip.
|
||||
|
||||
A fresh ``shielded_async_session`` is opened per poll so we never
|
||||
hold a transaction across the wait, and a failed poll is logged but
|
||||
does not abort the wait — transient DB hiccups should not collapse
|
||||
the tool call.
|
||||
|
||||
Returns
|
||||
-------
|
||||
``(terminal_status, columns, elapsed_seconds)``
|
||||
``columns`` mirrors the requested ``columns`` (including the
|
||||
status itself in position 0).
|
||||
"""
|
||||
if not columns:
|
||||
raise ValueError("wait_for_deliverable requires at least the status column")
|
||||
|
||||
start = time.monotonic()
|
||||
|
||||
while True:
|
||||
await asyncio.sleep(poll_interval_s)
|
||||
row: tuple[Any, ...] | None = None
|
||||
try:
|
||||
async with shielded_async_session() as session:
|
||||
result = await session.execute(
|
||||
select(*columns).where(model.id == row_id)
|
||||
)
|
||||
row = result.first()
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"[deliverable_wait] poll failed model=%s id=%s err=%r",
|
||||
getattr(model, "__name__", str(model)),
|
||||
row_id,
|
||||
exc,
|
||||
)
|
||||
|
||||
if row is not None:
|
||||
status_val = row[0]
|
||||
if status_val in terminal_statuses:
|
||||
return status_val, tuple(row), time.monotonic() - start
|
||||
|
|
@ -9,7 +9,7 @@ from typing import Any
|
|||
|
||||
from langchain_core.tools import BaseTool
|
||||
|
||||
from app.agents.shared.permissions import Ruleset
|
||||
from app.agents.multi_agent_chat.shared.permissions import Ruleset
|
||||
|
||||
from .generate_image import create_generate_image_tool
|
||||
from .podcast import create_generate_podcast_tool
|
||||
|
|
|
|||
|
|
@ -18,7 +18,9 @@ from sqlalchemy.ext.asyncio import AsyncSession
|
|||
|
||||
from app.agents.multi_agent_chat.shared.receipts.command import with_receipt
|
||||
from app.agents.multi_agent_chat.shared.receipts.receipt import make_receipt
|
||||
from app.agents.shared.deliverable_wait import wait_for_deliverable
|
||||
from app.agents.multi_agent_chat.subagents.builtins.deliverables.deliverable_wait import (
|
||||
wait_for_deliverable,
|
||||
)
|
||||
from app.db import Podcast, PodcastStatus, shielded_async_session
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
|
@ -96,7 +98,7 @@ def create_generate_podcast_tool(
|
|||
# Wait until the Celery worker flips the row to a terminal
|
||||
# state. The wait is bounded only by the subagent invoke
|
||||
# timeout (multi-agent) or HTTP lifetime (single-agent) —
|
||||
# see app.agents.shared.deliverable_wait for details.
|
||||
# see app.agents.multi_agent_chat.subagents.builtins.deliverables.deliverable_wait for details.
|
||||
terminal_status, columns, elapsed = await wait_for_deliverable(
|
||||
model=Podcast,
|
||||
row_id=podcast_id,
|
||||
|
|
|
|||
|
|
@ -19,7 +19,9 @@ from sqlalchemy.ext.asyncio import AsyncSession
|
|||
|
||||
from app.agents.multi_agent_chat.shared.receipts.command import with_receipt
|
||||
from app.agents.multi_agent_chat.shared.receipts.receipt import make_receipt
|
||||
from app.agents.shared.deliverable_wait import wait_for_deliverable
|
||||
from app.agents.multi_agent_chat.subagents.builtins.deliverables.deliverable_wait import (
|
||||
wait_for_deliverable,
|
||||
)
|
||||
from app.db import VideoPresentation, VideoPresentationStatus, shielded_async_session
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
|
@ -83,7 +85,7 @@ def create_generate_video_presentation_tool(
|
|||
# Wait until the Celery worker flips the row to a terminal
|
||||
# state. The wait is bounded only by the subagent invoke
|
||||
# timeout (multi-agent) or HTTP lifetime (single-agent) —
|
||||
# see app.agents.shared.deliverable_wait for details.
|
||||
# see app.agents.multi_agent_chat.subagents.builtins.deliverables.deliverable_wait for details.
|
||||
terminal_status, _columns, elapsed = await wait_for_deliverable(
|
||||
model=VideoPresentation,
|
||||
row_id=video_pres_id,
|
||||
|
|
|
|||
|
|
@ -13,9 +13,9 @@ from deepagents import SubAgent
|
|||
from langchain_core.language_models import BaseChatModel
|
||||
from langchain_core.tools import BaseTool
|
||||
|
||||
from app.agents.multi_agent_chat.shared.permissions import Rule, Ruleset
|
||||
from app.agents.multi_agent_chat.subagents.shared.spec import SurfSenseSubagentSpec
|
||||
from app.agents.shared.filesystem_selection import FilesystemMode
|
||||
from app.agents.shared.permissions import Rule, Ruleset
|
||||
|
||||
from .middleware_stack import build_kb_middleware
|
||||
from .prompts import load_description, load_readonly_system_prompt, load_system_prompt
|
||||
|
|
|
|||
|
|
@ -28,9 +28,9 @@ from app.agents.multi_agent_chat.shared.middleware.patch_tool_calls import (
|
|||
from app.agents.multi_agent_chat.shared.middleware.permissions import (
|
||||
build_permission_mw,
|
||||
)
|
||||
from app.agents.multi_agent_chat.shared.permissions import Ruleset
|
||||
from app.agents.shared.feature_flags import AgentFeatureFlags
|
||||
from app.agents.shared.filesystem_selection import FilesystemMode
|
||||
from app.agents.shared.permissions import Ruleset
|
||||
|
||||
|
||||
def _kb_user_allowlist(
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from typing import Any
|
|||
|
||||
from langchain_core.tools import BaseTool
|
||||
|
||||
from app.agents.shared.permissions import Ruleset
|
||||
from app.agents.multi_agent_chat.shared.permissions import Ruleset
|
||||
from app.db import ChatVisibility
|
||||
|
||||
from .update_memory import create_update_memory_tool, create_update_team_memory_tool
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from typing import Any
|
|||
|
||||
from langchain_core.tools import BaseTool
|
||||
|
||||
from app.agents.shared.permissions import Ruleset
|
||||
from app.agents.multi_agent_chat.shared.permissions import Ruleset
|
||||
|
||||
from .scrape_webpage import create_scrape_webpage_tool
|
||||
from .web_search import create_web_search_tool
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from app.agents.shared.permissions import Rule, Ruleset
|
||||
from app.agents.multi_agent_chat.shared.permissions import Rule, Ruleset
|
||||
|
||||
NAME = "airtable"
|
||||
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ from typing import Any
|
|||
|
||||
from langchain_core.tools import BaseTool
|
||||
|
||||
from app.agents.shared.permissions import Ruleset
|
||||
from app.agents.multi_agent_chat.shared.permissions import Ruleset
|
||||
|
||||
from .create_event import create_create_calendar_event_tool
|
||||
from .delete_event import create_delete_calendar_event_tool
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from app.agents.shared.permissions import Rule, Ruleset
|
||||
from app.agents.multi_agent_chat.shared.permissions import Rule, Ruleset
|
||||
|
||||
NAME = "clickup"
|
||||
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from typing import Any
|
|||
|
||||
from langchain_core.tools import BaseTool
|
||||
|
||||
from app.agents.shared.permissions import Ruleset
|
||||
from app.agents.multi_agent_chat.shared.permissions import Ruleset
|
||||
|
||||
from .create_page import create_create_confluence_page_tool
|
||||
from .delete_page import create_delete_confluence_page_tool
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from typing import Any
|
|||
|
||||
from langchain_core.tools import BaseTool
|
||||
|
||||
from app.agents.shared.permissions import Ruleset
|
||||
from app.agents.multi_agent_chat.shared.permissions import Ruleset
|
||||
|
||||
from .list_channels import create_list_discord_channels_tool
|
||||
from .read_messages import create_read_discord_messages_tool
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from typing import Any
|
|||
|
||||
from langchain_core.tools import BaseTool
|
||||
|
||||
from app.agents.shared.permissions import Ruleset
|
||||
from app.agents.multi_agent_chat.shared.permissions import Ruleset
|
||||
|
||||
from .create_file import create_create_dropbox_file_tool
|
||||
from .trash_file import create_delete_dropbox_file_tool
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from typing import Any
|
|||
|
||||
from langchain_core.tools import BaseTool
|
||||
|
||||
from app.agents.shared.permissions import Ruleset
|
||||
from app.agents.multi_agent_chat.shared.permissions import Ruleset
|
||||
|
||||
from .create_draft import create_create_gmail_draft_tool
|
||||
from .read_email import create_read_gmail_email_tool
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from typing import Any
|
|||
|
||||
from langchain_core.tools import BaseTool
|
||||
|
||||
from app.agents.shared.permissions import Ruleset
|
||||
from app.agents.multi_agent_chat.shared.permissions import Ruleset
|
||||
|
||||
from .create_file import create_create_google_drive_file_tool
|
||||
from .trash_file import create_delete_google_drive_file_tool
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from app.agents.shared.permissions import Rule, Ruleset
|
||||
from app.agents.multi_agent_chat.shared.permissions import Rule, Ruleset
|
||||
|
||||
NAME = "jira"
|
||||
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from app.agents.shared.permissions import Rule, Ruleset
|
||||
from app.agents.multi_agent_chat.shared.permissions import Rule, Ruleset
|
||||
|
||||
NAME = "linear"
|
||||
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from typing import Any
|
|||
|
||||
from langchain_core.tools import BaseTool
|
||||
|
||||
from app.agents.shared.permissions import Ruleset
|
||||
from app.agents.multi_agent_chat.shared.permissions import Ruleset
|
||||
|
||||
from .create_event import create_create_luma_event_tool
|
||||
from .list_events import create_list_luma_events_tool
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from typing import Any
|
|||
|
||||
from langchain_core.tools import BaseTool
|
||||
|
||||
from app.agents.shared.permissions import Ruleset
|
||||
from app.agents.multi_agent_chat.shared.permissions import Ruleset
|
||||
|
||||
from .create_page import create_create_notion_page_tool
|
||||
from .delete_page import create_delete_notion_page_tool
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from typing import Any
|
|||
|
||||
from langchain_core.tools import BaseTool
|
||||
|
||||
from app.agents.shared.permissions import Ruleset
|
||||
from app.agents.multi_agent_chat.shared.permissions import Ruleset
|
||||
|
||||
from .create_file import create_create_onedrive_file_tool
|
||||
from .trash_file import create_delete_onedrive_file_tool
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from app.agents.shared.permissions import Rule, Ruleset
|
||||
from app.agents.multi_agent_chat.shared.permissions import Rule, Ruleset
|
||||
|
||||
NAME = "slack"
|
||||
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from typing import Any
|
|||
|
||||
from langchain_core.tools import BaseTool
|
||||
|
||||
from app.agents.shared.permissions import Ruleset
|
||||
from app.agents.multi_agent_chat.shared.permissions import Ruleset
|
||||
|
||||
from .list_channels import create_list_teams_channels_tool
|
||||
from .read_messages import create_read_teams_messages_tool
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from typing import Any
|
|||
|
||||
from deepagents import SubAgent
|
||||
|
||||
from app.agents.shared.permissions import Ruleset
|
||||
from app.agents.multi_agent_chat.shared.permissions import Ruleset
|
||||
|
||||
# A context-hint provider receives the parent-agent ``runtime.state`` mapping
|
||||
# and the ``description`` the orchestrator wrote, and returns a short string
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ from langchain_core.tools import BaseTool
|
|||
from app.agents.multi_agent_chat.shared.middleware.permissions import (
|
||||
build_permission_mw,
|
||||
)
|
||||
from app.agents.multi_agent_chat.shared.permissions import Ruleset
|
||||
from app.agents.multi_agent_chat.subagents.shared.md_file_reader import (
|
||||
read_shared_snippet,
|
||||
)
|
||||
|
|
@ -22,7 +23,6 @@ from app.agents.multi_agent_chat.subagents.shared.spec import (
|
|||
ContextHintProvider,
|
||||
SurfSenseSubagentSpec,
|
||||
)
|
||||
from app.agents.shared.permissions import Ruleset
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue