Blame view

utils/cache.py 2.48 KB
5bac9649   tangwang   文本 embedding 与图片 ...
1
  """Small file-backed cache helpers."""
be52af70   tangwang   first commit
2
3
  
  import json
be52af70   tangwang   first commit
4
5
  from pathlib import Path
  from typing import Any, Optional
be52af70   tangwang   first commit
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
  
  
  class DictCache:
      """
      Simple dictionary-based cache for query rewrite rules, translations, etc.
      """
  
      def __init__(self, cache_file: str = ".cache/dict_cache.json"):
          self.cache_file = Path(cache_file)
          self.cache_file.parent.mkdir(parents=True, exist_ok=True)
          self.cache = self._load()
  
      def _load(self) -> dict:
          """Load cache from file."""
          if self.cache_file.exists():
              try:
                  with open(self.cache_file, 'r', encoding='utf-8') as f:
                      return json.load(f)
              except Exception as e:
                  print(f"Failed to load cache: {e}")
                  return {}
          return {}
  
      def _save(self) -> bool:
          """Save cache to file."""
          try:
              with open(self.cache_file, 'w', encoding='utf-8') as f:
                  json.dump(self.cache, f, ensure_ascii=False, indent=2)
              return True
          except Exception as e:
              print(f"Failed to save cache: {e}")
              return False
  
      def get(self, key: str, category: str = "default") -> Optional[Any]:
          """
          Get cached value.
  
          Args:
              key: Cache key
              category: Cache category (for organizing different types of data)
  
          Returns:
              Cached value or None
          """
          return self.cache.get(category, {}).get(key)
  
      def set(self, key: str, value: Any, category: str = "default") -> bool:
          """
          Store value in cache.
  
          Args:
              key: Cache key
              value: Value to cache
              category: Cache category
  
          Returns:
              True if successful
          """
          if category not in self.cache:
              self.cache[category] = {}
          self.cache[category][key] = value
          return self._save()
  
      def exists(self, key: str, category: str = "default") -> bool:
          """Check if key exists in cache."""
          return category in self.cache and key in self.cache[category]
  
      def clear(self, category: Optional[str] = None) -> bool:
          """
          Clear cache.
  
          Args:
              category: If specified, clear only this category. Otherwise clear all.
  
          Returns:
              True if successful
          """
          if category:
              if category in self.cache:
                  del self.cache[category]
          else:
              self.cache = {}
          return self._save()