feat: websocket

This commit is contained in:
Timothy J. Baek 2024-06-03 23:39:52 -07:00
parent 0495f01acb
commit 85484392b2
7 changed files with 153 additions and 4 deletions

View File

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

View File

@ -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
package-lock.json generated
View File

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

View File

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

View File

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

View File

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

View File

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