SurfSense/surfsense_backend/app/gateway/agent_invoke.py
CREDO23 7fb0707933 refactor(backend): rename search_space -> workspace across app bulk (Phase 2 Wave D)
Scoped codemod over surfsense_backend/app (excluding routes/, Wave E): renames
search_space_id -> workspace_id, search_space -> workspace, SearchSpace -> Workspace
across services, utils, tasks, agents, gateway, event_bus, notifications, podcasts,
automations, observability params, and prompt .md files. Also flips the camelCase
payload key searchSpaceId -> workspaceId (no backend reader; hard cutover).

Preserved carve-outs (verbatim): Celery task names "delete_search_space_background"
and "ai_sort_search_space" (wire names), and the OTel/metric key "search_space.id"
(dashboards depend on it). Enum values 'SEARCH_SPACE' and SearchSourceConnector
untouched.
2026-06-26 18:30:47 +02:00

105 lines
3.9 KiB
Python

"""Invoke SurfSense chat agent for external chat surfaces."""
from __future__ import annotations
import json
import logging
from collections.abc import AsyncIterator
from sqlalchemy import update
from sqlalchemy.ext.asyncio import AsyncSession
from app.auth.context import AuthContext
from app.db import ExternalChatBinding, NewChatMessage
from app.gateway.auth_invariant import assert_authorization_invariant
from app.gateway.base.translator import BaseStreamTranslator, GatewayStreamEvent
from app.gateway.bindings import get_or_create_thread_for_binding
from app.gateway.hitl_filter import DEFAULT_HITL_TOOL_NAMES
from app.gateway.thread_lock import acquire_thread_lock, release_thread_lock
from app.observability.metrics import record_gateway_turn_latency
from app.tasks.chat.streaming.flows import stream_new_chat
logger = logging.getLogger(__name__)
async def _events_from_sse(
chunks: AsyncIterator[str],
) -> AsyncIterator[GatewayStreamEvent]:
saw_text = False
async for chunk in chunks:
for raw_line in chunk.splitlines():
line = raw_line.strip()
if not line.startswith("data:"):
continue
payload = line.removeprefix("data:").strip()
if payload == "[DONE]":
logger.info("Gateway SSE normalized: done")
yield GatewayStreamEvent(type="done")
continue
try:
data = json.loads(payload)
except json.JSONDecodeError:
continue
event_type = str(data.get("type") or "")
if event_type == "text-delta":
delta = data.get("delta", "")
if delta and not saw_text:
logger.info("Gateway SSE normalized: text stream started")
saw_text = True
yield GatewayStreamEvent(type="text-delta", data={"delta": delta})
elif event_type in {"finish", "done"}:
logger.info("Gateway SSE normalized: %s", event_type)
yield GatewayStreamEvent(type="finish", data=data)
elif event_type == "data-interrupt-request":
logger.info("Gateway SSE normalized: interrupt request")
yield GatewayStreamEvent(type="data-interrupt-request", data=data)
async def call_agent_for_gateway(
*,
session: AsyncSession,
binding: ExternalChatBinding,
user_text: str,
translator: BaseStreamTranslator,
platform_label: str = "telegram",
request_id: str | None = None,
) -> None:
user = await assert_authorization_invariant(session, binding)
auth_context = AuthContext.system(user, source="gateway")
thread = await get_or_create_thread_for_binding(session, binding)
await session.commit()
if not acquire_thread_lock(thread.id):
raise RuntimeError("gateway_thread_busy")
try:
stream = stream_new_chat(
user_query=user_text,
workspace_id=binding.workspace_id,
chat_id=thread.id,
user_id=str(user.id),
needs_history_bootstrap=thread.needs_history_bootstrap,
thread_visibility=thread.visibility,
current_user_display_name=user.display_name or "A team member",
disabled_tools=sorted(DEFAULT_HITL_TOOL_NAMES),
request_id=request_id or "gateway",
auth_context=auth_context,
)
events = _events_from_sse(stream)
try:
await translator.translate(events)
finally:
await events.aclose()
await stream.aclose()
await session.execute(
update(NewChatMessage)
.where(
NewChatMessage.thread_id == thread.id,
NewChatMessage.source == "surfsense",
)
.values(source=platform_label)
)
await session.commit()
record_gateway_turn_latency(0, platform=platform_label)
finally:
release_thread_lock(thread.id)