chainlit--chainlit/backend/chainlit/cache.py

import importlib
import importlib.util
import os
import threading
from typing import Any

from chainlit.config import config
from chainlit.logger import logger


def init_lc_cache():
    use_cache = config.project.cache is True and config.run.no_cache is False

    if use_cache and importlib.util.find_spec("langchain") is not None:
        try:
            try:
                set_llm_cache = importlib.import_module(
                    "langchain_core.globals"
                ).set_llm_cache
            except ImportError:
                set_llm_cache = importlib.import_module(
                    "langchain.globals"
                ).set_llm_cache

            try:
                SQLiteCache = importlib.import_module(
                    "langchain_community.cache"
                ).SQLiteCache
            except ImportError:
                SQLiteCache = importlib.import_module("langchain.cache").SQLiteCache
        except (AttributeError, ImportError):
            return

        if config.project.lc_cache_path is not None:
            set_llm_cache(SQLiteCache(database_path=config.project.lc_cache_path))

            if not os.path.exists(config.project.lc_cache_path):
                logger.info(
                    f"LangChain cache created at: {config.project.lc_cache_path}"
                )


_cache: dict[tuple, Any] = {}
_cache_lock = threading.Lock()


def cache(func):
    def wrapper(*args, **kwargs):
        # Create a cache key based on the function name, arguments, and keyword arguments
        cache_key = (
            (func.__name__,) + args + tuple((k, v) for k, v in sorted(kwargs.items()))
        )

        with _cache_lock:
            # Check if the result is already in the cache
            if cache_key not in _cache:
                # If not, call the function and store the result in the cache
                _cache[cache_key] = func(*args, **kwargs)

        return _cache[cache_key]

    return wrapper