From ce8fc2739dd9e395ad81001c6b7f9878a00faf06 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Wed, 16 Aug 2023 22:36:34 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20chore(run.py):=20replace=20depre?= =?UTF-8?q?cated=20memoize=5Fdict=20decorator=20with=20Memoize=20class=20t?= =?UTF-8?q?o=20improve=20code=20maintainability=20and=20readability=20?= =?UTF-8?q?=F0=9F=94=A7=20chore(process.py):=20replace=20deprecated=20buil?= =?UTF-8?q?d=5Fsorted=5Fvertices=5Fwith=5Fcaching.hash=20with=20build=5Fso?= =?UTF-8?q?rted=5Fvertices=5Fwith=5Fcaching.session=5Fid=20to=20fix=20inco?= =?UTF-8?q?rrect=20session=5Fid=20assignment=20=F0=9F=94=A7=20chore(utils.?= =?UTF-8?q?py):=20add=20get=5Fcache=5Fmanager=20function=20to=20retrieve?= =?UTF-8?q?=20the=20cache=20manager=20from=20the=20service=20manager?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/backend/langflow/interface/run.py | 7 ++-- src/backend/langflow/processing/process.py | 2 +- src/backend/langflow/services/cache/utils.py | 41 +++++++++++++++++++- src/backend/langflow/services/utils.py | 4 ++ 4 files changed, 49 insertions(+), 5 deletions(-) diff --git a/src/backend/langflow/interface/run.py b/src/backend/langflow/interface/run.py index 42cea0e98..7b7a16b69 100644 --- a/src/backend/langflow/interface/run.py +++ b/src/backend/langflow/interface/run.py @@ -1,10 +1,11 @@ from typing import Any, Dict, Tuple -from langflow.services.cache.utils import memoize_dict +from langflow.services.cache.utils import Memoize from langflow.graph import Graph from langflow.utils.logger import logger +from langflow.services.utils import get_cache_manager -@memoize_dict(maxsize=10) +@Memoize(get_cache_manager=get_cache_manager) def build_langchain_object_with_caching(data_graph): """ Build langchain object from data_graph. @@ -15,7 +16,7 @@ def build_langchain_object_with_caching(data_graph): return graph.build() -@memoize_dict(maxsize=10) +@Memoize(get_cache_manager=get_cache_manager) def build_sorted_vertices_with_caching(data_graph) -> Tuple[Any, Dict]: """ Build langchain object from data_graph. diff --git a/src/backend/langflow/processing/process.py b/src/backend/langflow/processing/process.py index 396135e16..71904f79f 100644 --- a/src/backend/langflow/processing/process.py +++ b/src/backend/langflow/processing/process.py @@ -111,7 +111,7 @@ def load_langchain_object( data_graph: Dict[str, Any], session_id: str ) -> Tuple[Union[Chain, VectorStore], Dict[str, Any], str]: langchain_object, artifacts = get_build_result(data_graph, session_id) - session_id = build_sorted_vertices_with_caching.hash + session_id = build_sorted_vertices_with_caching.session_id logger.debug("Loaded LangChain object") if langchain_object is None: diff --git a/src/backend/langflow/services/cache/utils.py b/src/backend/langflow/services/cache/utils.py index 2333eb5f4..13e24dd6d 100644 --- a/src/backend/langflow/services/cache/utils.py +++ b/src/backend/langflow/services/cache/utils.py @@ -7,9 +7,12 @@ import os import tempfile from collections import OrderedDict from pathlib import Path -from typing import Any, Dict +from typing import TYPE_CHECKING, Any, Callable, Dict from appdirs import user_cache_dir +if TYPE_CHECKING: + from langflow.services.cache.base import BaseCacheManager + CACHE: Dict[str, Any] = {} CACHE_DIR = user_cache_dir("langflow", "langflow") @@ -191,3 +194,39 @@ def save_uploaded_file(file, folder_name): new_file.write(chunk) return file_path + + +class Memoize: + def __init__( + self, + get_cache_manager: Callable[[], "BaseCacheManager"], + ): + self.get_cache_manager = get_cache_manager + self.hash_func = compute_dict_hash + + def clear_cache(self, session_id): + cache_manager = self.get_cache_manager() + cache_manager.delete(session_id) + + def get_result_by_session_id(self, session_id): + cache_manager = self.get_cache_manager() + return cache_manager.get(session_id) + + def __call__(self, func: Callable[..., Any]): + @functools.wraps(func) + def wrapper(*args, **kwargs): + cache_manager = self.get_cache_manager() + session_id = self.hash_func(args[0]) + + result = cache_manager.get(session_id) + + if result is None: + result = func(*args, **kwargs) + cache_manager.set(session_id, result) + + wrapper.session_id = session_id + return result + + wrapper.clear_cache = self.clear_cache + wrapper.get_result_by_session_id = self.get_result_by_session_id + return wrapper diff --git a/src/backend/langflow/services/utils.py b/src/backend/langflow/services/utils.py index 049e82c0f..1c5796fb4 100644 --- a/src/backend/langflow/services/utils.py +++ b/src/backend/langflow/services/utils.py @@ -16,3 +16,7 @@ def get_db_manager(): def get_session(): db_manager = service_manager.get(ServiceType.DATABASE_MANAGER) yield from db_manager.get_session() + + +def get_cache_manager(): + return service_manager.get(ServiceType.CACHE_MANAGER)