diff --git a/src/backend/langflow/api/base.py b/src/backend/langflow/api/base.py index 084e04d65..8cddc52e4 100644 --- a/src/backend/langflow/api/base.py +++ b/src/backend/langflow/api/base.py @@ -3,6 +3,10 @@ from pydantic import BaseModel, validator from langflow.graph.utils import extract_input_variables_from_prompt +class CacheResponse(BaseModel): + data: dict + + class Code(BaseModel): code: str diff --git a/src/backend/langflow/cache/__init__.py b/src/backend/langflow/cache/__init__.py index e69de29bb..f7aac380b 100644 --- a/src/backend/langflow/cache/__init__.py +++ b/src/backend/langflow/cache/__init__.py @@ -0,0 +1 @@ +from langflow.cache.base import add_pandas, add_image, get # noqa diff --git a/src/backend/langflow/cache/utils.py b/src/backend/langflow/cache/base.py similarity index 80% rename from src/backend/langflow/cache/utils.py rename to src/backend/langflow/cache/base.py index 310f3be80..d96614218 100644 --- a/src/backend/langflow/cache/utils.py +++ b/src/backend/langflow/cache/base.py @@ -1,14 +1,18 @@ -import base64 import contextlib import functools import hashlib + import json import os import tempfile from collections import OrderedDict from pathlib import Path +from typing import Any +from PIL import Image +import dill +import pandas as pd # type: ignore -import dill # type: ignore +CACHE = {} def create_cache_folder(func): @@ -147,3 +151,33 @@ def load_cache(hash_val): with cache_path.open("rb") as cache_file: return dill.load(cache_file) return None + + +def add_pandas(name: str, obj: Any): + if isinstance(obj, (pd.DataFrame, pd.Series)): + CACHE[name] = {"obj": obj, "type": "pandas"} + else: + raise ValueError("Object is not a pandas DataFrame or Series") + + +def add_image(name: str, obj: Any): + if isinstance(obj, Image.Image): + CACHE[name] = {"obj": obj, "type": "image"} + else: + raise ValueError("Object is not a PIL Image") + + +def get(name: str): + return CACHE.get(name, {}).get("obj", None) + + +# get last added item +def get_last(): + obj_dict = list(CACHE.values())[-1] + if obj_dict["type"] == "pandas": + # return a csv string + return obj_dict["obj"].to_csv() + elif obj_dict["type"] == "image": + # return a base64 encoded string + return base64.b64encode(obj_dict["obj"].tobytes()).decode("utf-8") + return obj_dict["obj"] diff --git a/src/backend/langflow/graph/base.py b/src/backend/langflow/graph/base.py index ff586c6da..012250739 100644 --- a/src/backend/langflow/graph/base.py +++ b/src/backend/langflow/graph/base.py @@ -9,7 +9,7 @@ import warnings from copy import deepcopy from typing import Any, Dict, List, Optional -from langflow.cache import utils as cache_utils +from langflow.cache import base as cache_utils from langflow.graph.constants import DIRECT_TYPES from langflow.interface import loading from langflow.interface.listing import ALL_TYPES_DICT diff --git a/src/backend/langflow/interface/run.py b/src/backend/langflow/interface/run.py index 300e09e01..89fd7d784 100644 --- a/src/backend/langflow/interface/run.py +++ b/src/backend/langflow/interface/run.py @@ -2,7 +2,7 @@ import contextlib import io from typing import Any, Dict -from langflow.cache.utils import compute_dict_hash, load_cache, memoize_dict +from langflow.cache.base import compute_dict_hash, load_cache, memoize_dict from langflow.graph.graph import Graph from langflow.interface import loading from langflow.utils.logger import logger diff --git a/tests/test_cache.py b/tests/test_cache.py index 9c6ad30e3..131e015f3 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -3,7 +3,7 @@ import tempfile from pathlib import Path import pytest -from langflow.cache.utils import PREFIX, save_cache +from langflow.cache.base import PREFIX, save_cache from langflow.interface.run import load_langchain_object