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

chore: format

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2025-03-28 11:47:14 -07:00
parent b8c1bb0ec5
commit 1ac87c55ff
6 changed files with 432 additions and 244 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -8,7 +8,11 @@ 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.utils.redis import (
parse_redis_sentinel_url,
get_sentinels_from_env,
AsyncRedisSentinelManager,
)
from open_webui.env import (
ENABLE_WEBSOCKET_SUPPORT,
@@ -35,8 +39,15 @@ log.setLevel(SRC_LOG_LEVELS["SOCKET"])
if WEBSOCKET_MANAGER == "redis":
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"])
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(
@@ -64,10 +75,24 @@ TIMEOUT_DURATION = 3
if WEBSOCKET_MANAGER == "redis":
log.debug("Using Redis to manage websockets.")
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)
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,

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

@@ -2,13 +2,16 @@ import json
import uuid
from open_webui.utils.redis import get_redis_connection
class RedisLock:
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 = get_redis_connection(redis_url, redis_sentinels, 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
@@ -32,7 +35,9 @@ class RedisLock:
class RedisDict:
def __init__(self, name, redis_url, redis_sentinels=[]):
self.name = name
self.redis = get_redis_connection(redis_url, redis_sentinels, 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)