From 88f24d292dc5f4ce665301d1f8e93769c837e85a Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Wed, 27 Sep 2023 12:36:32 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(cache/manager.py):=20handle?= =?UTF-8?q?=20pickled=20values=20in=20get=20method=20to=20correctly=20retu?= =?UTF-8?q?rn=20the=20value=20=E2=9C=A8=20feat(cache/manager.py):=20add=20?= =?UTF-8?q?support=20for=20pickling=20values=20in=20set=20method=20to=20mi?= =?UTF-8?q?mic=20Redis=20behavior?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/backend/langflow/services/cache/manager.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/src/backend/langflow/services/cache/manager.py b/src/backend/langflow/services/cache/manager.py index 19f5034cf..da76a2b5c 100644 --- a/src/backend/langflow/services/cache/manager.py +++ b/src/backend/langflow/services/cache/manager.py @@ -74,13 +74,17 @@ class InMemoryCache(BaseCacheService, Service): ): # Move the key to the end to make it recently used self._cache.move_to_end(key) - unpickled = pickle.loads(item["value"]) - return unpickled + # Check if the value is pickled + if isinstance(item["value"], bytes): + value = pickle.loads(item["value"]) + else: + value = item["value"] + return value else: self.delete(key) return None - def set(self, key, value): + def set(self, key, value, pickle=False): """ Add an item to the cache. @@ -98,8 +102,10 @@ class InMemoryCache(BaseCacheService, Service): # Remove least recently used item self._cache.popitem(last=False) # pickle locally to mimic Redis - pickled = pickle.dumps(value) - self._cache[key] = {"value": pickled, "time": time.time()} + if pickle: + value = pickle.dumps(value) + + self._cache[key] = {"value": value, "time": time.time()} def upsert(self, key, value): """