mi
99 lines
3.5 KiB
Python
99 lines
3.5 KiB
Python
import asyncio
|
|
import json
|
|
import uuid
|
|
from collections.abc import AsyncIterator
|
|
from contextlib import suppress
|
|
from typing import Any
|
|
|
|
from redis.asyncio import Redis
|
|
|
|
DIALOG_CHANNEL_PREFIX = "han:rt:dialog:"
|
|
USER_CHANNEL_PREFIX = "han:rt:user:"
|
|
# Backward-compatible name used by existing chat integrations.
|
|
CHANNEL_PREFIX = DIALOG_CHANNEL_PREFIX
|
|
|
|
|
|
class LocalFanout:
|
|
def __init__(self) -> None:
|
|
self._queues: set[asyncio.Queue[dict[str, Any]]] = set()
|
|
|
|
async def publish(self, event: dict[str, Any]) -> None:
|
|
for queue in tuple(self._queues):
|
|
with suppress(asyncio.QueueFull):
|
|
queue.put_nowait(event)
|
|
|
|
async def subscribe(self) -> AsyncIterator[dict[str, Any]]:
|
|
queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue(maxsize=256)
|
|
self._queues.add(queue)
|
|
try:
|
|
while True:
|
|
yield await queue.get()
|
|
finally:
|
|
self._queues.discard(queue)
|
|
|
|
|
|
class RealtimeFanout:
|
|
def __init__(self, redis: Redis, local: LocalFanout | None = None) -> None:
|
|
self.redis = redis
|
|
self.local = local or LocalFanout()
|
|
|
|
async def publish(self, event: dict[str, Any]) -> None:
|
|
event = {"event_id": str(uuid.uuid4()), **event}
|
|
channel = DIALOG_CHANNEL_PREFIX + str(event["dialog_id"])
|
|
try:
|
|
await self.redis.publish(channel, json.dumps(event, default=str, separators=(",", ":")))
|
|
except Exception:
|
|
await self.local.publish(event)
|
|
|
|
async def publish_user(self, user_id: uuid.UUID, event: dict[str, Any]) -> None:
|
|
event = {"event_id": str(uuid.uuid4()), "_user_id": str(user_id), **event}
|
|
channel = USER_CHANNEL_PREFIX + str(user_id)
|
|
try:
|
|
await self.redis.publish(channel, json.dumps(event, default=str, separators=(",", ":")))
|
|
except Exception:
|
|
await self.local.publish(event)
|
|
|
|
async def events(
|
|
self,
|
|
dialog_ids: set[uuid.UUID],
|
|
user_id: uuid.UUID | None = None,
|
|
notifications: bool = False,
|
|
) -> AsyncIterator[dict[str, Any]]:
|
|
channels = [DIALOG_CHANNEL_PREFIX + str(dialog_id) for dialog_id in dialog_ids]
|
|
if notifications and user_id is not None:
|
|
channels.append(USER_CHANNEL_PREFIX + str(user_id))
|
|
if not channels:
|
|
await asyncio.Event().wait()
|
|
return
|
|
pubsub = self.redis.pubsub()
|
|
try:
|
|
await pubsub.subscribe(*channels)
|
|
except Exception:
|
|
await pubsub.aclose()
|
|
async for event in self.local.subscribe():
|
|
dialog_match = event.get("dialog_id") and uuid.UUID(
|
|
str(event["dialog_id"])
|
|
) in dialog_ids
|
|
user_match = (
|
|
notifications
|
|
and user_id is not None
|
|
and event.get("_user_id") == str(user_id)
|
|
)
|
|
if dialog_match or user_match:
|
|
payload = dict(event)
|
|
payload.pop("_user_id", None)
|
|
yield payload
|
|
return
|
|
try:
|
|
while True:
|
|
message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=1)
|
|
if message:
|
|
payload = json.loads(message["data"])
|
|
payload.pop("_user_id", None)
|
|
yield payload
|
|
else:
|
|
await asyncio.sleep(0)
|
|
finally:
|
|
await pubsub.unsubscribe(*channels)
|
|
await pubsub.aclose()
|