Browse Source

refac: usage event handling

Timothy Jaeryang Baek 1 week ago
parent
commit
423a35782b

+ 0 - 1
backend/open_webui/config.py

@@ -3094,4 +3094,3 @@ LDAP_ATTRIBUTE_FOR_GROUPS = PersistentConfig(
     "ldap.server.attribute_for_groups",
     os.environ.get("LDAP_ATTRIBUTE_FOR_GROUPS", "memberOf"),
 )
-

+ 15 - 0
backend/open_webui/main.py

@@ -57,6 +57,8 @@ from open_webui.utils.logger import start_logger
 from open_webui.socket.main import (
     app as socket_app,
     periodic_usage_pool_cleanup,
+    get_models_in_use,
+    get_active_user_ids,
 )
 from open_webui.routers import (
     audio,
@@ -1627,6 +1629,19 @@ async def get_app_changelog():
     return {key: CHANGELOG[key] for idx, key in enumerate(CHANGELOG) if idx < 5}
 
 
+@app.get("/api/usage")
+async def get_current_usage(user=Depends(get_verified_user)):
+    """
+    Get current usage statistics for Open WebUI.
+    This is an experimental endpoint and subject to change.
+    """
+    try:
+        return {"model_ids": get_models_in_use(), "user_ids": get_active_user_ids()}
+    except Exception as e:
+        log.error(f"Error getting usage statistics: {e}")
+        raise HTTPException(status_code=500, detail="Internal Server Error")
+
+
 ############################
 # OAuth Login & Callback
 ############################

+ 35 - 1
backend/open_webui/routers/users.py

@@ -14,7 +14,11 @@ from open_webui.models.users import (
 )
 
 
-from open_webui.socket.main import get_active_status_by_user_id
+from open_webui.socket.main import (
+    get_active_status_by_user_id,
+    get_active_user_ids,
+    get_user_active_status,
+)
 from open_webui.constants import ERROR_MESSAGES
 from open_webui.env import SRC_LOG_LEVELS
 from fastapi import APIRouter, Depends, HTTPException, Request, status
@@ -29,6 +33,24 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"])
 
 router = APIRouter()
 
+
+############################
+# GetActiveUsers
+############################
+
+
+@router.get("/active")
+async def get_active_users(
+    user=Depends(get_verified_user),
+):
+    """
+    Get a list of active users.
+    """
+    return {
+        "user_ids": get_active_user_ids(),
+    }
+
+
 ############################
 # GetUsers
 ############################
@@ -303,6 +325,18 @@ async def get_user_by_id(user_id: str, user=Depends(get_verified_user)):
         )
 
 
+############################
+# GetUserActiveStatusById
+############################
+
+
+@router.get("/{user_id}/active", response_model=dict)
+async def get_user_active_status_by_id(user_id: str, user=Depends(get_verified_user)):
+    return {
+        "active": get_user_active_status(user_id),
+    }
+
+
 ############################
 # UpdateUserById
 ############################

+ 37 - 51
backend/open_webui/socket/main.py

@@ -135,11 +135,6 @@ async def periodic_usage_pool_cleanup():
                     USAGE_POOL[model_id] = connections
 
                 send_usage = True
-
-            if send_usage:
-                # Emit updated usage information after cleaning
-                await sio.emit("usage", {"models": get_models_in_use()})
-
             await asyncio.sleep(TIMEOUT_DURATION)
     finally:
         release_func()
@@ -157,6 +152,43 @@ def get_models_in_use():
     return models_in_use
 
 
+def get_active_user_ids():
+    """Get the list of active user IDs."""
+    return list(USER_POOL.keys())
+
+
+def get_user_active_status(user_id):
+    """Check if a user is currently active."""
+    return user_id in USER_POOL
+
+
+def get_user_id_from_session_pool(sid):
+    user = SESSION_POOL.get(sid)
+    if user:
+        return user["id"]
+    return None
+
+
+def get_user_ids_from_room(room):
+    active_session_ids = sio.manager.get_participants(
+        namespace="/",
+        room=room,
+    )
+
+    active_user_ids = list(
+        set(
+            [SESSION_POOL.get(session_id[0])["id"] for session_id in active_session_ids]
+        )
+    )
+    return active_user_ids
+
+
+def get_active_status_by_user_id(user_id):
+    if user_id in USER_POOL:
+        return True
+    return False
+
+
 @sio.on("usage")
 async def usage(sid, data):
     if sid in SESSION_POOL:
@@ -170,9 +202,6 @@ async def usage(sid, data):
             sid: {"updated_at": current_time},
         }
 
