src/backend/base/langflow/services/chat/cache.py

from collections.abc import Awaitable, Callable
from contextlib import contextmanager
from typing import Any

import pandas as pd
from PIL import Image

from langflow.services.base import Service


class Subject:
    """Base class for implementing the observer pattern."""

    def __init__(self) -> None:
        self.observers: list[Callable[[], None]] = []

    def attach(self, observer: Callable[[], None]) -> None:
        """Attach an observer to the subject."""
        self.observers.append(observer)

    def detach(self, observer: Callable[[], None]) -> None:
        """Detach an observer from the subject."""
        self.observers.remove(observer)

    def notify(self) -> None:
        """Notify all observers about an event."""
        for observer in self.observers:
            if observer is None:
                continue
            observer()


class AsyncSubject:
    """Base class for implementing the async observer pattern."""

    def __init__(self) -> None:
        self.observers: list[Callable[[], Awaitable]] = []

    def attach(self, observer: Callable[[], Awaitable]) -> None:
        """Attach an observer to the subject."""
        self.observers.append(observer)

    def detach(self, observer: Callable[[], Awaitable]) -> None:
        """Detach an observer from the subject."""
        self.observers.remove(observer)

    async def notify(self) -> None:
        """Notify all observers about an event."""
        for observer in self.observers:
            if observer is None:
                continue
            await observer()


class CacheService(Subject, Service):
    """Manages cache for different clients and notifies observers on changes."""

    name = "cache_service"

    def __init__(self) -> None:
        super().__init__()
        self._cache: dict[str, Any] = {}
        self.current_client_id: str | None = None
        self.current_cache: dict[str, Any] = {}

    @contextmanager
    def set_client_id(self, client_id: str):
        """Context manager to set the current client_id and associated cache.

        Args:
            client_id (str): The client identifier.
        """
        previous_client_id = self.current_client_id
        self.current_client_id = client_id
        self.current_cache = self._cache.setdefault(client_id, {})
        try:
            yield
        finally:
            self.current_client_id = previous_client_id
            self.current_cache = self._cache.setdefault(previous_client_id, {}) if previous_client_id else {}

    def add(self, name: str, obj: Any, obj_type: str, extension: str | None = None) -> None:
        """Add an object to the current client's cache.

        Args:
            name (str): The cache key.
            obj (Any): The object to cache.
            obj_type (str): The type of the object.
            extension: The file extension of the object.
        """
        object_extensions = {
            "image": "png",
            "pandas": "csv",
        }
        extension_ = object_extensions[obj_type] if obj_type in object_extensions else type(obj).__name__.lower()
        self.current_cache[name] = {
            "obj": obj,
            "type": obj_type,
            "extension": extension or extension_,
        }
        self.notify()

    def add_pandas(self, name: str, obj: Any) -> None:
        """Add a pandas DataFrame or Series to the current client's cache.

        Args:
            name (str): The cache key.
            obj (Any): The pandas DataFrame or Series object.
        """
        if isinstance(obj, pd.DataFrame | pd.Series):
            self.add(name, obj.to_csv(), "pandas", extension="csv")
        else:
            msg = "Object is not a pandas DataFrame or Series"
            raise TypeError(msg)

    def add_image(self, name: str, obj: Any, extension: str = "png") -> None:
        """Add a PIL Image to the current client's cache.

        Args:
            name (str): The cache key.
            obj (Any): The PIL Image object.
            extension: The file extension of the image.
        """
        if isinstance(obj, Image.Image):
            self.add(name, obj, "image", extension=extension)
        else:
            msg = "Object is not a PIL Image"
            raise TypeError(msg)

    def get(self, name: str):
        """Get an object from the current client's cache.

        Args:
            name (str): The cache key.

        Returns:
            The cached object associated with the given cache key.
        """
        return self.current_cache[name]

    def get_last(self):
        """Get the last added item in the current client's cache.

        Returns:
            The last added item in the cache.
        """
        return list(self.current_cache.values())[-1]


cache_service = CacheService()