OpenWebUI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-06-26 18:26:48 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge pull request #11148 from Ithanil/redis_sentinel

Browse Source 
feat: Support for Redis Sentinel
This commit is contained in:
Timothy Jaeryang Baek
2025-03-27 20:36:53 -07:00
committed by GitHub
parent c137d3ad17 0615c11a53
commit 7fce3e2a9c
6 changed files with 128 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
backend/open_webui/socket/main.py
View File

@@ -3,16 +3,20 @@ import socketio
import logging
import sys
import time
from redis import asyncio as aioredis
from open_webui.models.users import Users, UserNameResponse
from open_webui.models.channels import Channels
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 (
ENABLE_WEBSOCKET_SUPPORT,
WEBSOCKET_MANAGER,
WEBSOCKET_REDIS_URL,
WEBSOCKET_REDIS_LOCK_TIMEOUT,
WEBSOCKET_SENTINEL_PORT,
WEBSOCKET_SENTINEL_HOSTS,
)
from open_webui.utils.auth import decode_token
from open_webui.socket.utils import RedisDict, RedisLock
@@ -29,7 +33,12 @@ log.setLevel(SRC_LOG_LEVELS["SOCKET"])
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(
cors_allowed_origins=[],
async_mode="asgi",
@@ -55,14 +64,16 @@ TIMEOUT_DURATION = 3
if WEBSOCKET_MANAGER == "redis":
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(
redis_url=WEBSOCKET_REDIS_URL,
lock_name="usage_cleanup_lock",
timeout_secs=WEBSOCKET_REDIS_LOCK_TIMEOUT,
redis_sentinels=redis_sentinels,
)
aquire_func = clean_up_lock.aquire_lock
renew_func = clean_up_lock.renew_lock

11
backend/open_webui/socket/utils.py
View File

@@ -1,15 +1,14 @@
import json
import redis
import uuid
from open_webui.utils.redis import get_redis_connection
class RedisLock:
def __init__(self, redis_url, lock_name, timeout_secs):
def __init__(self, redis_url, lock_name, timeout_secs, redis_sentinels=[]):
self.lock_name = lock_name
self.lock_id = str(uuid.uuid4())
self.timeout_secs = timeout_secs
self.lock_obtained = False
self.redis = redis.Redis.from_url(redis_url, decode_responses=True)
self.redis = get_redis_connection(redis_url, redis_sentinels, decode_responses=True)
def aquire_lock(self):
# nx=True will only set this key if it _hasn't_ already been set
@@ -31,9 +30,9 @@ class RedisLock:
class RedisDict:
def __init__(self, name, redis_url):
def __init__(self, name, redis_url, redis_sentinels=[]):
self.name = name
self.redis = redis.Redis.from_url(redis_url, decode_responses=True)
self.redis = get_redis_connection(redis_url, redis_sentinels, decode_responses=True)
def __setitem__(self, key, value):
serialized_value = json.dumps(value)