-        # Broadcast the usage data to all clients
-        await sio.emit("usage", {"models": get_models_in_use()})
-
 
 @sio.event
 async def connect(sid, environ, auth):
@@ -190,10 +219,6 @@ async def connect(sid, environ, auth):
             else:
                 USER_POOL[user.id] = [sid]
 
-            # print(f"user {user.name}({user.id}) connected with session ID {sid}")
-            await sio.emit("user-list", {"user_ids": list(USER_POOL.keys())})
-            await sio.emit("usage", {"models": get_models_in_use()})
-
 
 @sio.on("user-join")
 async def user_join(sid, data):
@@ -221,10 +246,6 @@ async def user_join(sid, data):
     log.debug(f"{channels=}")
     for channel in channels:
         await sio.enter_room(sid, f"channel:{channel.id}")
-
-    # print(f"user {user.name}({user.id}) connected with session ID {sid}")
-
-    await sio.emit("user-list", {"user_ids": list(USER_POOL.keys())})
     return {"id": user.id, "name": user.name}
 
 
@@ -277,12 +298,6 @@ async def channel_events(sid, data):
         )
 
 
-@sio.on("user-list")
-async def user_list(sid):
-    if sid in SESSION_POOL:
-        await sio.emit("user-list", {"user_ids": list(USER_POOL.keys())})
-
-
 @sio.event
 async def disconnect(sid):
     if sid in SESSION_POOL:
@@ -294,8 +309,6 @@ async def disconnect(sid):
 
         if len(USER_POOL[user_id]) == 0:
             del USER_POOL[user_id]
-
-        await sio.emit("user-list", {"user_ids": list(USER_POOL.keys())})
     else:
         pass
         # print(f"Unknown session ID {sid} disconnected")
@@ -388,30 +401,3 @@ def get_event_call(request_info):
 
 
 get_event_caller = get_event_call
-
-
-def get_user_id_from_session_pool(sid):
-    user = SESSION_POOL.get(sid)
-    if user:
-        return user["id"]
-    return None
-
-
-def get_user_ids_from_room(room):
-    active_session_ids = sio.manager.get_participants(
-        namespace="/",
-        room=room,
-    )
-
-    active_user_ids = list(
-        set(
-            [SESSION_POOL.get(session_id[0])["id"] for session_id in active_session_ids]
-        )
-    )
-    return active_user_ids
-
-
-def get_active_status_by_user_id(user_id):
-    if user_id in USER_POOL:
-        return True
-    return False

+ 27 - 0
src/lib/apis/index.ts

@@ -1271,6 +1271,33 @@ export const updatePipelineValves = async (
 	return res;
 };
 
+export const getUsage = async (token: string = '') => {
+	let error = null;
+
+	const res = await fetch(`${WEBUI_BASE_URL}/api/usage`, {
+		method: 'GET',
+		headers: {
+			'Content-Type': 'application/json',
+			...(token && { Authorization: `Bearer ${token}` })
+		}
+	})
+		.then(async (res) => {
+			if (!res.ok) throw await res.json();
+			return res.json();
+		})
+		.catch((err) => {
+			console.error(err);
+			error = err;
+			return null;
+		});
+
+	if (error) {
+		throw error;
+	}
+
+	return res;
+};
+
 export const getBackendConfig = async () => {
 	let error = null;
 

+ 27 - 0
src/lib/apis/users/index.ts

@@ -348,6 +348,33 @@ export const getAndUpdateUserLocation = async (token: string) => {
 	}
 };
 
