|
@@ -3,16 +3,20 @@ import socketio
|
|
import logging
|
|
import logging
|
|
import sys
|
|
import sys
|
|
import time
|
|
import time
|
|
|
|
+from redis import asyncio as aioredis
|
|
|
|
|
|
from open_webui.models.users import Users, UserNameResponse
|
|
from open_webui.models.users import Users, UserNameResponse
|
|
from open_webui.models.channels import Channels
|
|
from open_webui.models.channels import Channels
|
|
from open_webui.models.chats import Chats
|
|
from open_webui.models.chats import Chats
|
|
|
|
+from open_webui.utils.redis import parse_redis_sentinel_url, get_sentinels_from_env, AsyncRedisSentinelManager
|
|
|
|
|
|
from open_webui.env import (
|
|
from open_webui.env import (
|
|
ENABLE_WEBSOCKET_SUPPORT,
|
|
ENABLE_WEBSOCKET_SUPPORT,
|
|
WEBSOCKET_MANAGER,
|
|
WEBSOCKET_MANAGER,
|
|
WEBSOCKET_REDIS_URL,
|
|
WEBSOCKET_REDIS_URL,
|
|
WEBSOCKET_REDIS_LOCK_TIMEOUT,
|
|
WEBSOCKET_REDIS_LOCK_TIMEOUT,
|
|
|
|
+ WEBSOCKET_SENTINEL_PORT,
|
|
|
|
+ WEBSOCKET_SENTINEL_HOSTS,
|
|
)
|
|
)
|
|
from open_webui.utils.auth import decode_token
|
|
from open_webui.utils.auth import decode_token
|
|
from open_webui.socket.utils import RedisDict, RedisLock
|
|
from open_webui.socket.utils import RedisDict, RedisLock
|
|
@@ -29,7 +33,12 @@ log.setLevel(SRC_LOG_LEVELS["SOCKET"])
|
|
|
|
|
|
|
|
|
|
if WEBSOCKET_MANAGER == "redis":
|
|
if WEBSOCKET_MANAGER == "redis":
|
|
- mgr = socketio.AsyncRedisManager(WEBSOCKET_REDIS_URL)
|
|
|
|
|
|
+ if WEBSOCKET_SENTINEL_HOSTS:
|
|
|
|
+ redis_config = parse_redis_sentinel_url(WEBSOCKET_REDIS_URL)
|
|
|
|
+ mgr = AsyncRedisSentinelManager(WEBSOCKET_SENTINEL_HOSTS.split(','), sentinel_port=int(WEBSOCKET_SENTINEL_PORT), redis_port=redis_config["port"],
|
|
|
|
+ service=redis_config["service"], db=redis_config["db"], username=redis_config["username"], password=redis_config["password"])
|
|
|
|
+ else:
|
|
|
|
+ mgr = socketio.AsyncRedisManager(WEBSOCKET_REDIS_URL)
|
|
sio = socketio.AsyncServer(
|
|
sio = socketio.AsyncServer(
|
|
cors_allowed_origins=[],
|
|
cors_allowed_origins=[],
|
|
async_mode="asgi",
|
|
async_mode="asgi",
|
|
@@ -55,14 +64,16 @@ TIMEOUT_DURATION = 3
|
|
|
|
|
|
if WEBSOCKET_MANAGER == "redis":
|
|
if WEBSOCKET_MANAGER == "redis":
|
|
log.debug("Using Redis to manage websockets.")
|
|
log.debug("Using Redis to manage websockets.")
|
|
- SESSION_POOL = RedisDict("open-webui:session_pool", redis_url=WEBSOCKET_REDIS_URL)
|
|
|
|
- USER_POOL = RedisDict("open-webui:user_pool", redis_url=WEBSOCKET_REDIS_URL)
|
|
|
|
- USAGE_POOL = RedisDict("open-webui:usage_pool", redis_url=WEBSOCKET_REDIS_URL)
|
|
|
|
|
|
+ redis_sentinels=get_sentinels_from_env(WEBSOCKET_SENTINEL_HOSTS, WEBSOCKET_SENTINEL_PORT)
|
|
|
|
+ SESSION_POOL = RedisDict("open-webui:session_pool", redis_url=WEBSOCKET_REDIS_URL, redis_sentinels=redis_sentinels)
|
|
|
|
+ USER_POOL = RedisDict("open-webui:user_pool", redis_url=WEBSOCKET_REDIS_URL, redis_sentinels=redis_sentinels)
|
|
|
|
+ USAGE_POOL = RedisDict("open-webui:usage_pool", redis_url=WEBSOCKET_REDIS_URL, redis_sentinels=redis_sentinels)
|
|
|
|
|
|
clean_up_lock = RedisLock(
|
|
clean_up_lock = RedisLock(
|
|
redis_url=WEBSOCKET_REDIS_URL,
|
|
redis_url=WEBSOCKET_REDIS_URL,
|
|
lock_name="usage_cleanup_lock",
|
|
lock_name="usage_cleanup_lock",
|
|
timeout_secs=WEBSOCKET_REDIS_LOCK_TIMEOUT,
|
|
timeout_secs=WEBSOCKET_REDIS_LOCK_TIMEOUT,
|
|
|
|
+ redis_sentinels=redis_sentinels,
|
|
)
|
|
)
|
|
aquire_func = clean_up_lock.aquire_lock
|
|
aquire_func = clean_up_lock.aquire_lock
|
|
renew_func = clean_up_lock.renew_lock
|
|
renew_func = clean_up_lock.renew_lock
|