refactor(cache): move cache-related functions to base.py module

feat(cache): add support for pandas and PIL Image objects caching
fix(interface): import cache-related functions from base.py module
test(cache): update import statements in cache-related test file

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
 

src/backend/langflow/cache/__init__.py

@@ -0,0 +1 @@
+from langflow.cache.base import add_pandas, add_image, get  # noqa

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"]

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

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

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