refac: usage event handling

This commit is contained in:
Timothy Jaeryang Baek 2025-06-16 10:42:34 +04:00
parent deaa7133a2
commit 423a35782b
12 changed files with 219 additions and 152 deletions

View File

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

View File

@ -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
############################

View File

@ -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
############################

View File

@ -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

View File

@ -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;

View File

@ -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;

View File

@ -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

View File

@ -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" />
{#if usage}
{#if usage?.user_ids?.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>
</div>
<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">
{$activeUserIds?.length}
</span>
<div class=" ">
<span class="">
{$i18n.t('Active Users')}:
</span>
<span class=" font-semibold">
{usage?.user_ids?.length}
</span>
</div>
</div>
</div>
</Tooltip>
</Tooltip>
{/if}
{/if}
<!-- <DropdownMenu.Item class="flex items-center py-1.5 px-3 text-sm ">

View File

@ -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';

View File

@ -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';

View File

@ -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']);
});
};

View File

@ -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) => {