diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index a38acc0c00..9eda1d55c3 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -286,14 +286,14 @@ def get_session_ids_from_room(room): def get_user_ids_from_room(room): active_session_ids = get_session_ids_from_room(room) + # Single pool lookup per session (each .get is a Redis round trip + # when the session pool is Redis-backed). active_user_ids = list( - set( - [ - SESSION_POOL.get(session_id)['id'] - for session_id in active_session_ids - if SESSION_POOL.get(session_id) is not None - ] - ) + { + entry['id'] + for entry in (SESSION_POOL.get(session_id) for session_id in active_session_ids) + if entry is not None + } ) return active_user_ids @@ -357,7 +357,7 @@ async def usage(sid, data): # Store the new usage data and task USAGE_POOL[model_id] = { - **(USAGE_POOL[model_id] if model_id in USAGE_POOL else {}), + **(USAGE_POOL.get(model_id) or {}), sid: {'updated_at': current_time}, } diff --git a/backend/open_webui/utils/chat.py b/backend/open_webui/utils/chat.py index 196689f61d..2850c92b93 100644 --- a/backend/open_webui/utils/chat.py +++ b/backend/open_webui/utils/chat.py @@ -179,8 +179,10 @@ async def generate_chat_completion( # Merge the direct connection model into server models so that # task functions (title, tags, etc.) can resolve a server-side # task model while still having the direct model available. + # dict(...items()) is one HGETALL on a Redis-backed pool; ``{**pool}`` + # would issue HKEYS plus one HGET per model. models = { - **request.app.state.MODELS, + **dict(request.app.state.MODELS.items()), request.state.model['id']: request.state.model, } log.debug(f'direct connection to model: {request.state.model["id"]}') @@ -188,11 +190,12 @@ async def generate_chat_completion( models = request.app.state.MODELS model_id = form_data['model'] - if model_id not in models: + # Single lookup — membership check plus getitem would be two Redis + # round trips on a Redis-backed model pool. + model = models.get(model_id) + if model is None: raise Exception('Model not found') - model = models[model_id] - if getattr(request.state, 'direct', False) and model_id == getattr(request.state, 'model', {}).get('id'): return await generate_direct_chat_completion(request, form_data, user=user, models=models) else: