enh/refac: redis cluster support

backend/open_webui/socket/main.py

@@ -22,6 +22,7 @@ from open_webui.env import (
     ENABLE_WEBSOCKET_SUPPORT,
     WEBSOCKET_MANAGER,
     WEBSOCKET_REDIS_URL,
+    WEBSOCKET_REDIS_CLUSTER,
     WEBSOCKET_REDIS_LOCK_TIMEOUT,
     WEBSOCKET_SENTINEL_PORT,
     WEBSOCKET_SENTINEL_HOSTS,
@@ -86,6 +87,7 @@ if WEBSOCKET_MANAGER == "redis":
         redis_sentinels=get_sentinels_from_env(
             WEBSOCKET_SENTINEL_HOSTS, WEBSOCKET_SENTINEL_PORT
         ),
+        redis_cluster=WEBSOCKET_REDIS_CLUSTER,
         async_mode=True,
     )
 
@@ -96,16 +98,19 @@ if WEBSOCKET_MANAGER == "redis":
         f"{REDIS_KEY_PREFIX}:session_pool",
         redis_url=WEBSOCKET_REDIS_URL,
         redis_sentinels=redis_sentinels,
+        redis_cluster=WEBSOCKET_REDIS_CLUSTER,
     )
     USER_POOL = RedisDict(
         f"{REDIS_KEY_PREFIX}:user_pool",
         redis_url=WEBSOCKET_REDIS_URL,
         redis_sentinels=redis_sentinels,
+        redis_cluster=WEBSOCKET_REDIS_CLUSTER,
     )
     USAGE_POOL = RedisDict(
         f"{REDIS_KEY_PREFIX}:usage_pool",
         redis_url=WEBSOCKET_REDIS_URL,
         redis_sentinels=redis_sentinels,
+        redis_cluster=WEBSOCKET_REDIS_CLUSTER,
     )
 
     clean_up_lock = RedisLock(
@@ -113,6 +118,7 @@ if WEBSOCKET_MANAGER == "redis":
         lock_name="usage_cleanup_lock",
         timeout_secs=WEBSOCKET_REDIS_LOCK_TIMEOUT,
         redis_sentinels=redis_sentinels,
+        redis_cluster=WEBSOCKET_REDIS_CLUSTER,
     )
     aquire_func = clean_up_lock.aquire_lock
     renew_func = clean_up_lock.renew_lock

backend/open_webui/socket/utils.py

@@ -7,13 +7,24 @@ import pycrdt as Y
 
 
 class RedisLock:
-    def __init__(self, redis_url, lock_name, timeout_secs, redis_sentinels=[]):
+    def __init__(
+        self,
+        redis_url,
+        lock_name,
+        timeout_secs,
+        redis_sentinels=[],
+        redis_cluster=False,
+    ):
+
         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
+            redis_url,
+            redis_sentinels,
+            redis_cluster=redis_cluster,
+            decode_responses=True,
         )
 
     def aquire_lock(self):
@@ -36,10 +47,13 @@ class RedisLock:
 
 
 class RedisDict:
-    def __init__(self, name, redis_url, redis_sentinels=[]):
+    def __init__(self, name, redis_url, redis_sentinels=[], redis_cluster=False):
         self.name = name
         self.redis = get_redis_connection(
-            redis_url, redis_sentinels, decode_responses=True
+            redis_url,
+            redis_sentinels,
+            redis_cluster=redis_cluster,
+            decode_responses=True,
         )
 
     def __setitem__(self, key, value):