mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-08-28 05:25:04 +00:00
198 lines
7.3 KiB
Python
198 lines
7.3 KiB
Python
"""Local file-based cache utility for DataFrames, lists, dicts, and strings."""
|
|
|
|
import json
|
|
from datetime import datetime, timedelta
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pandas as pd
|
|
from loguru import logger
|
|
|
|
|
|
class CacheHandler:
|
|
"""Handles persistent data caching with expiration and type support."""
|
|
|
|
_EXTENSIONS = {
|
|
pd.DataFrame: ".csv",
|
|
dict: ".json",
|
|
list: ".jsonl",
|
|
str: ".txt",
|
|
}
|
|
|
|
_TYPE_NAMES = {
|
|
"DataFrame": pd.DataFrame,
|
|
"dict": dict,
|
|
"list": list,
|
|
"str": str,
|
|
}
|
|
|
|
def __init__(self, cache_dir: str | Path = "cache"):
|
|
"""Initialize cache directory and load existing metadata."""
|
|
self.cache_dir = Path(cache_dir)
|
|
self.cache_dir.mkdir(parents=True, exist_ok=True)
|
|
self.metadata_file = self.cache_dir / "metadata.json"
|
|
self.metadata: dict[str, Any] = self._load_metadata()
|
|
|
|
def set_cache_dir(self, cache_dir: str | Path) -> None:
|
|
"""Change the cache directory and reload metadata."""
|
|
self.cache_dir = Path(cache_dir)
|
|
self.cache_dir.mkdir(parents=True, exist_ok=True)
|
|
self.metadata_file = self.cache_dir / "metadata.json"
|
|
self.metadata = self._load_metadata()
|
|
logger.info(f"Cache directory moved to: {self.cache_dir}")
|
|
|
|
def _load_metadata(self) -> dict[str, Any]:
|
|
"""Load metadata from the JSON file."""
|
|
if self.metadata_file.exists():
|
|
try:
|
|
with open(self.metadata_file, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except (json.JSONDecodeError, OSError) as e:
|
|
logger.warning(f"Metadata load failed: {e}")
|
|
return {}
|
|
|
|
def _save_metadata(self) -> None:
|
|
"""Persist metadata to the disk."""
|
|
try:
|
|
with open(self.metadata_file, "w", encoding="utf-8") as f:
|
|
json.dump(self.metadata, f, ensure_ascii=False, indent=2)
|
|
except OSError as e:
|
|
logger.error(f"Metadata save failed: {e}")
|
|
|
|
def _get_path(self, key: str, data_type: type | None = None) -> Path:
|
|
"""Resolve the file path based on data type or metadata."""
|
|
ext = ".dat"
|
|
if data_type in self._EXTENSIONS:
|
|
ext = self._EXTENSIONS[data_type]
|
|
elif key in self.metadata:
|
|
stored_type = self.metadata[key].get("data_type")
|
|
ext = self._EXTENSIONS.get(self._TYPE_NAMES.get(stored_type, None), ".dat")
|
|
return self.cache_dir / f"{key}{ext}"
|
|
|
|
@staticmethod
|
|
def _execute_save(data: Any, path: Path, dtype: type, **kwargs) -> dict:
|
|
"""Execute type-specific save operations."""
|
|
if dtype is pd.DataFrame:
|
|
data.to_csv(path, index=kwargs.get("index", False), encoding="utf-8")
|
|
return {"row_count": len(data), "file_size": path.stat().st_size}
|
|
|
|
if dtype is dict:
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
json.dump(data, f, ensure_ascii=False, indent=2)
|
|
return {"item_count": len(data), "file_size": path.stat().st_size}
|
|
|
|
if dtype is list:
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
for item in data:
|
|
f.write(json.dumps(item, ensure_ascii=False) + "\n")
|
|
return {"item_count": len(data), "file_size": path.stat().st_size}
|
|
|
|
if dtype is str:
|
|
path.write_text(data, encoding=kwargs.get("encoding", "utf-8"))
|
|
return {"char_count": len(data), "file_size": path.stat().st_size}
|
|
|
|
raise ValueError(f"Unsupported type: {dtype}")
|
|
|
|
@staticmethod
|
|
def _execute_load(path: Path, type_name: str, **kwargs) -> Any:
|
|
"""Execute type-specific load operations."""
|
|
if type_name == "DataFrame":
|
|
return pd.read_csv(path, encoding=kwargs.get("encoding", "utf-8"))
|
|
if type_name == "dict":
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
if type_name == "list":
|
|
result = []
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if line:
|
|
result.append(json.loads(line))
|
|
return result
|
|
if type_name == "str":
|
|
return path.read_text(encoding=kwargs.get("encoding", "utf-8"))
|
|
raise ValueError(f"Unknown data type in metadata: {type_name}")
|
|
|
|
def save(self, key: str, data: Any, expire_hours: float | None = None, **kwargs) -> bool:
|
|
"""Save data to cache with optional expiration."""
|
|
try:
|
|
dtype = type(data)
|
|
path = self._get_path(key, dtype)
|
|
stats = self._execute_save(data, path, dtype, **kwargs)
|
|
|
|
now = datetime.now()
|
|
self.metadata[key] = {
|
|
"created_at": now.isoformat(),
|
|
"expire_at": (now + timedelta(hours=expire_hours)).isoformat() if expire_hours else None,
|
|
"data_type": dtype.__name__,
|
|
**stats,
|
|
}
|
|
self._save_metadata()
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Save failed for {key}: {e}")
|
|
return False
|
|
|
|
def load(self, key: str, auto_clean: bool = True, **kwargs) -> Any | None:
|
|
"""Load data from cache if not expired."""
|
|
if self._is_expired(key):
|
|
if auto_clean:
|
|
self.delete(key)
|
|
return None
|
|
|
|
path = self._get_path(key)
|
|
if not path.exists() or key not in self.metadata:
|
|
return None
|
|
|
|
try:
|
|
return self._execute_load(path, self.metadata[key]["data_type"], **kwargs)
|
|
except Exception as e:
|
|
logger.error(f"Load failed for {key}: {e}")
|
|
return None
|
|
|
|
def _is_expired(self, key: str) -> bool:
|
|
"""Check if the cached entry has expired."""
|
|
entry = self.metadata.get(key)
|
|
if not entry or not entry.get("expire_at"):
|
|
return False
|
|
return datetime.now() > datetime.fromisoformat(entry["expire_at"])
|
|
|
|
def delete(self, key: str) -> bool:
|
|
"""Remove a specific cache entry and its file."""
|
|
try:
|
|
path = self._get_path(key)
|
|
if path.exists():
|
|
path.unlink()
|
|
if key in self.metadata:
|
|
del self.metadata[key]
|
|
self._save_metadata()
|
|
return True
|
|
except OSError as e:
|
|
logger.error(f"Delete failed for {key}: {e}")
|
|
return False
|
|
|
|
def exists(self, key: str) -> bool:
|
|
"""Check if a valid cache entry exists."""
|
|
return key in self.metadata and not self._is_expired(key)
|
|
|
|
def clear_all(self) -> bool:
|
|
"""Purge all cache files and reset metadata."""
|
|
try:
|
|
for file in self.cache_dir.iterdir():
|
|
if file.is_file():
|
|
file.unlink()
|
|
self.metadata = {}
|
|
self._save_metadata()
|
|
return True
|
|
except OSError as e:
|
|
logger.error(f"Clear all failed: {e}")
|
|
return False
|
|
|
|
def get_stats(self) -> dict[str, Any]:
|
|
"""Return cache usage statistics."""
|
|
total_size = sum(f.stat().st_size for f in self.cache_dir.glob("*") if f.is_file())
|
|
return {
|
|
"count": len(self.metadata),
|
|
"size_mb": round(total_size / (1024 * 1024), 2),
|
|
"dir": str(self.cache_dir),
|
|
}
|