diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index a3bdf4f327..c95aad90c3 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -182,7 +182,9 @@ from open_webui.routers.retrieval import ( get_rf, ) from open_webui.socket.main import ( + DIRECT_COMPLETION_RELAY_ENABLED, MODELS, + direct_completion_listener, get_event_emitter, get_models_in_use, get_user_id_from_session_pool, @@ -386,6 +388,9 @@ async def lifespan(app: FastAPI): if app.state.redis is not None: app.state.redis_task_command_listener = asyncio.create_task(redis_task_command_listener(app)) + if DIRECT_COMPLETION_RELAY_ENABLED: + app.state.direct_completion_listener = asyncio.create_task(direct_completion_listener()) + app.state.periodic_usage_pool_cleanup = asyncio.create_task(periodic_usage_pool_cleanup()) app.state.periodic_session_pool_cleanup = asyncio.create_task(periodic_session_pool_cleanup()) @@ -472,6 +477,9 @@ async def lifespan(app: FastAPI): if hasattr(app.state, 'redis_task_command_listener'): app.state.redis_task_command_listener.cancel() + if hasattr(app.state, 'direct_completion_listener'): + app.state.direct_completion_listener.cancel() + app.state.periodic_usage_pool_cleanup.cancel() app.state.periodic_session_pool_cleanup.cancel() app.state.scheduler_worker_loop.cancel() diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index c3678f592d..30ef579c26 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -5,6 +5,7 @@ import logging import random import sys import time +from contextlib import suppress from typing import Any import pycrdt as Y @@ -37,17 +38,23 @@ from open_webui.models.folders import Folders from open_webui.models.notes import Notes, NoteUpdateForm from open_webui.models.users import UserNameResponse, Users from open_webui.socket.utils import RedisDict, RedisLock, YdocManager -from open_webui.tasks import create_task, stop_item_tasks +from open_webui.tasks import ( + REDIS_PUBSUB_MAX_RECONNECT_INTERVAL, + REDIS_PUBSUB_RECONNECT_INTERVAL, + create_task, + stop_item_tasks, +) from open_webui.utils.access_control import has_permission from open_webui.utils.auth import get_verified_user_by_token from open_webui.utils.chat_id import is_saved_chat_id -from open_webui.utils.json_codec import SOCKETIO_JSON +from open_webui.utils.json_codec import SOCKETIO_JSON, JSONCodec, dumps_bytes from open_webui.utils.misc import get_output_text from open_webui.utils.redis import ( build_sentinel_url, get_redis_connection, get_sentinels_from_env, ) +from redis.exceptions import RedisError from socketio.packet import Packet logging.basicConfig(stream=sys.stdout, level=GLOBAL_LOG_LEVEL) @@ -189,6 +196,13 @@ YDOC_MANAGER = YdocManager( redis_key_prefix=f'{REDIS_KEY_PREFIX}:ydoc:documents', ) +REDIS_DIRECT_COMPLETION_CHANNEL = f'{REDIS_KEY_PREFIX}:direct_completion' + +DIRECT_COMPLETION_RELAY_ENABLED = WEBSOCKET_MANAGER == 'redis' + +DIRECT_COMPLETION_QUEUES: dict[str, asyncio.Queue] = {} +DIRECT_COMPLETION_PUBLISH_LOCK = asyncio.Lock() + def get_session_pool_batches(): """All session pool entries, in bounded batches for the Redis backing.""" @@ -943,6 +957,78 @@ async def disconnect(sid, reason=None): # print(f"Unknown session ID {sid} disconnected") +async def direct_completion_listener() -> None: + """Hand every relayed chunk to the queue of the direct completion it belongs to.""" + reconnect_interval = REDIS_PUBSUB_RECONNECT_INTERVAL + + while True: + pubsub = None + try: + # RedisCluster can't route a pubsub subscribe until initialize() fills its slot cache. + await REDIS.initialize() + + pubsub = REDIS.pubsub() + await pubsub.subscribe(REDIS_DIRECT_COMPLETION_CHANNEL) + reconnect_interval = REDIS_PUBSUB_RECONNECT_INTERVAL + + async for message in pubsub.listen(): + if message['type'] != 'message': + continue + chunk = JSONCodec.loads(message['data']) + queue = DIRECT_COMPLETION_QUEUES.get(chunk['channel']) + if queue is not None: + await queue.put(chunk['data']) + log.warning('Direct completion relay listener stopped. Retrying.') + except asyncio.CancelledError: + raise + except Exception: + log.exception('Direct completion relay listener failed. Retrying.') + finally: + if pubsub: + with suppress(Exception): + await pubsub.aclose() + + await asyncio.sleep(reconnect_interval) + reconnect_interval = min(reconnect_interval * 2, REDIS_PUBSUB_MAX_RECONNECT_INTERVAL) + + +def open_direct_completion_queue(channel: str) -> asyncio.Queue: + """Open the queue receiving the client's chunk emits for one direct completion.""" + queue = asyncio.Queue() + DIRECT_COMPLETION_QUEUES[channel] = queue + return queue + + +def close_direct_completion_queue(channel: str) -> None: + DIRECT_COMPLETION_QUEUES.pop(channel, None) + + +async def publish_direct_completion_chunk(channel: str, data: Any) -> None: + """Publish one chunk to the relay channel.""" + async with DIRECT_COMPLETION_PUBLISH_LOCK: + await REDIS.publish(REDIS_DIRECT_COMPLETION_CHANNEL, dumps_bytes({'channel': channel, 'data': data})) + + +@sio.on('*') +async def handle_direct_completion_chunk(event: Any, sid: str, *args: Any) -> None: + """Intake for the chunks a browser streams back for its own direct completion.""" + if not isinstance(event, str) or event.count(':') != 2 or not args: + return + + user = await get_socket_session_user(sid) + if not user or user.get('id') != event.split(':', 1)[0]: + return + + queue = DIRECT_COMPLETION_QUEUES.get(event) + if queue is not None: + await queue.put(args[0]) + elif DIRECT_COMPLETION_RELAY_ENABLED: + try: + await publish_direct_completion_chunk(event, args[0]) + except RedisError as e: + log.debug('Failed to relay direct completion chunk on %s: %s', event, e) + + async def _make_channel_emitter(request_info): """Event emitter that routes pipeline output to a channel message. diff --git a/backend/open_webui/utils/chat.py b/backend/open_webui/utils/chat.py index 0b64a52165..8ae5a69599 100644 --- a/backend/open_webui/utils/chat.py +++ b/backend/open_webui/utils/chat.py @@ -1,4 +1,3 @@ -import asyncio import logging import random import sys @@ -23,9 +22,10 @@ from open_webui.routers.pipelines import ( process_pipeline_outlet_filter, ) from open_webui.socket.main import ( + close_direct_completion_queue, get_event_call, get_event_emitter, - sio, + open_direct_completion_queue, ) from open_webui.utils.filter import ( get_filter_functions, @@ -71,19 +71,7 @@ async def generate_direct_chat_completion( logging.info('WebSocket channel: %s', channel) if form_data.get('stream'): - q = asyncio.Queue() - - async def message_listener(sid, data): - """ - Handle received socket messages and push them into the queue. - """ - await q.put(data) - - def remove_message_listener(): - sio.handlers['/'].pop(channel, None) - - # Register the listener - sio.on(channel, message_listener) + queue = open_direct_completion_queue(channel) # Start processing chat completion in background try: @@ -103,16 +91,15 @@ async def generate_direct_chat_completion( status = res.get('status', False) except BaseException: - remove_message_listener() + close_direct_completion_queue(channel) raise if status: # Define a generator to stream responses async def event_generator(): - nonlocal q try: while True: - data = await q.get() # Wait for new messages + data = await queue.get() # Wait for new messages if isinstance(data, dict): if 'done' in data and data['done']: break # Stop streaming when 'done' is received @@ -127,16 +114,16 @@ async def generate_direct_chat_completion( log.debug('Error in event generator: %s', e) pass finally: - remove_message_listener() + close_direct_completion_queue(channel) # Define a background task to run the event generator async def background(): - remove_message_listener() + close_direct_completion_queue(channel) # Return the streaming response return StreamingResponse(event_generator(), media_type='text/event-stream', background=background) else: - remove_message_listener() + close_direct_completion_queue(channel) raise Exception(str(res)) else: res = await event_caller(