import json from hashlib import md5 from pathlib import Path from typing import Any, Callable, Generator, no_type_check class Cache: """ Decorator class that adds caching functionality to a function. """ def __init__(self, length: int, cache_path: Path) -> None: """ Initialize the Cache decorator. :param length: Integer, maximum number of cache files to keep. """ self.length = length self.cache_path = cache_path self.cache_path.mkdir(parents=True, exist_ok=True) def __call__(self, func: Callable[..., Any]) -> Callable[..., Any]: """ The Cache decorator. :param func: The function to cache. :return: Wrapped function with caching. """ def wrapper(*args: Any, **kwargs: Any) -> Generator[str, None, None]: # Exclude self instance from hashing. cache_key = md5(json.dumps((args[1:], kwargs)).encode("utf-8")).hexdigest() cache_file = self.cache_path / cache_key # TODO: Fix caching for chat, should hash last user message, (not entire history). if kwargs.pop("caching", True) and cache_file.exists(): yield cache_file.read_text() return result = "" for i in func(*args, **kwargs): result += i yield i cache_file.write_text(result) self._delete_oldest_files(self.length) # type: ignore return wrapper @no_type_check def _delete_oldest_files(self, max_files: int) -> None: """ Class method to delete the oldest cached files in the CACHE_DIR folder. :param max_files: Integer, the maximum number of files to keep in the CACHE_DIR folder. """ # Get all files in the folder. files = self.cache_path.glob("*") # Sort files by last modification time in ascending order. files = sorted(files, key=lambda f: f.stat().st_mtime) # Delete the oldest files if the number of files exceeds the limit. if len(files) > max_files: num_files_to_delete = len(files) - max_files for i in range(num_files_to_delete): files[i].unlink()