Browse Source

refac: async redis

Timothy Jaeryang Baek 2 weeks ago
parent
commit
db3c26ab7a
2 changed files with 39 additions and 21 deletions
  1. 1 0
      backend/open_webui/main.py
  2. 38 21
      backend/open_webui/utils/redis.py

+ 1 - 0
backend/open_webui/main.py

@@ -510,6 +510,7 @@ async def lifespan(app: FastAPI):
         redis_sentinels=get_sentinels_from_env(
             REDIS_SENTINEL_HOSTS, REDIS_SENTINEL_PORT
         ),
+        async_mode=True,
     )
 
     if isinstance(app.state.redis, Redis):

+ 38 - 21
backend/open_webui/utils/redis.py

@@ -1,6 +1,4 @@
 import socketio
-import redis
-from redis import asyncio as aioredis
 from urllib.parse import urlparse
 from typing import Optional
 
@@ -20,26 +18,45 @@ def parse_redis_service_url(redis_url):
 
 
 def get_redis_connection(
-    redis_url, redis_sentinels, decode_responses=True
-) -> Optional[redis.Redis]:
-    if redis_sentinels:
-        redis_config = parse_redis_service_url(redis_url)
-        sentinel = redis.sentinel.Sentinel(
-            redis_sentinels,
-            port=redis_config["port"],
-            db=redis_config["db"],
-            username=redis_config["username"],
-            password=redis_config["password"],
-            decode_responses=decode_responses,
-        )
-
-        # Get a master connection from Sentinel
-        return sentinel.master_for(redis_config["service"])
-    elif redis_url:
-        # Standard Redis connection
-        return redis.Redis.from_url(redis_url, decode_responses=decode_responses)
+    redis_url, redis_sentinels, async_mode=False, decode_responses=True
+):
+    if async_mode:
+        import redis.asyncio as redis
+
+        # If using sentinel in async mode
+        if redis_sentinels:
+            redis_config = parse_redis_service_url(redis_url)
+            sentinel = redis.sentinel.Sentinel(
+                redis_sentinels,
+                port=redis_config["port"],
+                db=redis_config["db"],
+                username=redis_config["username"],
+                password=redis_config["password"],
+                decode_responses=decode_responses,
+            )
+            return sentinel.master_for(redis_config["service"])
+        elif redis_url:
+            return redis.from_url(redis_url, decode_responses=decode_responses)
+        else:
+            return None
     else:
-        return None
+        import redis
+
+        if redis_sentinels:
+            redis_config = parse_redis_service_url(redis_url)
+            sentinel = redis.sentinel.Sentinel(
+                redis_sentinels,
+                port=redis_config["port"],
+                db=redis_config["db"],
+                username=redis_config["username"],
+                password=redis_config["password"],
+                decode_responses=decode_responses,
+            )
+            return sentinel.master_for(redis_config["service"])
+        elif redis_url:
+            return redis.Redis.from_url(redis_url, decode_responses=decode_responses)
+        else:
+            return None
 
 
 def get_sentinels_from_env(sentinel_hosts_env, sentinel_port_env):