MemoryOS 1.0.1__py3-none-any.whl → 1.1.1__py3-none-any.whl

memos/memos_tools/singleton.py

@@ -0,0 +1,174 @@
+"""
+Singleton decorator module for caching factory instances to avoid excessive memory usage
+from repeated initialization.
+"""
+
+import hashlib
+import json
+
+from collections.abc import Callable
+from functools import wraps
+from typing import Any, TypeVar
+from weakref import WeakValueDictionary
+
+
+T = TypeVar("T")
+
+
+class FactorySingleton:
+    """Factory singleton manager that caches instances based on configuration parameters"""
+
+    def __init__(self):
+        # Use weak reference dictionary for automatic cleanup when instances are no longer referenced
+        self._instances: dict[str, WeakValueDictionary] = {}
+
+    def _generate_cache_key(self, config: Any, *args, **kwargs) -> str:
+        """Generate cache key based on configuration only (ignoring other parameters)"""
+
+        # Handle configuration objects - only use the config parameter
+        if hasattr(config, "model_dump"):  # Pydantic model
+            config_data = config.model_dump()
+        elif hasattr(config, "dict"):  # Legacy Pydantic model
+            config_data = config.dict()
+        elif isinstance(config, dict):
+            config_data = config
+        else:
+            # For other types, try to convert to string
+            config_data = str(config)
+
+        # Filter out time-related fields that shouldn't affect caching
+        filtered_config = self._filter_temporal_fields(config_data)
+
+        # Generate hash key based only on config
+        try:
+            cache_str = json.dumps(filtered_config, sort_keys=True, ensure_ascii=False, default=str)
+        except (TypeError, ValueError):
+            # If JSON serialization fails, convert the entire config to string
+            cache_str = str(filtered_config)
+
+        return hashlib.md5(cache_str.encode("utf-8")).hexdigest()
+
+    def _filter_temporal_fields(self, config_data: Any) -> Any:
+        """Filter out temporal fields that shouldn't affect instance caching"""
+        if isinstance(config_data, dict):
+            filtered = {}
+            for key, value in config_data.items():
+                # Skip common temporal field names
+                if key.lower() in {
+                    "created_at",
+                    "updated_at",
+                    "timestamp",
+                    "time",
+                    "date",
+                    "created_time",
+                    "updated_time",
+                    "last_modified",
+                    "modified_at",
+                    "start_time",
+                    "end_time",
+                    "execution_time",
+                    "run_time",
+                }:
+                    continue
+                # Recursively filter nested dictionaries
+                filtered[key] = self._filter_temporal_fields(value)
+            return filtered
+        elif isinstance(config_data, list):
+            # Recursively filter lists
+            return [self._filter_temporal_fields(item) for item in config_data]
+        else:
+            # For primitive types, return as-is
+            return config_data
+
+    def get_or_create(self, factory_class: type, cache_key: str, creator_func: Callable) -> Any:
+        """Get or create instance"""
+        class_name = factory_class.__name__
+
+        if class_name not in self._instances:
+            self._instances[class_name] = WeakValueDictionary()
+
+        class_cache = self._instances[class_name]
+
+        if cache_key in class_cache:
+            return class_cache[cache_key]
+
+        # Create new instance
+        instance = creator_func()
+        class_cache[cache_key] = instance
+        return instance
+
+    def clear_cache(self, factory_class: type | None = None):
+        """Clear cache"""
+        if factory_class:
+            class_name = factory_class.__name__
+            if class_name in self._instances:
+                self._instances[class_name].clear()
+        else:
+            for cache in self._instances.values():
+                cache.clear()
+
+
+# Global singleton manager
+_factory_singleton = FactorySingleton()
+
+
+def singleton_factory(factory_class: type | str | None = None):
+    """
+    Factory singleton decorator
+
+    Usage:
+    @singleton_factory()
+    def from_config(cls, config):
+        return SomeClass(config)
+
+    Or specify factory class:
+    @singleton_factory(EmbedderFactory)
+    def from_config(cls, config):
+        return SomeClass(config)
+    """
+
+    def decorator(func: Callable[..., T]) -> Callable[..., T]:
+        @wraps(func)
+        def wrapper(*args, **kwargs) -> T:
+            # Determine factory class and config parameter
+            target_factory_class = factory_class
+            config = None
+
+            # Simple logic: check if first parameter is a class or config
+            if args:
+                if hasattr(args[0], "__name__") and hasattr(args[0], "__module__"):
+                    # First parameter is a class (cls), so this is a @classmethod
+                    if target_factory_class is None:
+                        target_factory_class = args[0]
+                    config = args[1] if len(args) > 1 else None
+                else:
+                    # First parameter is config, so this is a @staticmethod
+                    if target_factory_class is None:
+                        raise ValueError(
+                            "Factory class must be explicitly specified for static methods"
+                        )
+                    if isinstance(target_factory_class, str):
+                        # Convert string to a mock class for caching purposes
+                        class MockFactoryClass:
+                            __name__ = target_factory_class
+
+                        target_factory_class = MockFactoryClass
+                    config = args[0]
+
+            if config is None:
+                # If no configuration parameter, call original function directly
+                return func(*args, **kwargs)
+
+            # Generate cache key based only on config
+            cache_key = _factory_singleton._generate_cache_key(config)
+
+            # Function to create instance
+            def creator():
+                return func(*args, **kwargs)
+
+            # Get or create instance
+            return _factory_singleton.get_or_create(target_factory_class, cache_key, creator)
+
+        return wrapper
+
+    return decorator

memos/memos_tools/thread_safe_dict.py

@@ -7,10 +7,15 @@ import threading
 from collections.abc import ItemsView, Iterator, KeysView, ValuesView
 from typing import Generic, TypeVar
 
+from memos.log import get_logger
+from memos.utils import timed
+
 
 K = TypeVar("K")
 V = TypeVar("V")
 
+logger = get_logger(__name__)
+
 
 class ReadWriteLock:
     """A simple read-write lock implementation. use for product-server scenario"""
@@ -19,6 +24,7 @@ class ReadWriteLock:
         self._read_ready = threading.Condition(threading.RLock())
         self._readers = 0
 
+    @timed
     def acquire_read(self):
         """Acquire a read lock. Multiple readers can hold the lock simultaneously."""
         self._read_ready.acquire()
@@ -37,6 +43,7 @@ class ReadWriteLock:
         finally:
             self._read_ready.release()
 
+    @timed
     def acquire_write(self):
         """Acquire a write lock. Only one writer can hold the lock."""
         self._read_ready.acquire()
@@ -67,6 +74,7 @@ class ThreadSafeDict(Generic[K, V]):
         self._dict: dict[K, V] = initial_dict.copy() if initial_dict else {}
         self._lock = ReadWriteLock()
 
+    @timed
     def __getitem__(self, key: K) -> V:
         """Get item by key."""
         self._lock.acquire_read()
@@ -75,6 +83,7 @@ class ThreadSafeDict(Generic[K, V]):
         finally:
             self._lock.release_read()
 
+    @timed
     def __setitem__(self, key: K, value: V) -> None:
         """Set item by key."""
         self._lock.acquire_write()
@@ -83,6 +92,7 @@ class ThreadSafeDict(Generic[K, V]):
         finally:
             self._lock.release_write()
 
+    @timed
     def __delitem__(self, key: K) -> None:
         """Delete item by key."""
         self._lock.acquire_write()
@@ -91,6 +101,7 @@ class ThreadSafeDict(Generic[K, V]):
         finally:
             self._lock.release_write()
 
+    @timed
     def __contains__(self, key: K) -> bool:
         """Check if key exists in dictionary."""
         self._lock.acquire_read()
@@ -99,6 +110,7 @@ class ThreadSafeDict(Generic[K, V]):
         finally:
             self._lock.release_read()
 
+    @timed
     def __len__(self) -> int:
         """Get length of dictionary."""
         self._lock.acquire_read()
@@ -115,6 +127,7 @@ class ThreadSafeDict(Generic[K, V]):
         finally:
             self._lock.release_read()
 
+    @timed
     def __iter__(self) -> Iterator[K]:
         """Iterate over keys. Returns a snapshot to avoid iteration issues."""
         self._lock.acquire_read()
@@ -124,6 +137,7 @@ class ThreadSafeDict(Generic[K, V]):
         finally:
             self._lock.release_read()
 
+    @timed
     def get(self, key: K, default: V | None = None) -> V:
         """Get item by key with optional default."""
         self._lock.acquire_read()
@@ -132,6 +146,7 @@ class ThreadSafeDict(Generic[K, V]):
         finally:
             self._lock.release_read()
 
+    @timed
     def pop(self, key: K, *args) -> V:
         """Pop item by key."""
         self._lock.acquire_write()
@@ -140,6 +155,7 @@ class ThreadSafeDict(Generic[K, V]):
         finally:
             self._lock.release_write()
 
+    @timed
     def update(self, *args, **kwargs) -> None:
         """Update dictionary."""
         self._lock.acquire_write()
@@ -148,6 +164,7 @@ class ThreadSafeDict(Generic[K, V]):
         finally:
             self._lock.release_write()
 
+    @timed
     def clear(self) -> None:
         """Clear all items."""
         self._lock.acquire_write()
@@ -156,6 +173,7 @@ class ThreadSafeDict(Generic[K, V]):
         finally:
             self._lock.release_write()
 
+    @timed
     def keys(self) -> KeysView[K]:
         """Get dictionary keys view (snapshot)."""
         self._lock.acquire_read()
@@ -164,6 +182,7 @@ class ThreadSafeDict(Generic[K, V]):
         finally:
             self._lock.release_read()
 
+    @timed
     def values(self) -> ValuesView[V]:
         """Get dictionary values view (snapshot)."""
         self._lock.acquire_read()
@@ -172,6 +191,7 @@ class ThreadSafeDict(Generic[K, V]):
         finally:
             self._lock.release_read()
 
+    @timed
     def items(self) -> ItemsView[K, V]:
         """Get dictionary items view (snapshot)."""
         self._lock.acquire_read()
@@ -180,6 +200,7 @@ class ThreadSafeDict(Generic[K, V]):
         finally:
             self._lock.release_read()
 
+    @timed
     def copy(self) -> dict[K, V]:
         """Create a copy of the dictionary."""
         self._lock.acquire_read()
@@ -188,6 +209,7 @@ class ThreadSafeDict(Generic[K, V]):
         finally:
             self._lock.release_read()
 
+    @timed
     def setdefault(self, key: K, default: V | None = None) -> V:
         """Set default value for key if not exists."""
         self._lock.acquire_write()

memos/memos_tools/thread_safe_dict_segment.py

@@ -0,0 +1,382 @@
+import threading
+import time
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import Any, Generic, TypeVar
+
+
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+class FastReadWriteLock:
+    """Read-write lock optimized for FastAPI scenarios:
+    reader priority with writer starvation prevention"""
+
+    def __init__(self):
+        self._readers = 0
+        self._writers = 0
+        self._waiting_writers = 0
+        self._lock = threading.RLock()
+        self._read_ready = threading.Condition(self._lock)
+        self._write_ready = threading.Condition(self._lock)
+        # Writer starvation detection
+        self._last_write_time = 0
+        self._write_starvation_threshold = 0.1  # 100ms
+
+    def acquire_read(self) -> bool:
+        """Fast read lock acquisition"""
+        with self._lock:
+            # Check if writers are starving
+            current_time = time.time()
+            write_starving = (
+                self._waiting_writers > 0
+                and current_time - self._last_write_time > self._write_starvation_threshold
+            )
+
+            # If no writers are active and no starvation, allow readers to continue
+            if self._writers == 0 and not write_starving:
+                self._readers += 1
+                return True
+
+            # Otherwise wait
+            while self._writers > 0 or write_starving:
+                self._read_ready.wait()
+                current_time = time.time()
+                write_starving = (
+                    self._waiting_writers > 0
+                    and current_time - self._last_write_time > self._write_starvation_threshold
+                )
+
+            self._readers += 1
+            return True
+
+    def release_read(self):
+        """Release read lock"""
+        with self._lock:
+            self._readers -= 1
+            if self._readers == 0:
+                self._write_ready.notify()
+
+    def acquire_write(self) -> bool:
+        """Write lock acquisition"""
+        with self._lock:
+            self._waiting_writers += 1
+            try:
+                while self._readers > 0 or self._writers > 0:
+                    self._write_ready.wait()
+
+                self._writers = 1
+                self._waiting_writers -= 1
+                self._last_write_time = time.time()
+                return True
+            except:
+                self._waiting_writers -= 1
+                raise
+
+    def release_write(self):
+        """Release write lock"""
+        with self._lock:
+            self._writers = 0
+            # Prioritize notifying readers (reader priority strategy)
+            self._read_ready.notify_all()
+            self._write_ready.notify()
+
+
+class SegmentedLock:
+    """Segmented lock, segments based on key hash"""
+
+    def __init__(self, segment_count: int = 64):
+        self.segment_count = segment_count
+        self.locks = [FastReadWriteLock() for _ in range(segment_count)]
+
+    def get_lock(self, key: K) -> FastReadWriteLock:
+        """Get the corresponding lock based on key"""
+        segment = hash(key) % self.segment_count
+        return self.locks[segment]
+
+    @contextmanager
+    def read_lock(self, key: K):
+        """Read lock context manager"""
+        lock = self.get_lock(key)
+        lock.acquire_read()
+        try:
+            yield
+        finally:
+            lock.release_read()
+
+    @contextmanager
+    def write_lock(self, key: K):
+        """Write lock context manager"""
+        lock = self.get_lock(key)
+        lock.acquire_write()
+        try:
+            yield
+        finally:
+            lock.release_write()
+
+
+class OptimizedThreadSafeDict(Generic[K, V]):
+    """
+    Thread-safe dictionary optimized for FastAPI scenarios:
+    - Segmented locks to reduce contention
+    - Reader priority with writer starvation prevention
+    - Support for large object storage
+    - Strong consistency guarantee
+    """
+
+    def __init__(
+        self, initial_dict: dict[K, V] | None = None, segment_count: int = 128
+    ):  # More segments for high concurrency
+        self._segments: list[dict[K, V]] = [{} for _ in range(segment_count)]
+        self._segment_count = segment_count
+        self._segmented_lock = SegmentedLock(segment_count)
+
+        # Initialize data
+        if initial_dict:
+            for k, v in initial_dict.items():
+                segment_idx = self._get_segment(k)
+                self._segments[segment_idx][k] = v
+
+    def _get_segment(self, key: K) -> int:
+        """Calculate the segment corresponding to the key"""
+        return hash(key) % self._segment_count
+
+    def __getitem__(self, key: K) -> V:
+        """Get element"""
+        segment_idx = self._get_segment(key)
+        with self._segmented_lock.read_lock(key):
+            return self._segments[segment_idx][key]
+
+    def __setitem__(self, key: K, value: V) -> None:
+        """Set element - key optimization point"""
+        segment_idx = self._get_segment(key)
+        with self._segmented_lock.write_lock(key):
+            self._segments[segment_idx][key] = value
+
+    def __delitem__(self, key: K) -> None:
+        """Delete element"""
+        segment_idx = self._get_segment(key)
+        with self._segmented_lock.write_lock(key):
+            del self._segments[segment_idx][key]
+
+    def __contains__(self, key: K) -> bool:
+        """Check if key is contained"""
+        segment_idx = self._get_segment(key)
+        with self._segmented_lock.read_lock(key):
+            return key in self._segments[segment_idx]
+
+    def get(self, key: K, default: V | None = None) -> V | None:
+        """Safely get element"""
+        segment_idx = self._get_segment(key)
+        with self._segmented_lock.read_lock(key):
+            return self._segments[segment_idx].get(key, default)
+
+    def pop(self, key: K, *args) -> V:
+        """Pop element"""
+        segment_idx = self._get_segment(key)
+        with self._segmented_lock.write_lock(key):
+            return self._segments[segment_idx].pop(key, *args)
+
+    def setdefault(self, key: K, default: V | None = None) -> V:
+        """Set default value"""
+        segment_idx = self._get_segment(key)
+        with self._segmented_lock.write_lock(key):
+            return self._segments[segment_idx].setdefault(key, default)
+
+    def update(self, other=None, **kwargs) -> None:
+        """Batch update - optimized batch operation"""
+        items = (other.items() if hasattr(other, "items") else other) if other is not None else []
+
+        # Group update items by segment
+        segment_updates: dict[int, list[tuple[K, V]]] = {}
+
+        for k, v in items:
+            segment_idx = self._get_segment(k)
+            if segment_idx not in segment_updates:
+                segment_updates[segment_idx] = []
+            segment_updates[segment_idx].append((k, v))
+
+        for k, v in kwargs.items():
+            segment_idx = self._get_segment(k)
+            if segment_idx not in segment_updates:
+                segment_updates[segment_idx] = []
+            segment_updates[segment_idx].append((k, v))
+
+        # Update segment by segment to reduce lock holding time
+        for segment_idx, updates in segment_updates.items():
+            # Use the first key to get the lock (all keys in the same segment map to the same lock)
+            first_key = updates[0][0]
+            with self._segmented_lock.write_lock(first_key):
+                for k, v in updates:
+                    self._segments[segment_idx][k] = v
+
+    def clear(self) -> None:
+        """Clear all elements - need to acquire all locks"""
+        # Acquire all locks in order to avoid deadlock
+        acquired_locks = []
+        try:
+            for i in range(self._segment_count):
+                lock = self._segmented_lock.locks[i]
+                lock.acquire_write()
+                acquired_locks.append(lock)
+
+            # Clear all segments
+            for segment in self._segments:
+                segment.clear()
+
+        finally:
+            # Release locks in reverse order
+            for lock in reversed(acquired_locks):
+                lock.release_write()
+
+    def __len__(self) -> int:
+        """Get total length - snapshot read"""
+        total = 0
+        acquired_locks = []
+        try:
+            # Acquire all read locks
+            for i in range(self._segment_count):
+                lock = self._segmented_lock.locks[i]
+                lock.acquire_read()
+                acquired_locks.append(lock)
+
+            # Calculate total length
+            for segment in self._segments:
+                total += len(segment)
+
+            return total
+
+        finally:
+            # Release all read locks
+            for lock in reversed(acquired_locks):
+                lock.release_read()
+
+    def __bool__(self) -> bool:
+        """Check if empty"""
+        return len(self) > 0
+
+    def keys(self) -> list[K]:
+        """Get snapshot of all keys"""
+        all_keys = []
+        acquired_locks = []
+
+        try:
+            # Acquire all read locks
+            for i in range(self._segment_count):
+                lock = self._segmented_lock.locks[i]
+                lock.acquire_read()
+                acquired_locks.append(lock)
+
+            # Collect all keys
+            for segment in self._segments:
+                all_keys.extend(segment.keys())
+
+            return all_keys
+
+        finally:
+            for lock in reversed(acquired_locks):
+                lock.release_read()
+
+    def values(self) -> list[V]:
+        """Get snapshot of all values"""
+        all_values = []
+        acquired_locks = []
+
+        try:
+            for i in range(self._segment_count):
+                lock = self._segmented_lock.locks[i]
+                lock.acquire_read()
+                acquired_locks.append(lock)
+
+            for segment in self._segments:
+                all_values.extend(segment.values())
+
+            return all_values
+
+        finally:
+            for lock in reversed(acquired_locks):
+                lock.release_read()
+
+    def items(self) -> list[tuple[K, V]]:
+        """Get snapshot of all items"""
+        all_items = []
+        acquired_locks = []
+
+        try:
+            for i in range(self._segment_count):
+                lock = self._segmented_lock.locks[i]
+                lock.acquire_read()
+                acquired_locks.append(lock)
+
+            for segment in self._segments:
+                all_items.extend(segment.items())
+
+            return all_items
+
+        finally:
+            for lock in reversed(acquired_locks):
+                lock.release_read()
+
+    def copy(self) -> dict[K, V]:
+        """Create dictionary copy"""
+        result = {}
+        acquired_locks = []
+
+        try:
+            for i in range(self._segment_count):
+                lock = self._segmented_lock.locks[i]
+                lock.acquire_read()
+                acquired_locks.append(lock)
+
+            for segment in self._segments:
+                result.update(segment)
+
+            return result
+
+        finally:
+            for lock in reversed(acquired_locks):
+                lock.release_read()
+
+    def __iter__(self) -> Iterator[K]:
+        """Iterator - returns snapshot"""
+        return iter(self.keys())
+
+    def __repr__(self) -> str:
+        """String representation"""
+        return f"OptimizedThreadSafeDict({dict(self.items())})"
+
+    def stats(self) -> dict[str, Any]:
+        """Get statistics"""
+        segment_sizes = []
+        total_items = 0
+
+        acquired_locks = []
+        try:
+            for i in range(self._segment_count):
+                lock = self._segmented_lock.locks[i]
+                lock.acquire_read()
+                acquired_locks.append(lock)
+
+            for segment in self._segments:
+                size = len(segment)
+                segment_sizes.append(size)
+                total_items += size
+
+            avg_size = total_items / self._segment_count if self._segment_count > 0 else 0
+            max_size = max(segment_sizes) if segment_sizes else 0
+            min_size = min(segment_sizes) if segment_sizes else 0
+
+            return {
+                "total_items": total_items,
+                "segment_count": self._segment_count,
+                "avg_segment_size": avg_size,
+                "max_segment_size": max_size,
+                "min_segment_size": min_size,
+                "load_balance_ratio": min_size / max_size if max_size > 0 else 1.0,
+            }
+
+        finally:
+            for lock in reversed(acquired_locks):
+                lock.release_read()