UniqAI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2024-11-16 21:42:58 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

refac: active users

Browse Source
This commit is contained in:
Timothy J. Baek 2024-06-07 17:35:01 -07:00
parent 694d1708a3
commit da47c2dfa3
1 changed files with 19 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
backend/apps/socket/main.py
View File

@ -10,7 +10,7 @@ app = socketio.ASGIApp(sio, socketio_path="/ws/socket.io")
# Dictionary to maintain the user pool # Dictionary to maintain the user pool
SESSION_POOL = {}
USER_POOL = {} USER_POOL = {}
USAGE_POOL = {} USAGE_POOL = {}
# Timeout duration in seconds # Timeout duration in seconds
@ -29,7 +29,12 @@ async def connect(sid, environ, auth):
user = Users.get_user_by_id(data["id"]) user = Users.get_user_by_id(data["id"])
if user: if user:
USER_POOL[sid] = user.id SESSION_POOL[sid] = user.id
if user.id in USER_POOL:
USER_POOL[user.id].append(sid)
else:
USER_POOL[user.id] = [sid]
print(f"user {user.name}({user.id}) connected with session ID {sid}") print(f"user {user.name}({user.id}) connected with session ID {sid}")
print(len(set(USER_POOL))) print(len(set(USER_POOL)))
@ -50,7 +55,13 @@ async def user_join(sid, data):
user = Users.get_user_by_id(data["id"]) user = Users.get_user_by_id(data["id"])
if user: if user:
USER_POOL[sid] = user.id
SESSION_POOL[sid] = user.id
if user.id in USER_POOL:
USER_POOL[user.id].append(sid)
else:
USER_POOL[user.id] = [sid]
print(f"user {user.name}({user.id}) connected with session ID {sid}") print(f"user {user.name}({user.id}) connected with session ID {sid}")
print(len(set(USER_POOL))) print(len(set(USER_POOL)))
@ -124,7 +135,11 @@ async def remove_after_timeout(sid, model_id):
@sio.event @sio.event
async def disconnect(sid): async def disconnect(sid):
if sid in USER_POOL: if sid in USER_POOL:
disconnected_user = USER_POOL.pop(sid) disconnected_user = SESSION_POOL.pop(sid)
USER_POOL[disconnected_user].remove(sid)
if len(USER_POOL[disconnected_user]) == 0:
del USER_POOL[disconnected_user]
print(f"user {disconnected_user} disconnected with session ID {sid}") print(f"user {disconnected_user} disconnected with session ID {sid}")
await sio.emit("user-count", {"count": len(USER_POOL)}) await sio.emit("user-count", {"count": len(USER_POOL)})