Merge ae5abfc53c into ee25b7d42d

backend/open_webui/main.py

@@ -183,7 +183,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,
@@ -387,6 +389,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())
 
@@ -473,6 +478,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()

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.
 

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(