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)