浏览代码

feat: websocket

Timothy J. Baek 11 月之前
父节点
当前提交
85484392b2
共有 7 个文件被更改,包括 153 次插入4 次删除
  1. 47 0
      backend/apps/socket/main.py
  2. 5 0
      backend/main.py
  3. 80 0
      package-lock.json
  4. 1 0
      package.json
  5. 2 1
      src/lib/constants.ts
  6. 3 0
      src/lib/stores/index.ts
  7. 15 3
      src/routes/+layout.svelte

+ 47 - 0
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)

+ 5 - 0
backend/main.py

@@ -20,6 +20,8 @@ from starlette.exceptions import HTTPException as StarletteHTTPException
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.responses import StreamingResponse, Response
 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.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
 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
     return response
 
 
 
 
+app.mount("/ws", socket_app)
+
+
 app.mount("/ollama", ollama_app)
 app.mount("/ollama", ollama_app)
 app.mount("/openai", openai_app)
 app.mount("/openai", openai_app)
 
 

+ 80 - 0
package-lock.json

@@ -25,6 +25,7 @@
 				"marked": "^9.1.0",
 				"marked": "^9.1.0",
 				"mermaid": "^10.9.1",
 				"mermaid": "^10.9.1",
 				"pyodide": "^0.26.0-alpha.4",
 				"pyodide": "^0.26.0-alpha.4",
+				"socket.io-client": "^4.7.5",
 				"sortablejs": "^1.15.2",
 				"sortablejs": "^1.15.2",
 				"svelte-sonner": "^0.3.19",
 				"svelte-sonner": "^0.3.19",
 				"tippy.js": "^6.3.7",
 				"tippy.js": "^6.3.7",
@@ -1214,6 +1215,11 @@
 			"integrity": "sha512-+Fj43pSMwJs4KRrH/938Uf+uAELIgVBmQzg/q1YG10djyfA3TnrU8N8XzqCh/okZdszqBQTZf96idMfE5lnwTA==",
 			"integrity": "sha512-+Fj43pSMwJs4KRrH/938Uf+uAELIgVBmQzg/q1YG10djyfA3TnrU8N8XzqCh/okZdszqBQTZf96idMfE5lnwTA==",
 			"dev": true
 			"dev": true
 		},
 		},
+		"node_modules/@socket.io/component-emitter": {
+			"version": "3.1.2",
+			"resolved": "https://registry.npmjs.org/@socket.io/component-emitter/-/component-emitter-3.1.2.tgz",
+			"integrity": "sha512-9BCxFwvbGg/RsZK9tjXd8s4UcwR0MWeFQ1XEKIQVVvAGJyINdrqKMcTRyLoK8Rse1GjzLV9cwjWV1olXRWEXVA=="
+		},
 		"node_modules/@sveltejs/adapter-auto": {
 		"node_modules/@sveltejs/adapter-auto": {
 			"version": "2.1.1",
 			"version": "2.1.1",
 			"resolved": "https://registry.npmjs.org/@sveltejs/adapter-auto/-/adapter-auto-2.1.1.tgz",
 			"resolved": "https://registry.npmjs.org/@sveltejs/adapter-auto/-/adapter-auto-2.1.1.tgz",
@@ -3800,6 +3806,46 @@
 				"once": "^1.4.0"
 				"once": "^1.4.0"
 			}
 			}
 		},
 		},