+export const getUserActiveStatusById = async (token: string, userId: string) => {
+	let error = null;
+
+	const res = await fetch(`${WEBUI_API_BASE_URL}/users/${userId}/active`, {
+		method: 'GET',
+		headers: {
+			'Content-Type': 'application/json',
+			Authorization: `Bearer ${token}`
+		}
+	})
+		.then(async (res) => {
+			if (!res.ok) throw await res.json();
+			return res.json();
+		})
+		.catch((err) => {
+			console.error(err);
+			error = err.detail;
+			return null;
+		});
+
+	if (error) {
+		throw error;
+	}
+
+	return res;
+};
+
 export const deleteUserById = async (token: string, userId: string) => {
 	let error = null;
 

+ 20 - 7
src/lib/components/channel/Messages/Message/ProfilePreview.svelte

@@ -1,10 +1,9 @@
 <script lang="ts">
 	import { DropdownMenu } from 'bits-ui';
-	import { createEventDispatcher } from 'svelte';
 
 	import { flyAndScale } from '$lib/utils/transitions';
 	import { WEBUI_BASE_URL } from '$lib/constants';
-	import { activeUserIds } from '$lib/stores';
+	import { getUserActiveStatusById } from '$lib/apis/users';
 
 	export let side = 'right';
 	export let align = 'top';
@@ -12,15 +11,29 @@
 	export let user = null;
 	let show = false;
 
-	const dispatch = createEventDispatcher();
+	let active = false;
+
+	const getActiveStatus = async () => {
+		const res = await getUserActiveStatusById(localStorage.token, user.id).catch((error) => {
+			console.error('Error fetching user active status:', error);
+		});
+
+		if (res) {
+			active = res.active;
+		} else {
+			active = false;
+		}
+	};
+
+	$: if (show) {
+		getActiveStatus();
+	}
 </script>
 
 <DropdownMenu.Root
 	bind:open={show}
 	closeFocus={false}
-	onOpenChange={(state) => {
-		dispatch('change', state);
-	}}
+	onOpenChange={(state) => {}}
 	typeahead={false}
 >
 	<DropdownMenu.Trigger>
@@ -52,7 +65,7 @@
 						</div>
 
 						<div class=" flex items-center gap-2">
-							{#if $activeUserIds.includes(user.id)}
+							{#if active}
 								<div>
 									<span class="relative flex size-2">
 										<span

+ 58 - 29
src/lib/components/layout/Sidebar/UserMenu.svelte

@@ -4,10 +4,14 @@
 
 	import { flyAndScale } from '$lib/utils/transitions';
 	import { goto } from '$app/navigation';
-	import { showSettings, activeUserIds, USAGE_POOL, mobile, showSidebar, user } from '$lib/stores';
 	import { fade, slide } from 'svelte/transition';
-	import Tooltip from '$lib/components/common/Tooltip.svelte';
+
+	import { getUsage } from '$lib/apis';
 	import { userSignOut } from '$lib/apis/auths';
+
+	import { showSettings, mobile, showSidebar, user } from '$lib/stores';
+
+	import Tooltip from '$lib/components/common/Tooltip.svelte';
 	import ArchiveBox from '$lib/components/icons/ArchiveBox.svelte';
 	import QuestionMarkCircle from '$lib/components/icons/QuestionMarkCircle.svelte';
 	import Map from '$lib/components/icons/Map.svelte';
@@ -28,10 +32,28 @@
 	let showShortcuts = false;
 
 	const dispatch = createEventDispatcher();
+
+	let usage = null;
+	const getUsageInfo = async () => {
+		const res = await getUsage(localStorage.token).catch((error) => {
+			console.error('Error fetching usage info:', error);
+		});
+
+		if (res) {
+			usage = res;
+		} else {
+			usage = null;
+		}
+	};
+
+	$: if (show) {
+		getUsageInfo();
+	}
 </script>
 
 <ShortcutsModal bind:show={showShortcuts} />
 
+<!-- svelte-ignore a11y-no-static-element-interactions -->
 <DropdownMenu.Root
 	bind:open={show}
 	onOpenChange={(state) => {
@@ -181,34 +203,41 @@
 				<div class=" self-center truncate">{$i18n.t('Sign Out')}</div>
 			</button>
 
-			{#if $activeUserIds?.length > 0}
-				<hr class=" border-gray-100 dark:border-gray-800 my-1 p-0" />
-
-				<Tooltip
-					content={$USAGE_POOL && $USAGE_POOL.length > 0
-						? `${$i18n.t('Running')}: ${$USAGE_POOL.join(', ')} ✨`
-						: ''}
-				>
-					<div class="flex rounded-md py-1 px-3 text-xs gap-2.5 items-center">
-						<div class=" flex items-center">
-							<span class="relative flex size-2">
-								<span
-									class="animate-ping absolute inline-flex h-full w-full rounded-full bg-green-400 opacity-75"
-								/>
-								<span class="relative inline-flex rounded-full size-2 bg-green-500" />
-							</span>
+			{#if usage}
+				{#if usage?.user_ids?.length > 0}
+					<hr class=" border-gray-100 dark:border-gray-800 my-1 p-0" />
+
+					<Tooltip
+						content={usage?.model_ids && usage?.model_ids.length > 0
+							? `${$i18n.t('Running')}: ${usage.model_ids.join(', ')} ✨`
+							: ''}
+					>
+						<div
+							class="flex rounded-md py-1 px-3 text-xs gap-2.5 items-center"
+							on:mouseenter={() => {
+								getUsageInfo();
+							}}
+						>
+							<div class=" flex items-center">
+								<span class="relative flex size-2">
+									<span
+										class="animate-ping absolute inline-flex h-full w-full rounded-full bg-green-400 opacity-75"
+									/>
+									<span class="relative inline-flex rounded-full size-2 bg-green-500" />
+								</span>
+							</div>
+
+							<div class=" ">
+								<span class="">
+									{$i18n.t('Active Users')}:
+								</span>
+								<span class=" font-semibold">
+									{usage?.user_ids?.length}
+								</span>
+							</div>
 						</div>
-
-						<div class=" ">
-							<span class="">
-								{$i18n.t('Active Users')}:
-							</span>
-							<span class=" font-semibold">
-								{$activeUserIds?.length}
-							</span>
-						</div>
-					</div>
-				</Tooltip>
+					</Tooltip>
+				{/if}
 			{/if}
 
 			<!-- <DropdownMenu.Item class="flex items-center py-1.5 px-3 text-sm ">

+ 1 - 1
src/lib/components/notes/Notes/NoteMenu.svelte

@@ -5,7 +5,7 @@
 	import { flyAndScale } from '$lib/utils/transitions';
 	import { fade, slide } from 'svelte/transition';
 
-	import { showSettings, activeUserIds, USAGE_POOL, mobile, showSidebar, user } from '$lib/stores';
+	import { showSettings, mobile, showSidebar, user } from '$lib/stores';
 
 	import Tooltip from '$lib/components/common/Tooltip.svelte';
 	import ArchiveBox from '$lib/components/icons/ArchiveBox.svelte';

+ 1 - 1
src/lib/components/notes/RecordMenu.svelte

@@ -2,7 +2,7 @@
 	import { DropdownMenu } from 'bits-ui';
 	import { createEventDispatcher, getContext, onMount } from 'svelte';
 
-	import { showSettings, activeUserIds, USAGE_POOL, mobile, showSidebar, user } from '$lib/stores';
+	import { showSettings, mobile, showSidebar, user } from '$lib/stores';
 	import { fade, slide } from 'svelte/transition';
 
 	import Mic from '../icons/Mic.svelte';

+ 0 - 51
src/lib/utils/websocket.ts

@@ -1,51 +0,0 @@
-import { io } from 'socket.io-client';
-
-import { socket, activeUserIds, USAGE_POOL } from '$lib/stores';
-import { WEBUI_BASE_URL } from '$lib/constants';
-
-export const setupSocket = async (enableWebsocket) => {
-	const _socket = io(`${WEBUI_BASE_URL}` || undefined, {
-		reconnection: true,
-		reconnectionDelay: 1000,
-		reconnectionDelayMax: 5000,
-		randomizationFactor: 0.5,
-		path: '/ws/socket.io',
-		transports: enableWebsocket ? ['websocket'] : ['polling', 'websocket'],
-		auth: { token: localStorage.token }
-	});
-
-	await socket.set(_socket);
-
-	_socket.on('connect_error', (err) => {
-		console.log('connect_error', err);
-	});
-
-	_socket.on('connect', () => {
-		console.log('connected', _socket.id);
-	});
-
-	_socket.on('reconnect_attempt', (attempt) => {
-		console.log('reconnect_attempt', attempt);
-	});
-
-	_socket.on('reconnect_failed', () => {
-		console.log('reconnect_failed');
-	});
-
-	_socket.on('disconnect', (reason, details) => {
-		console.log(`Socket ${_socket.id} disconnected due to ${reason}`);
-		if (details) {
-			console.log('Additional details:', details);
-		}
-	});
-
-	_socket.on('user-list', (data) => {
-		console.log('user-list', data);
-		activeUserIds.set(data.user_ids);
-	});
-
-	_socket.on('usage', (data) => {
-		console.log('usage', data);
-		USAGE_POOL.set(data['models']);
-	});
-};

+ 0 - 12
src/routes/+layout.svelte

@@ -16,8 +16,6 @@
 		WEBUI_NAME,
 		mobile,
 		socket,
-		activeUserIds,
-		USAGE_POOL,
 		chatId,
 		chats,
 		currentChatPage,
@@ -103,16 +101,6 @@
 				console.log('Additional details:', details);
 			}
 		});
-
-		_socket.on('user-list', (data) => {
-			console.log('user-list', data);
-			activeUserIds.set(data.user_ids);
-		});
-
-		_socket.on('usage', (data) => {
-			console.log('usage', data);
-			USAGE_POOL.set(data['models']);
-		});
 	};
 
 	const executePythonAsWorker = async (id, code, cb) => {