scriptoriumadmin/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-08-28 05:27:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

refac

Browse source
This commit is contained in:
Timothy Jaeryang Baek 2026-08-23 13:40:13 -04:00
parent 5093a99389
commit bf3a58dbcd
1 changed files with 34 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

46
backend/open_webui/tasks.py
View file

@ -1,6 +1,7 @@
# tasks.py
import asyncio
import logging
from contextlib import suppress
from uuid import uuid4
from redis.asyncio import Redis
@ -20,25 +21,46 @@ REDIS_TASKS_KEY = f'{REDIS_KEY_PREFIX}:tasks'
REDIS_ITEM_TASKS_KEY = f'{REDIS_KEY_PREFIX}:tasks:item'
REDIS_RESPONSE_STREAMS_KEY = f'{REDIS_KEY_PREFIX}:tasks:response_streams'
REDIS_PUBSUB_CHANNEL = f'{REDIS_KEY_PREFIX}:tasks:commands'
REDIS_PUBSUB_RECONNECT_INTERVAL = 1.0
REDIS_PUBSUB_MAX_RECONNECT_INTERVAL = 30.0
async def redis_task_command_listener(app):
redis: Redis = app.state.redis
pubsub = redis.pubsub()
await pubsub.subscribe(REDIS_PUBSUB_CHANNEL)
reconnect_interval = REDIS_PUBSUB_RECONNECT_INTERVAL
async for message in pubsub.listen():
if message['type'] != 'message':
continue
while True:
pubsub = None
try:
command = JSONCodec.loads(message['data'])
if command.get('action') == 'stop':
task_id = command.get('task_id')
local_task = tasks.get(task_id)
if local_task:
local_task.cancel()
pubsub = redis.pubsub()
await pubsub.subscribe(REDIS_PUBSUB_CHANNEL)
reconnect_interval = REDIS_PUBSUB_RECONNECT_INTERVAL
async for message in pubsub.listen():
if message['type'] != 'message':
continue
try:
command = JSONCodec.loads(message['data'])
if command.get('action') != 'stop':
continue
local_task = tasks.get(command.get('task_id'))
if local_task:
local_task.cancel()
except Exception as e:
log.exception(f'Error handling distributed task command: {e}')
log.warning('Redis task command listener stopped. Retrying.')
except asyncio.CancelledError:
raise
except Exception as e:
log.exception(f'Error handling distributed task command: {e}')
log.exception(f'Redis task command listener failed. Retrying: {e}')
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)
### ------------------------------