diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index c060db6de9..29af592458 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -38,7 +38,7 @@ from open_webui.models.chats import Chats from open_webui.models.folders import Folders from open_webui.models.notes import Notes, NoteUpdateForm from open_webui.models.users import UserNameResponse, Users -from open_webui.socket.utils import RedisDict, RedisLock, YdocManager +from open_webui.socket.utils import CachedRedisDict, RedisDict, RedisLock, YdocManager from open_webui.tasks import ( REDIS_PUBSUB_MAX_RECONNECT_INTERVAL, REDIS_PUBSUB_RECONNECT_INTERVAL, @@ -140,12 +140,11 @@ if WEBSOCKET_MANAGER == 'redis': async_mode=True, ) - MODELS = RedisDict( + MODELS = CachedRedisDict( f'{REDIS_KEY_PREFIX}:models', redis_url=WEBSOCKET_REDIS_URL, redis_sentinels=ws_sentinels, redis_cluster=WEBSOCKET_REDIS_CLUSTER, - cache_set_signature=True, ) SESSION_POOL = RedisDict( diff --git a/backend/open_webui/socket/utils.py b/backend/open_webui/socket/utils.py index 049052a1e1..f976a94e77 100644 --- a/backend/open_webui/socket/utils.py +++ b/backend/open_webui/socket/utils.py @@ -134,7 +134,8 @@ class RedisDict: """Delete fields in one HDEL; no keys is a no-op (HDEL rejects an empty field list).""" if keys: self.redis.hdel(self.name, *keys) - self._last_signature = None + if self._signature_name: + self.redis.delete(self._signature_name) def set(self, mapping: dict): if not mapping: @@ -149,10 +150,14 @@ class RedisDict: digest.update(b'\0') digest.update(serialized[key].encode()) digest.update(b'\0') - signature = digest.hexdigest() + content_digest = digest.hexdigest() - if self._signature_name and self.redis.get(self._signature_name) == signature: - return + if self._signature_name: + stored_signature = self.redis.get(self._signature_name) + if stored_signature and stored_signature.startswith(f'{content_digest}:'): + return + # Cleared first so readers refetch while the hash is being rewritten. + self.redis.delete(self._signature_name) # Fetch existing keys before writing so we know which ones to remove. # HKEYS is cheap — it transfers only short key strings, not large JSON values. @@ -168,7 +173,7 @@ class RedisDict: self.redis.hdel(self.name, *keys_to_remove) if self._signature_name: - self.redis.set(self._signature_name, signature) + self.redis.set(self._signature_name, f'{content_digest}:{uuid.uuid4().hex}') def get(self, key, default=None): try: @@ -196,6 +201,43 @@ class RedisDict: return self[key] +class CachedRedisDict(RedisDict): + """Answers reads from a per-worker cache of the hash, refetched whenever its signature changes.""" + + def __init__(self, name: str, redis_url: str, redis_sentinels: list = [], redis_cluster: bool = False): + super().__init__(name, redis_url, redis_sentinels, redis_cluster, cache_set_signature=True) + self._cache: dict = {} + self._cached_signature: str | None = None + + def _refresh_cache(self) -> dict: + stored_signature = self.redis.get(self._signature_name) + if stored_signature is None or stored_signature != self._cached_signature: + self._cache = self.redis.hgetall(self.name) + self._cached_signature = stored_signature + return self._cache + + def __getitem__(self, key): + value = self._refresh_cache().get(key) + if value is None: + raise KeyError(key) + return JSONCodec.loads(value) + + def __contains__(self, key): + return key in self._refresh_cache() + + def __len__(self): + return len(self._refresh_cache()) + + def keys(self): + return list(self._refresh_cache().keys()) + + def values(self): + return [JSONCodec.loads(v) for v in self._refresh_cache().values()] + + def items(self): + return [(k, JSONCodec.loads(v)) for k, v in self._refresh_cache().items()] + + class YdocManager: COMPACTION_THRESHOLD = 500