+		"node_modules/engine.io-client": {
+			"version": "6.5.3",
+			"resolved": "https://registry.npmjs.org/engine.io-client/-/engine.io-client-6.5.3.tgz",
+			"integrity": "sha512-9Z0qLB0NIisTRt1DZ/8U2k12RJn8yls/nXMZLn+/N8hANT3TcYjKFKcwbw5zFQiN4NTde3TSY9zb79e1ij6j9Q==",
+			"dependencies": {
+				"@socket.io/component-emitter": "~3.1.0",
+				"debug": "~4.3.1",
+				"engine.io-parser": "~5.2.1",
+				"ws": "~8.11.0",
+				"xmlhttprequest-ssl": "~2.0.0"
+			}
+		},
+		"node_modules/engine.io-client/node_modules/ws": {
+			"version": "8.11.0",
+			"resolved": "https://registry.npmjs.org/ws/-/ws-8.11.0.tgz",
+			"integrity": "sha512-HPG3wQd9sNQoT9xHyNCXoDUa+Xw/VevmY9FoHyQ+g+rrMn4j6FB4np7Z0OhdTgjx6MgQLK7jwSy1YecU1+4Asg==",
+			"engines": {
+				"node": ">=10.0.0"
+			},
+			"peerDependencies": {
+				"bufferutil": "^4.0.1",
+				"utf-8-validate": "^5.0.2"
+			},
+			"peerDependenciesMeta": {
+				"bufferutil": {
+					"optional": true
+				},
+				"utf-8-validate": {
+					"optional": true
+				}
+			}
+		},
+		"node_modules/engine.io-parser": {
+			"version": "5.2.2",
+			"resolved": "https://registry.npmjs.org/engine.io-parser/-/engine.io-parser-5.2.2.tgz",
+			"integrity": "sha512-RcyUFKA93/CXH20l4SoVvzZfrSDMOTUS3bWVpTt2FuFP+XYrL8i8oonHP7WInRyVHXh0n/ORtoeiE1os+8qkSw==",
+			"engines": {
+				"node": ">=10.0.0"
+			}
+		},
 		"node_modules/enquirer": {
 		"node_modules/enquirer": {
 			"version": "2.4.1",
 			"version": "2.4.1",
 			"resolved": "https://registry.npmjs.org/enquirer/-/enquirer-2.4.1.tgz",
 			"resolved": "https://registry.npmjs.org/enquirer/-/enquirer-2.4.1.tgz",
@@ -7949,6 +7995,32 @@
 				"node": ">=8"
 				"node": ">=8"
 			}
 			}
 		},
 		},
+		"node_modules/socket.io-client": {
+			"version": "4.7.5",
+			"resolved": "https://registry.npmjs.org/socket.io-client/-/socket.io-client-4.7.5.tgz",
+			"integrity": "sha512-sJ/tqHOCe7Z50JCBCXrsY3I2k03iOiUe+tj1OmKeD2lXPiGH/RUCdTZFoqVyN7l1MnpIzPrGtLcijffmeouNlQ==",
+			"dependencies": {
+				"@socket.io/component-emitter": "~3.1.0",
+				"debug": "~4.3.2",
+				"engine.io-client": "~6.5.2",
+				"socket.io-parser": "~4.2.4"
+			},
+			"engines": {
+				"node": ">=10.0.0"
+			}
+		},
+		"node_modules/socket.io-parser": {
+			"version": "4.2.4",
+			"resolved": "https://registry.npmjs.org/socket.io-parser/-/socket.io-parser-4.2.4.tgz",
+			"integrity": "sha512-/GbIKmo8ioc+NIWIhwdecY0ge+qVBSMdgxGygevmdHj24bsfgtCmcUUcQ5ZzcylGFHsN3k4HB4Cgkl96KVnuew==",
+			"dependencies": {
+				"@socket.io/component-emitter": "~3.1.0",
+				"debug": "~4.3.1"
+			},
+			"engines": {
+				"node": ">=10.0.0"
+			}
+		},
 		"node_modules/sorcery": {
 		"node_modules/sorcery": {
 			"version": "0.11.0",
 			"version": "0.11.0",
 			"resolved": "https://registry.npmjs.org/sorcery/-/sorcery-0.11.0.tgz",
 			"resolved": "https://registry.npmjs.org/sorcery/-/sorcery-0.11.0.tgz",
@@ -10142,6 +10214,14 @@
 				}
 				}
 			}
 			}
 		},
 		},
