diff --git a/backend/apps/socket/main.py b/backend/apps/socket/main.py index 58f48025b..13ad15c41 100644 --- a/backend/apps/socket/main.py +++ b/backend/apps/socket/main.py @@ -1,4 +1,6 @@ import socketio +import asyncio + from apps.webui.models.users import Users from utils.utils import decode_token @@ -10,6 +12,9 @@ app = socketio.ASGIApp(sio, socketio_path="/ws/socket.io") USER_POOL = {} +USAGE_POOL = {} +# Timeout duration in seconds +TIMEOUT_DURATION = 3 @sio.event @@ -57,6 +62,66 @@ async def user_count(sid): await sio.emit("user-count", {"count": len(set(USER_POOL))}) +def get_models_in_use(): + # Aggregate all models in use + + models_in_use = [] + for sid, data in USAGE_POOL.items(): + models_in_use.extend(data["models"]) + print(f"Models in use: {models_in_use}") + + return models_in_use + + +@sio.on("usage") +async def usage(sid, data): + print(f'Received "usage" event from {sid}: {data}') + + # Cancel previous task if there is one + if sid in USAGE_POOL: + USAGE_POOL[sid]["task"].cancel() + + # Store the new usage data and task + model_id = data["model"] + + if sid in USAGE_POOL and "models" in USAGE_POOL[sid]: + + print(USAGE_POOL[sid]) + + models = USAGE_POOL[sid]["models"] + if model_id not in models: + models.append(model_id) + USAGE_POOL[sid] = {"models": models} + + else: + USAGE_POOL[sid] = {"models": [model_id]} + + # Schedule a task to remove the usage data after TIMEOUT_DURATION + USAGE_POOL[sid]["task"] = asyncio.create_task(remove_after_timeout(sid, model_id)) + + models_in_use = get_models_in_use() + # Broadcast the usage data to all clients + await sio.emit("usage", {"models": models_in_use}) + + +async def remove_after_timeout(sid, model_id): + try: + await asyncio.sleep(TIMEOUT_DURATION) + if sid in USAGE_POOL: + if model_id in USAGE_POOL[sid]["models"]: + USAGE_POOL[sid]["models"].remove(model_id) + if len(USAGE_POOL[sid]["models"]) == 0: + del USAGE_POOL[sid] + print(f"Removed usage data for {sid} due to timeout") + + models_in_use = get_models_in_use() + # Broadcast the usage data to all clients + await sio.emit("usage", {"models": models_in_use}) + except asyncio.CancelledError: + # Task was cancelled due to new 'usage' event + pass + + @sio.event async def disconnect(sid): if sid in USER_POOL: diff --git a/src/lib/components/chat/Chat.svelte b/src/lib/components/chat/Chat.svelte index bcb5d59a6..37ec7eeca 100644 --- a/src/lib/components/chat/Chat.svelte +++ b/src/lib/components/chat/Chat.svelte @@ -18,7 +18,8 @@ tags as _tags, WEBUI_NAME, banners, - user + user, + socket } from '$lib/stores'; import { convertMessagesToHistory, @@ -280,6 +281,16 @@ } }; + const getChatEventEmitter = async (modelId: string, chatId: string = '') => { + return setInterval(() => { + $socket?.emit('usage', { + action: 'chat', + model: modelId, + chat_id: chatId + }); + }, 1000); + }; + ////////////////////////// // Ollama functions ////////////////////////// @@ -451,6 +462,8 @@ } responseMessage.userContext = userContext; + const chatEventEmitter = await getChatEventEmitter(model.id, _chatId); + if (webSearchEnabled) { await getWebSearchResults(model.id, parentId, responseMessageId); } @@ -460,6 +473,10 @@ } else if (model) { await sendPromptOllama(model, prompt, responseMessageId, _chatId); } + + console.log('chatEventEmitter', chatEventEmitter); + + if (chatEventEmitter) clearInterval(chatEventEmitter); } else { toast.error($i18n.t(`Model {{modelId}} not found`, { modelId })); } @@ -542,6 +559,7 @@ const sendPromptOllama = async (model, userPrompt, responseMessageId, _chatId) => { model = model.id; + const responseMessage = history.messages[responseMessageId]; // Wait until history/message have been updated diff --git a/src/lib/components/common/Tooltip.svelte b/src/lib/components/common/Tooltip.svelte index f19f8450f..0b5d52d25 100644 --- a/src/lib/components/common/Tooltip.svelte +++ b/src/lib/components/common/Tooltip.svelte @@ -21,6 +21,10 @@ touch: touch }); } + } else if (tooltipInstance && content === '') { + if (tooltipInstance) { + tooltipInstance.destroy(); + } } onDestroy(() => { diff --git a/src/lib/components/layout/Sidebar/UserMenu.svelte b/src/lib/components/layout/Sidebar/UserMenu.svelte index 1525727c0..5b445cbdb 100644 --- a/src/lib/components/layout/Sidebar/UserMenu.svelte +++ b/src/lib/components/layout/Sidebar/UserMenu.svelte @@ -5,8 +5,9 @@ import { flyAndScale } from '$lib/utils/transitions'; import { goto } from '$app/navigation'; import ArchiveBox from '$lib/components/icons/ArchiveBox.svelte'; - import { showSettings, activeUserCount } from '$lib/stores'; + import { showSettings, activeUserCount, USAGE_POOL } from '$lib/stores'; import { fade, slide } from 'svelte/transition'; + import Tooltip from '$lib/components/common/Tooltip.svelte'; const i18n = getContext('i18n'); @@ -142,25 +143,31 @@ {#if $activeUserCount}