feat: websocket

backend/apps/socket/main.py

@@ -0,0 +1,47 @@
+import socketio
+
+from apps.webui.models.users import Users
+from utils.utils import decode_token
+
+sio = socketio.AsyncServer(cors_allowed_origins=[], async_mode="asgi")
+app = socketio.ASGIApp(sio, socketio_path="/ws/socket.io")
+
+# Dictionary to maintain the user pool
+USER_POOL = {}
+
+
+@sio.event
+async def connect(sid, environ, auth):
+    print("connect ", sid)
+
+    user = None
+    data = decode_token(auth["token"])
+
+    if data is not None and "id" in data:
+        user = Users.get_user_by_id(data["id"])
+
+    if user:
+        USER_POOL[sid] = {
+            "id": user.id,
+            "name": user.name,
+            "email": user.email,
+            "role": user.role,
+        }
+        print(f"user {user.name}({user.id}) connected with session ID {sid}")
+    else:
+        print("Authentication failed. Disconnecting.")
+        await sio.disconnect(sid)
+
+
+@sio.event
+def disconnect(sid):
+    if sid in USER_POOL:
+        disconnected_user = USER_POOL.pop(sid)
+        print(f"user {disconnected_user} disconnected with session ID {sid}")
+    else:
+        print(f"Unknown session ID {sid} disconnected")
+
+
+@sio.event
+def disconnect(sid):
+    print("disconnect", sid)

backend/main.py

@@ -20,6 +20,8 @@ from starlette.exceptions import HTTPException as StarletteHTTPException
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.responses import StreamingResponse, Response
 
+
+from apps.socket.main import app as socket_app
 from apps.ollama.main import app as ollama_app, get_all_models as get_ollama_models
 from apps.openai.main import app as openai_app, get_all_models as get_openai_models
 
@@ -376,6 +378,9 @@ async def update_embedding_function(request: Request, call_next):
     return response
 
 
+app.mount("/ws", socket_app)
+
+
 app.mount("/ollama", ollama_app)
 app.mount("/openai", openai_app)