mirror of
https://github.com/langgenius/dify.git
synced 2026-10-04 13:47:55 +08:00
58 lines
2.3 KiB
Python
58 lines
2.3 KiB
Python
"""Session progress bus over the Redis broadcast channel.
|
|
|
|
The Celery advance task publishes progress events here and the POST-SSE
|
|
endpoints subscribe to the same topic. A streaming request must establish its
|
|
delivery boundary before dispatching the task; merely constructing the Redis
|
|
subscription is not sufficient because subscriptions activate lazily.
|
|
|
|
Uses ``extensions.ext_redis.get_pubsub_broadcast_channel()`` rather than
|
|
constructing ``StreamsBroadcastChannel`` directly: that helper builds the
|
|
channel on the raw pubsub client (avoiding a double key-prefix bug when
|
|
``REDIS_KEY_PREFIX`` is set) and respects the deployment's configured
|
|
``PUBSUB_REDIS_CHANNEL_TYPE`` (streams / sharded / pubsub), so publishers
|
|
and subscribers always agree on channel type and keys. It asserts the
|
|
pubsub client is initialized, so it must be called lazily inside
|
|
``publish``/``subscribe`` rather than at module import time.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
from extensions.ext_redis import get_pubsub_broadcast_channel
|
|
from libs.broadcast_channel.channel import Subscription, SupportsPreparedSubscription
|
|
|
|
_TOPIC_FMT = "dify_builder:{session_id}"
|
|
|
|
|
|
def _topic(session_id: str) -> str:
|
|
return _TOPIC_FMT.format(session_id=session_id)
|
|
|
|
|
|
def publish(session_id: str, event: dict[str, Any]) -> None:
|
|
"""Publish ``event`` to ``session_id``'s progress topic, JSON-encoded."""
|
|
get_pubsub_broadcast_channel().topic(_topic(session_id)).publish(json.dumps(event).encode())
|
|
|
|
|
|
def subscribe(session_id: str) -> Subscription:
|
|
"""Establish a lossless delivery boundary for ``session_id`` progress.
|
|
|
|
Redis Streams can fix the boundary without starting their listener, so use
|
|
``prepare_subscription`` when the topic exposes that capability. Redis
|
|
Pub/Sub has no replay boundary and must be entered eagerly so its SUBSCRIBE
|
|
command completes before the Celery task can publish. The SSE generator
|
|
consumes the returned subscription as-is and owns closing it.
|
|
"""
|
|
subscriber = get_pubsub_broadcast_channel().topic(_topic(session_id)).as_subscriber()
|
|
if isinstance(subscriber, SupportsPreparedSubscription):
|
|
return subscriber.prepare_subscription()
|
|
|
|
subscription = subscriber.subscribe()
|
|
try:
|
|
subscription.__enter__()
|
|
except BaseException:
|
|
subscription.close()
|
|
raise
|
|
return subscription
|