github-starred/open-webui
2
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-05-04 03:16:03 -05:00
Code Issues 915 Packages Projects Releases 126 Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2026-03-17 17:58:01 -05:00
parent fcf7208352
commit de3317e26b
220 changed files with 17200 additions and 22836 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
backend/open_webui/socket/utils.py
View File

@@ -15,7 +15,6 @@ class RedisLock:
redis_sentinels=[],
redis_cluster=False,
):
self.lock_name = lock_name
self.lock_id = str(uuid.uuid4())
self.timeout_secs = timeout_secs
@@ -29,16 +28,12 @@ class RedisLock:
def aquire_lock(self):
# nx=True will only set this key if it _hasn't_ already been set
self.lock_obtained = self.redis.set(
self.lock_name, self.lock_id, nx=True, ex=self.timeout_secs
)
self.lock_obtained = self.redis.set(self.lock_name, self.lock_id, nx=True, ex=self.timeout_secs)
return self.lock_obtained
def renew_lock(self):
# xx=True will only set this key if it _has_ already been set
return self.redis.set(
self.lock_name, self.lock_id, xx=True, ex=self.timeout_secs
)
return self.redis.set(self.lock_name, self.lock_id, xx=True, ex=self.timeout_secs)
def release_lock(self):
lock_value = self.redis.get(self.lock_name)
@@ -106,7 +101,7 @@ class RedisDict:
def update(self, other=None, **kwargs):
if other is not None:
for k, v in other.items() if hasattr(other, "items") else other:
for k, v in other.items() if hasattr(other, 'items') else other:
self[k] = v
for k, v in kwargs.items():
self[k] = v
@@ -123,7 +118,7 @@ class YdocManager:
def __init__(
self,
redis=None,
redis_key_prefix: str = f"{REDIS_KEY_PREFIX}:ydoc:documents",
redis_key_prefix: str = f'{REDIS_KEY_PREFIX}:ydoc:documents',
):
self._updates = {}
self._users = {}
@@ -131,9 +126,9 @@ class YdocManager:
self._redis_key_prefix = redis_key_prefix
async def append_to_updates(self, document_id: str, update: bytes):
document_id = document_id.replace(":", "_")
document_id = document_id.replace(':', '_')
if self._redis:
redis_key = f"{self._redis_key_prefix}:{document_id}:updates"
redis_key = f'{self._redis_key_prefix}:{document_id}:updates'
await self._redis.rpush(redis_key, json.dumps(list(update)))
list_len = await self._redis.llen(redis_key)
if list_len >= self.COMPACTION_THRESHOLD:
@@ -147,7 +142,7 @@ class YdocManager:
async def _compact_updates_redis(self, document_id: str):
"""Rolling compaction: squash oldest half into one snapshot."""
redis_key = f"{self._redis_key_prefix}:{document_id}:updates"
redis_key = f'{self._redis_key_prefix}:{document_id}:updates'
all_updates = await self._redis.lrange(redis_key, 0, -1)
if len(all_updates) <= 1:
return
@@ -173,39 +168,39 @@ class YdocManager:
self._updates[document_id] = [ydoc.get_update()] + updates[mid:]
async def get_updates(self, document_id: str) -> List[bytes]:
document_id = document_id.replace(":", "_")
document_id = document_id.replace(':', '_')
if self._redis:
redis_key = f"{self._redis_key_prefix}:{document_id}:updates"
redis_key = f'{self._redis_key_prefix}:{document_id}:updates'
updates = await self._redis.lrange(redis_key, 0, -1)
return [bytes(json.loads(update)) for update in updates]
else:
return self._updates.get(document_id, [])
async def document_exists(self, document_id: str) -> bool:
document_id = document_id.replace(":", "_")
document_id = document_id.replace(':', '_')
if self._redis:
redis_key = f"{self._redis_key_prefix}:{document_id}:updates"
redis_key = f'{self._redis_key_prefix}:{document_id}:updates'
return await self._redis.exists(redis_key) > 0
else:
return document_id in self._updates
async def get_users(self, document_id: str) -> List[str]:
document_id = document_id.replace(":", "_")
document_id = document_id.replace(':', '_')
if self._redis:
redis_key = f"{self._redis_key_prefix}:{document_id}:users"
redis_key = f'{self._redis_key_prefix}:{document_id}:users'
users = await self._redis.smembers(redis_key)
return list(users)
else:
return self._users.get(document_id, [])
async def add_user(self, document_id: str, user_id: str):
document_id = document_id.replace(":", "_")
document_id = document_id.replace(':', '_')
if self._redis:
redis_key = f"{self._redis_key_prefix}:{document_id}:users"
redis_key = f'{self._redis_key_prefix}:{document_id}:users'
await self._redis.sadd(redis_key, user_id)
else:
if document_id not in self._users:
@@ -213,10 +208,10 @@ class YdocManager:
self._users[document_id].add(user_id)
async def remove_user(self, document_id: str, user_id: str):
document_id = document_id.replace(":", "_")
document_id = document_id.replace(':', '_')
if self._redis:
redis_key = f"{self._redis_key_prefix}:{document_id}:users"
redis_key = f'{self._redis_key_prefix}:{document_id}:users'
await self._redis.srem(redis_key, user_id)
else:
if document_id in self._users and user_id in self._users[document_id]:
@@ -225,15 +220,13 @@ class YdocManager:
async def remove_user_from_all_documents(self, user_id: str):
if self._redis:
keys = []
async for key in self._redis.scan_iter(
match=f"{self._redis_key_prefix}:*", count=100
):
async for key in self._redis.scan_iter(match=f'{self._redis_key_prefix}:*', count=100):
keys.append(key)
for key in keys:
if key.endswith(":users"):
if key.endswith(':users'):
await self._redis.srem(key, user_id)
document_id = key.split(":")[-2]
document_id = key.split(':')[-2]
if len(await self.get_users(document_id)) == 0:
await self.clear_document(document_id)
@@ -247,12 +240,12 @@ class YdocManager:
await self.clear_document(document_id)
async def clear_document(self, document_id: str):
document_id = document_id.replace(":", "_")
document_id = document_id.replace(':', '_')
if self._redis:
redis_key = f"{self._redis_key_prefix}:{document_id}:updates"
redis_key = f'{self._redis_key_prefix}:{document_id}:updates'
await self._redis.delete(redis_key)
redis_users_key = f"{self._redis_key_prefix}:{document_id}:users"
redis_users_key = f'{self._redis_key_prefix}:{document_id}:users'
await self._redis.delete(redis_users_key)
else:
if document_id in self._updates: