codekingpro/portable-devtools

1import threading
2import uuid
3from typing import Any, Callable
4from chromadb.types import Segment
5from overrides import override
6from typing import Dict, Optional
7from abc import ABC, abstractmethod
8
9
10class SegmentCache(ABC):
11    @abstractmethod
12    def get(self, key: uuid.UUID) -> Optional[Segment]:
13        pass
14
15    @abstractmethod
16    def pop(self, key: uuid.UUID) -> Optional[Segment]:
17        pass
18
19    @abstractmethod
20    def set(self, key: uuid.UUID, value: Segment) -> None:
21        pass
22
23    @abstractmethod
24    def reset(self) -> None:
25        pass
26
27
28class BasicCache(SegmentCache):
29    def __init__(self):
30        self.cache: Dict[uuid.UUID, Segment] = {}
31        self.lock = threading.RLock()
32
33    @override
34    def get(self, key: uuid.UUID) -> Optional[Segment]:
35        with self.lock:
36            return self.cache.get(key)
37
38    @override
39    def pop(self, key: uuid.UUID) -> Optional[Segment]:
40        with self.lock:
41            return self.cache.pop(key, None)
42
43    @override
44    def set(self, key: uuid.UUID, value: Segment) -> None:
45        with self.lock:
46            self.cache[key] = value
47
48    @override
49    def reset(self) -> None:
50        with self.lock:
51            self.cache = {}
52
53
54class SegmentLRUCache(BasicCache):
55    """A simple LRU cache implementation that handles objects with dynamic sizes.
56    The size of each object is determined by a user-provided size function."""
57
58    def __init__(
59        self,
60        capacity: int,
61        size_func: Callable[[uuid.UUID], int],
62        callback: Optional[Callable[[uuid.UUID, Segment], Any]] = None,
63    ):
64        self.capacity = capacity
65        self.size_func = size_func
66        self.cache: Dict[uuid.UUID, Segment] = {}
67        self.history = []
68        self.callback = callback
69        self.lock = threading.RLock()
70
71    def _upsert_key(self, key: uuid.UUID):
72        if key in self.history:
73            self.history.remove(key)
74            self.history.append(key)
75        else:
76            self.history.append(key)
77
78    @override
79    def get(self, key: uuid.UUID) -> Optional[Segment]:
80        with self.lock:
81            self._upsert_key(key)
82            if key in self.cache:
83                return self.cache[key]
84            else:
85                return None
86
87    @override
88    def pop(self, key: uuid.UUID) -> Optional[Segment]:
89        with self.lock:
90            if key in self.history:
91                self.history.remove(key)
92            return self.cache.pop(key, None)
93
94    @override
95    def set(self, key: uuid.UUID, value: Segment) -> None:
96        with self.lock:
97            if key in self.cache:
98                return
99            item_size = self.size_func(key)
100            key_sizes = {key: self.size_func(key) for key in self.cache}
101            total_size = sum(key_sizes.values())
102            index = 0
103            # Evict items if capacity is exceeded
104            while total_size + item_size > self.capacity and len(self.history) > index:
105                key_delete = self.history[index]
106                if key_delete in self.cache:
107                    self.callback(key_delete, self.cache[key_delete])
108                    del self.cache[key_delete]
109                    total_size -= key_sizes[key_delete]
110                index += 1
111
112            self.cache[key] = value
113            self._upsert_key(key)
114
115    @override
116    def reset(self):
117        with self.lock:
118            self.cache = {}
119            self.history = []
120