+		"node_modules/xmlhttprequest-ssl": {
+			"version": "2.0.0",
+			"resolved": "https://registry.npmjs.org/xmlhttprequest-ssl/-/xmlhttprequest-ssl-2.0.0.tgz",
+			"integrity": "sha512-QKxVRxiRACQcVuQEYFsI1hhkrMlrXHPegbbd1yn9UHOmRxY+si12nQYzri3vbzt8VdTTRviqcKxcyllFas5z2A==",
+			"engines": {
+				"node": ">=0.4.0"
+			}
+		},
 		"node_modules/xtend": {
 		"node_modules/xtend": {
 			"version": "4.0.2",
 			"version": "4.0.2",
 			"resolved": "https://registry.npmjs.org/xtend/-/xtend-4.0.2.tgz",
 			"resolved": "https://registry.npmjs.org/xtend/-/xtend-4.0.2.tgz",

+ 1 - 0
package.json

@@ -65,6 +65,7 @@
 		"marked": "^9.1.0",
 		"marked": "^9.1.0",
 		"mermaid": "^10.9.1",
 		"mermaid": "^10.9.1",
 		"pyodide": "^0.26.0-alpha.4",
 		"pyodide": "^0.26.0-alpha.4",
+		"socket.io-client": "^4.7.5",
 		"sortablejs": "^1.15.2",
 		"sortablejs": "^1.15.2",
 		"svelte-sonner": "^0.3.19",
 		"svelte-sonner": "^0.3.19",
 		"tippy.js": "^6.3.7",
 		"tippy.js": "^6.3.7",

+ 2 - 1
src/lib/constants.ts

@@ -2,8 +2,9 @@ import { browser, dev } from '$app/environment';
 // import { version } from '../../package.json';
 // import { version } from '../../package.json';
 
 
 export const APP_NAME = 'Open WebUI';
 export const APP_NAME = 'Open WebUI';
-export const WEBUI_BASE_URL = browser ? (dev ? `http://${location.hostname}:8080` : ``) : ``;
 
 
+export const WEBUI_HOSTNAME = browser ? (dev ? `${location.hostname}:8080` : ``) : '';
+export const WEBUI_BASE_URL = browser ? (dev ? `http://${WEBUI_HOSTNAME}` : ``) : ``;
 export const WEBUI_API_BASE_URL = `${WEBUI_BASE_URL}/api/v1`;
 export const WEBUI_API_BASE_URL = `${WEBUI_BASE_URL}/api/v1`;
 
 
 export const OLLAMA_API_BASE_URL = `${WEBUI_BASE_URL}/ollama`;
 export const OLLAMA_API_BASE_URL = `${WEBUI_BASE_URL}/ollama`;

+ 3 - 0
src/lib/stores/index.ts

@@ -2,6 +2,7 @@ import { APP_NAME } from '$lib/constants';
 import { type Writable, writable } from 'svelte/store';
 import { type Writable, writable } from 'svelte/store';
 import type { GlobalModelConfig, ModelConfig } from '$lib/apis';
 import type { GlobalModelConfig, ModelConfig } from '$lib/apis';
 import type { Banner } from '$lib/types';
 import type { Banner } from '$lib/types';
+import type { Socket } from 'socket.io-client';
 
 
 // Backend
 // Backend
 export const WEBUI_NAME = writable(APP_NAME);
 export const WEBUI_NAME = writable(APP_NAME);
@@ -13,6 +14,8 @@ export const MODEL_DOWNLOAD_POOL = writable({});
 
 
 export const mobile = writable(false);
 export const mobile = writable(false);
 
 
+export const socket: Writable<null | Socket> = writable(null);
+
 export const theme = writable('system');
 export const theme = writable('system');
 export const chatId = writable('');
 export const chatId = writable('');
 
 

+ 15 - 3
src/routes/+layout.svelte

@@ -1,6 +1,8 @@
 <script>
 <script>
+	import { io } from 'socket.io-client';
+
 	import { onMount, tick, setContext } from 'svelte';
 	import { onMount, tick, setContext } from 'svelte';
-	import { config, user, theme, WEBUI_NAME, mobile } from '$lib/stores';
+	import { config, user, theme, WEBUI_NAME, mobile, socket } from '$lib/stores';
 	import { goto } from '$app/navigation';
 	import { goto } from '$app/navigation';
 	import { Toaster, toast } from 'svelte-sonner';
 	import { Toaster, toast } from 'svelte-sonner';
 
 
@@ -12,7 +14,7 @@
 
 
 	import 'tippy.js/dist/tippy.css';
 	import 'tippy.js/dist/tippy.css';
 
 
-	import { WEBUI_BASE_URL } from '$lib/constants';
+	import { WEBUI_BASE_URL, WEBUI_HOSTNAME } from '$lib/constants';
 	import i18n, { initI18n, getLanguages } from '$lib/i18n';
 	import i18n, { initI18n, getLanguages } from '$lib/i18n';
 
 
 	setContext('i18n', i18n);
 	setContext('i18n', i18n);
@@ -55,10 +57,20 @@
 		if (backendConfig) {
 		if (backendConfig) {
 			// Save Backend Status to Store
 			// Save Backend Status to Store
 			await config.set(backendConfig);
 			await config.set(backendConfig);
-
 			await WEBUI_NAME.set(backendConfig.name);
 			await WEBUI_NAME.set(backendConfig.name);
 
 
 			if ($config) {
 			if ($config) {
+				const _socket = io(`${WEBUI_BASE_URL}`, {
+					path: '/ws/socket.io',
+					auth: { token: localStorage.token }
+				});
+
+				_socket.on('connect', () => {
+					console.log('connected');
+				});
+
+				socket.set(_socket);
+
 				if (localStorage.token) {
 				if (localStorage.token) {
 					// Get Session User Info
 					// Get Session User Info
 					const sessionUser = await getSessionUser(localStorage.token).catch((error) => {
 					const sessionUser = await getSessionUser(localStorage.token).catch((error) => {