chore: format

This commit is contained in:
Timothy Jaeryang Baek 2025-03-28 11:47:14 -07:00
parent b8c1bb0ec5
commit 1ac87c55ff
6 changed files with 432 additions and 244 deletions

View File

@ -32,6 +32,7 @@ from open_webui.env import (
from open_webui.internal.db import Base, get_db from open_webui.internal.db import Base, get_db
from open_webui.utils.redis import get_redis_connection from open_webui.utils.redis import get_redis_connection
class EndpointFilter(logging.Filter): class EndpointFilter(logging.Filter):
def filter(self, record: logging.LogRecord) -> bool: def filter(self, record: logging.LogRecord) -> bool:
return record.getMessage().find("/health") == -1 return record.getMessage().find("/health") == -1
@ -254,11 +255,14 @@ class AppConfig:
_state: dict[str, PersistentConfig] _state: dict[str, PersistentConfig]
_redis: Optional[redis.Redis] = None _redis: Optional[redis.Redis] = None
def __init__(self, redis_url: Optional[str] = None, redis_sentinels: Optional[list] = []): def __init__(
self, redis_url: Optional[str] = None, redis_sentinels: Optional[list] = []
):
super().__setattr__("_state", {}) super().__setattr__("_state", {})
if redis_url: if redis_url:
super().__setattr__( super().__setattr__(
"_redis", get_redis_connection(redis_url, redis_sentinels, decode_responses=True) "_redis",
get_redis_connection(redis_url, redis_sentinels, decode_responses=True),
) )
def __setattr__(self, key, value): def __setattr__(self, key, value):

View File

@ -8,7 +8,11 @@ from redis import asyncio as aioredis
from open_webui.models.users import Users, UserNameResponse from open_webui.models.users import Users, UserNameResponse
from open_webui.models.channels import Channels from open_webui.models.channels import Channels
from open_webui.models.chats import Chats from open_webui.models.chats import Chats
from open_webui.utils.redis import parse_redis_sentinel_url, get_sentinels_from_env, AsyncRedisSentinelManager from open_webui.utils.redis import (
parse_redis_sentinel_url,
get_sentinels_from_env,
AsyncRedisSentinelManager,
)
from open_webui.env import ( from open_webui.env import (
ENABLE_WEBSOCKET_SUPPORT, ENABLE_WEBSOCKET_SUPPORT,
@ -35,8 +39,15 @@ log.setLevel(SRC_LOG_LEVELS["SOCKET"])
if WEBSOCKET_MANAGER == "redis": if WEBSOCKET_MANAGER == "redis":
if WEBSOCKET_SENTINEL_HOSTS: if WEBSOCKET_SENTINEL_HOSTS:
redis_config = parse_redis_sentinel_url(WEBSOCKET_REDIS_URL) redis_config = parse_redis_sentinel_url(WEBSOCKET_REDIS_URL)
mgr = AsyncRedisSentinelManager(WEBSOCKET_SENTINEL_HOSTS.split(','), sentinel_port=int(WEBSOCKET_SENTINEL_PORT), redis_port=redis_config["port"], mgr = AsyncRedisSentinelManager(
service=redis_config["service"], db=redis_config["db"], username=redis_config["username"], password=redis_config["password"]) WEBSOCKET_SENTINEL_HOSTS.split(","),
sentinel_port=int(WEBSOCKET_SENTINEL_PORT),
redis_port=redis_config["port"],
service=redis_config["service"],
db=redis_config["db"],
username=redis_config["username"],
password=redis_config["password"],
)
else: else:
mgr = socketio.AsyncRedisManager(WEBSOCKET_REDIS_URL) mgr = socketio.AsyncRedisManager(WEBSOCKET_REDIS_URL)
sio = socketio.AsyncServer( sio = socketio.AsyncServer(
@ -64,10 +75,24 @@ TIMEOUT_DURATION = 3
if WEBSOCKET_MANAGER == "redis": if WEBSOCKET_MANAGER == "redis":
log.debug("Using Redis to manage websockets.") log.debug("Using Redis to manage websockets.")
redis_sentinels=get_sentinels_from_env(WEBSOCKET_SENTINEL_HOSTS, WEBSOCKET_SENTINEL_PORT) redis_sentinels = get_sentinels_from_env(
SESSION_POOL = RedisDict("open-webui:session_pool", redis_url=WEBSOCKET_REDIS_URL, redis_sentinels=redis_sentinels) WEBSOCKET_SENTINEL_HOSTS, WEBSOCKET_SENTINEL_PORT
USER_POOL = RedisDict("open-webui:user_pool", redis_url=WEBSOCKET_REDIS_URL, redis_sentinels=redis_sentinels) )
USAGE_POOL = RedisDict("open-webui:usage_pool", redis_url=WEBSOCKET_REDIS_URL, redis_sentinels=redis_sentinels) SESSION_POOL = RedisDict(
"open-webui:session_pool",
redis_url=WEBSOCKET_REDIS_URL,
redis_sentinels=redis_sentinels,
)
USER_POOL = RedisDict(
"open-webui:user_pool",
redis_url=WEBSOCKET_REDIS_URL,
redis_sentinels=redis_sentinels,
)
USAGE_POOL = RedisDict(
"open-webui:usage_pool",
redis_url=WEBSOCKET_REDIS_URL,
redis_sentinels=redis_sentinels,
)
clean_up_lock = RedisLock( clean_up_lock = RedisLock(
redis_url=WEBSOCKET_REDIS_URL, redis_url=WEBSOCKET_REDIS_URL,

View File

@ -2,13 +2,16 @@ import json
import uuid import uuid
from open_webui.utils.redis import get_redis_connection from open_webui.utils.redis import get_redis_connection
class RedisLock: class RedisLock:
def __init__(self, redis_url, lock_name, timeout_secs, redis_sentinels=[]): def __init__(self, redis_url, lock_name, timeout_secs, redis_sentinels=[]):
self.lock_name = lock_name self.lock_name = lock_name
self.lock_id = str(uuid.uuid4()) self.lock_id = str(uuid.uuid4())
self.timeout_secs = timeout_secs self.timeout_secs = timeout_secs
self.lock_obtained = False self.lock_obtained = False
self.redis = get_redis_connection(redis_url, redis_sentinels, decode_responses=True) self.redis = get_redis_connection(
redis_url, redis_sentinels, decode_responses=True
)
def aquire_lock(self): def aquire_lock(self):
# nx=True will only set this key if it _hasn't_ already been set # nx=True will only set this key if it _hasn't_ already been set
@ -32,7 +35,9 @@ class RedisLock:
class RedisDict: class RedisDict:
def __init__(self, name, redis_url, redis_sentinels=[]): def __init__(self, name, redis_url, redis_sentinels=[]):
self.name = name self.name = name
self.redis = get_redis_connection(redis_url, redis_sentinels, decode_responses=True) self.redis = get_redis_connection(
redis_url, redis_sentinels, decode_responses=True
)
def __setitem__(self, key, value): def __setitem__(self, key, value):
serialized_value = json.dumps(value) serialized_value = json.dumps(value)

View File

@ -3,6 +3,7 @@ import redis
from redis import asyncio as aioredis from redis import asyncio as aioredis
from urllib.parse import urlparse from urllib.parse import urlparse
def parse_redis_sentinel_url(redis_url): def parse_redis_sentinel_url(redis_url):
parsed_url = urlparse(redis_url) parsed_url = urlparse(redis_url)
if parsed_url.scheme != "redis": if parsed_url.scheme != "redis":
@ -11,39 +12,54 @@ def parse_redis_sentinel_url(redis_url):
return { return {
"username": parsed_url.username or None, "username": parsed_url.username or None,
"password": parsed_url.password or None, "password": parsed_url.password or None,
"service": parsed_url.hostname or 'mymaster', "service": parsed_url.hostname or "mymaster",
"port": parsed_url.port or 6379, "port": parsed_url.port or 6379,
"db": int(parsed_url.path.lstrip("/") or 0), "db": int(parsed_url.path.lstrip("/") or 0),
} }
def get_redis_connection(redis_url, redis_sentinels, decode_responses=True): def get_redis_connection(redis_url, redis_sentinels, decode_responses=True):
if redis_sentinels: if redis_sentinels:
redis_config = parse_redis_sentinel_url(redis_url) redis_config = parse_redis_sentinel_url(redis_url)
sentinel = redis.sentinel.Sentinel( sentinel = redis.sentinel.Sentinel(
redis_sentinels, redis_sentinels,
port=redis_config['port'], port=redis_config["port"],
db=redis_config['db'], db=redis_config["db"],
username=redis_config['username'], username=redis_config["username"],
password=redis_config['password'], password=redis_config["password"],
decode_responses=decode_responses decode_responses=decode_responses,
) )
# Get a master connection from Sentinel # Get a master connection from Sentinel
return sentinel.master_for(redis_config['service']) return sentinel.master_for(redis_config["service"])
else: else:
# Standard Redis connection # Standard Redis connection
return redis.Redis.from_url(redis_url, decode_responses=decode_responses) return redis.Redis.from_url(redis_url, decode_responses=decode_responses)
def get_sentinels_from_env(sentinel_hosts_env, sentinel_port_env): def get_sentinels_from_env(sentinel_hosts_env, sentinel_port_env):
if sentinel_hosts_env: if sentinel_hosts_env:
sentinel_hosts=sentinel_hosts_env.split(',') sentinel_hosts = sentinel_hosts_env.split(",")
sentinel_port=int(sentinel_port_env) sentinel_port = int(sentinel_port_env)
return [(host, sentinel_port) for host in sentinel_hosts] return [(host, sentinel_port) for host in sentinel_hosts]
return [] return []
class AsyncRedisSentinelManager(socketio.AsyncRedisManager): class AsyncRedisSentinelManager(socketio.AsyncRedisManager):
def __init__(self, sentinel_hosts, sentinel_port=26379, redis_port=6379, service="mymaster", db=0, def __init__(
username=None, password=None, channel='socketio', write_only=False, logger=None, redis_options=None): self,
sentinel_hosts,
sentinel_port=26379,
redis_port=6379,
service="mymaster",
db=0,
username=None,
password=None,
channel="socketio",
write_only=False,
logger=None,
redis_options=None,
):
""" """
Initialize the Redis Sentinel Manager. Initialize the Redis Sentinel Manager.
This implementation mostly replicates the __init__ of AsyncRedisManager and This implementation mostly replicates the __init__ of AsyncRedisManager and
@ -65,7 +81,7 @@ class AsyncRedisSentinelManager(socketio.AsyncRedisManager):
``aioredis.from_url()``. ``aioredis.from_url()``.
""" """
self._sentinels = [(host, sentinel_port) for host in sentinel_hosts] self._sentinels = [(host, sentinel_port) for host in sentinel_hosts]
self._redis_port=redis_port self._redis_port = redis_port
self._service = service self._service = service
self._db = db self._db = db
self._username = username self._username = username
@ -75,7 +91,9 @@ class AsyncRedisSentinelManager(socketio.AsyncRedisManager):
# connect and call grandparent constructor # connect and call grandparent constructor
self._redis_connect() self._redis_connect()
super(socketio.AsyncRedisManager, self).__init__(channel=channel, write_only=write_only, logger=logger) super(socketio.AsyncRedisManager, self).__init__(
channel=channel, write_only=write_only, logger=logger
)
def _redis_connect(self): def _redis_connect(self):
"""Establish connections to Redis through Sentinel.""" """Establish connections to Redis through Sentinel."""
@ -84,7 +102,7 @@ class AsyncRedisSentinelManager(socketio.AsyncRedisManager):
port=self._redis_port, port=self._redis_port,
db=self._db, db=self._db,
password=self._password, password=self._password,
**self.redis_options **self.redis_options,
) )
self.redis = sentinel.master_for(self._service) self.redis = sentinel.master_for(self._service)

571
package-lock.json generated

File diff suppressed because it is too large Load Diff

View File

@ -1,6 +1,6 @@
{ {
"name": "open-webui", "name": "open-webui",
"version": "0.5.20", "version": "0.6.0",
"private": true, "private": true,
"scripts": { "scripts": {
"dev": "npm run pyodide:fetch && vite dev --host", "dev": "npm run pyodide:fetch && vite dev --host",
@ -80,6 +80,7 @@
"file-saver": "^2.0.5", "file-saver": "^2.0.5",
"fuse.js": "^7.0.0", "fuse.js": "^7.0.0",
"highlight.js": "^11.9.0", "highlight.js": "^11.9.0",
"html-entities": "^2.5.3",
"html2canvas-pro": "^1.5.8", "html2canvas-pro": "^1.5.8",
"i18next": "^23.10.0", "i18next": "^23.10.0",
"i18next-browser-languagedetector": "^7.2.0", "i18next-browser-languagedetector": "^7.2.0",