From a44e9a8dda7213bcdddc2a7f006f3bedca3c0fb3 Mon Sep 17 00:00:00 2001 From: "Timothy J. Baek" Date: Tue, 24 Sep 2024 15:41:23 +0200 Subject: [PATCH] refac --- backend/open_webui/apps/socket/main.py | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/backend/open_webui/apps/socket/main.py b/backend/open_webui/apps/socket/main.py index ac80e6259..1c9ff505c 100644 --- a/backend/open_webui/apps/socket/main.py +++ b/backend/open_webui/apps/socket/main.py @@ -48,6 +48,18 @@ else: ) +# Dictionary to maintain the user pool + +if WEBSOCKET_MANAGER == "redis": + 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) +else: + SESSION_POOL = {} + USER_POOL = {} + USAGE_POOL = {} + + # Timeout duration in seconds TIMEOUT_DURATION = 3 @@ -84,17 +96,6 @@ app = socketio.ASGIApp( on_startup=asyncio.create_task(periodic_usage_pool_cleanup()), ) -# Dictionary to maintain the user pool - -if WEBSOCKET_MANAGER == "redis": - 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) -else: - SESSION_POOL = {} - USER_POOL = {} - USAGE_POOL = {} - def get_models_in_use(): # List models that are currently in use