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

memos/memos_tools/notification_utils.py

@@ -2,6 +2,7 @@
 Notification utilities for MemOS product.
 """
 
+import asyncio
 import logging
 
 from collections.abc import Callable
@@ -51,6 +52,51 @@ def send_online_bot_notification(
         logger.warning(f"Failed to send online bot notification: {e}")
 
 
+async def send_online_bot_notification_async(
+    online_bot: Callable | None,
+    header_name: str,
+    sub_title_name: str,
+    title_color: str,
+    other_data1: dict[str, Any],
+    other_data2: dict[str, Any],
+    emoji: dict[str, str],
+) -> None:
+    """
+    Send notification via online_bot asynchronously if available.
+
+    Args:
+        online_bot: The online_bot function or None
+        header_name: Header name for the report
+        sub_title_name: Subtitle for the report
+        title_color: Title color
+        other_data1: First data dict
+        other_data2: Second data dict
+        emoji: Emoji configuration dict
+    """
+    if online_bot is None:
+        return
+
+    try:
+        # Run the potentially blocking notification in a thread pool
+        loop = asyncio.get_event_loop()
+        await loop.run_in_executor(
+            None,
+            lambda: online_bot(
+                header_name=header_name,
+                sub_title_name=sub_title_name,
+                title_color=title_color,
+                other_data1=other_data1,
+                other_data2=other_data2,
+                emoji=emoji,
+            ),
+        )
+
+        logger.info(f"Online bot notification sent successfully (async): {header_name}")
+
+    except Exception as e:
+        logger.warning(f"Failed to send online bot notification (async): {e}")
+
+
 def send_error_bot_notification(
     error_bot: Callable | None,
     err: str,

memos/memos_tools/thread_safe_dict.py

@@ -0,0 +1,288 @@
+"""
+Thread-safe dictionary wrapper for concurrent access with optimized read-write locks.
+"""
+
+import threading
+
+from collections.abc import ItemsView, Iterator, KeysView, ValuesView
+from typing import Generic, TypeVar
+
+
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+class ReadWriteLock:
+    """A simple read-write lock implementation. use for product-server scenario"""
+
+    def __init__(self):
+        self._read_ready = threading.Condition(threading.RLock())
+        self._readers = 0
+
+    def acquire_read(self):
+        """Acquire a read lock. Multiple readers can hold the lock simultaneously."""
+        self._read_ready.acquire()
+        try:
+            self._readers += 1
+        finally:
+            self._read_ready.release()
+
+    def release_read(self):
+        """Release a read lock."""
+        self._read_ready.acquire()
+        try:
+            self._readers -= 1
+            if self._readers == 0:
+                self._read_ready.notify_all()
+        finally:
+            self._read_ready.release()
+
+    def acquire_write(self):
+        """Acquire a write lock. Only one writer can hold the lock."""
+        self._read_ready.acquire()
+        while self._readers > 0:
+            self._read_ready.wait()
+
+    def release_write(self):
+        """Release a write lock."""
+        self._read_ready.release()
+
+
+class ThreadSafeDict(Generic[K, V]):
+    """
+    A thread-safe dictionary wrapper with optimized read-write locks.
+
+    This class allows multiple concurrent readers while ensuring exclusive access for writers.
+    Read operations (get, contains, iteration) can happen concurrently.
+    Write operations (set, delete, update) are exclusive.
+    """
+
+    def __init__(self, initial_dict: dict[K, V] | None = None):
+        """
+        Initialize the thread-safe dictionary.
+
+        Args:
+            initial_dict: Optional initial dictionary to copy from
+        """
+        self._dict: dict[K, V] = initial_dict.copy() if initial_dict else {}
+        self._lock = ReadWriteLock()
+
+    def __getitem__(self, key: K) -> V:
+        """Get item by key."""
+        self._lock.acquire_read()
+        try:
+            return self._dict[key]
+        finally:
+            self._lock.release_read()
+
+    def __setitem__(self, key: K, value: V) -> None:
+        """Set item by key."""
+        self._lock.acquire_write()
+        try:
+            self._dict[key] = value
+        finally:
+            self._lock.release_write()
+
+    def __delitem__(self, key: K) -> None:
+        """Delete item by key."""
+        self._lock.acquire_write()
+        try:
+            del self._dict[key]
+        finally:
+            self._lock.release_write()
+
+    def __contains__(self, key: K) -> bool:
+        """Check if key exists in dictionary."""
+        self._lock.acquire_read()
+        try:
+            return key in self._dict
+        finally:
+            self._lock.release_read()
+
+    def __len__(self) -> int:
+        """Get length of dictionary."""
+        self._lock.acquire_read()
+        try:
+            return len(self._dict)
+        finally:
+            self._lock.release_read()
+
+    def __bool__(self) -> bool:
+        """Check if dictionary is not empty."""
+        self._lock.acquire_read()
+        try:
+            return bool(self._dict)
+        finally:
+            self._lock.release_read()
+
+    def __iter__(self) -> Iterator[K]:
+        """Iterate over keys. Returns a snapshot to avoid iteration issues."""
+        self._lock.acquire_read()
+        try:
+            # Return a snapshot of keys to avoid iteration issues
+            return iter(list(self._dict.keys()))
+        finally:
+            self._lock.release_read()
+
+    def get(self, key: K, default: V | None = None) -> V:
+        """Get item by key with optional default."""
+        self._lock.acquire_read()
+        try:
+            return self._dict.get(key, default)
+        finally:
+            self._lock.release_read()
+
+    def pop(self, key: K, *args) -> V:
+        """Pop item by key."""
+        self._lock.acquire_write()
+        try:
+            return self._dict.pop(key, *args)
+        finally:
+            self._lock.release_write()
+
+    def update(self, *args, **kwargs) -> None:
+        """Update dictionary."""
+        self._lock.acquire_write()
+        try:
+            self._dict.update(*args, **kwargs)
+        finally:
+            self._lock.release_write()
+
+    def clear(self) -> None:
+        """Clear all items."""
+        self._lock.acquire_write()
+        try:
+            self._dict.clear()
+        finally:
+            self._lock.release_write()
+
+    def keys(self) -> KeysView[K]:
+        """Get dictionary keys view (snapshot)."""
+        self._lock.acquire_read()
+        try:
+            return list(self._dict.keys())
+        finally:
+            self._lock.release_read()
+
+    def values(self) -> ValuesView[V]:
+        """Get dictionary values view (snapshot)."""
+        self._lock.acquire_read()
+        try:
+            return list(self._dict.values())
+        finally:
+            self._lock.release_read()
+
+    def items(self) -> ItemsView[K, V]:
+        """Get dictionary items view (snapshot)."""
+        self._lock.acquire_read()
+        try:
+            return list(self._dict.items())
+        finally:
+            self._lock.release_read()
+
+    def copy(self) -> dict[K, V]:
+        """Create a copy of the dictionary."""
+        self._lock.acquire_read()
+        try:
+            return self._dict.copy()
+        finally:
+            self._lock.release_read()
+
+    def setdefault(self, key: K, default: V | None = None) -> V:
+        """Set default value for key if not exists."""
+        self._lock.acquire_write()
+        try:
+            return self._dict.setdefault(key, default)
+        finally:
+            self._lock.release_write()
+
+    def __repr__(self) -> str:
+        """String representation."""
+        self._lock.acquire_read()
+        try:
+            return f"ThreadSafeDict({self._dict})"
+        finally:
+            self._lock.release_read()
+
+    def __str__(self) -> str:
+        """String representation."""
+        self._lock.acquire_read()
+        try:
+            return str(self._dict)
+        finally:
+            self._lock.release_read()
+
+
+class SimpleThreadSafeDict(Generic[K, V]):
+    """
+    Simple thread-safe dictionary with exclusive locks for all operations.
+    Use this if you prefer simplicity over performance.
+    """
+
+    def __init__(self, initial_dict: dict[K, V] | None = None):
+        self._dict: dict[K, V] = initial_dict.copy() if initial_dict else {}
+        self._lock = threading.RLock()
+
+    def __getitem__(self, key: K) -> V:
+        with self._lock:
+            return self._dict[key]
+
+    def __setitem__(self, key: K, value: V) -> None:
+        with self._lock:
+            self._dict[key] = value
+
+    def __delitem__(self, key: K) -> None:
+        with self._lock:
+            del self._dict[key]
+
+    def __contains__(self, key: K) -> bool:
+        with self._lock:
+            return key in self._dict
+
+    def __len__(self) -> int:
+        with self._lock:
+            return len(self._dict)
+
+    def __bool__(self) -> bool:
+        with self._lock:
+            return bool(self._dict)
+
+    def __iter__(self) -> Iterator[K]:
+        with self._lock:
+            return iter(list(self._dict.keys()))
+
+    def get(self, key: K, default: V | None = None) -> V:
+        with self._lock:
+            return self._dict.get(key, default)
+
+    def pop(self, key: K, *args) -> V:
+        with self._lock:
+            return self._dict.pop(key, *args)
+
+    def update(self, *args, **kwargs) -> None:
+        with self._lock:
+            self._dict.update(*args, **kwargs)
+
+    def clear(self) -> None:
+        with self._lock:
+            self._dict.clear()
+
+    def keys(self):
+        with self._lock:
+            return list(self._dict.keys())
+
+    def values(self):
+        with self._lock:
+            return list(self._dict.values())
+
+    def items(self):
+        with self._lock:
+            return list(self._dict.items())
+
+    def copy(self) -> dict[K, V]:
+        with self._lock:
+            return self._dict.copy()
+
+    def setdefault(self, key: K, default: V | None = None) -> V:
+        with self._lock:
+            return self._dict.setdefault(key, default)

memos/reranker/__init__.py

@@ -0,0 +1,4 @@
+from .factory import RerankerFactory
+
+
+__all__ = ["RerankerFactory"]

memos/reranker/base.py

@@ -0,0 +1,24 @@
+# memos/reranker/base.py
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from memos.memories.textual.item import TextualMemoryItem
+
+
+class BaseReranker(ABC):
+    """Abstract interface for memory rerankers."""
+
+    @abstractmethod
+    def rerank(
+        self,
+        query: str,
+        graph_results: list,
+        top_k: int,
+        **kwargs,
+    ) -> list[tuple[TextualMemoryItem, float]]:
+        """Return top_k (item, score) sorted by score desc."""
+        raise NotImplementedError

memos/reranker/cosine_local.py

@@ -0,0 +1,95 @@
+# memos/reranker/cosine_local.py
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from .base import BaseReranker
+
+
+if TYPE_CHECKING:
+    from memos.memories.textual.item import TextualMemoryItem
+
+try:
+    import numpy as _np
+
+    _HAS_NUMPY = True
+except Exception:
+    _HAS_NUMPY = False
+
+
+def _cosine_one_to_many(q: list[float], m: list[list[float]]) -> list[float]:
+    """
+    Compute cosine similarities between a single vector q and a matrix m (rows are candidates).
+    """
+    if not _HAS_NUMPY:
+
+        def dot(a, b):  # lowercase per N806
+            return sum(x * y for x, y in zip(a, b, strict=False))
+
+        def norm(a):  # lowercase per N806
+            return sum(x * x for x in a) ** 0.5
+
+        qn = norm(q) or 1e-10
+        sims = []
+        for v in m:
+            vn = norm(v) or 1e-10
+            sims.append(dot(q, v) / (qn * vn))
+        return sims
+
+    qv = _np.asarray(q, dtype=float)  # lowercase
+    mv = _np.asarray(m, dtype=float)  # lowercase
+    qn = _np.linalg.norm(qv) or 1e-10
+    mn = _np.linalg.norm(mv, axis=1)  # lowercase
+    dots = mv @ qv
+    return (dots / (mn * qn + 1e-10)).tolist()
+
+
+class CosineLocalReranker(BaseReranker):
+    def __init__(
+        self,
+        level_weights: dict[str, float] | None = None,
+        level_field: str = "background",
+    ):
+        self.level_weights = level_weights or {"topic": 1.0, "concept": 1.0, "fact": 1.0}
+        self.level_field = level_field
+
+    def rerank(
+        self,
+        query: str,
+        graph_results: list,
+        top_k: int,
+        **kwargs,
+    ) -> list[tuple[TextualMemoryItem, float]]:
+        if not graph_results:
+            return []
+
+        query_embedding: list[float] | None = kwargs.get("query_embedding")
+        if not query_embedding:
+            return [(item, 0.0) for item in graph_results[:top_k]]
+
+        items_with_emb = [
+            it
+            for it in graph_results
+            if getattr(it, "metadata", None) and getattr(it.metadata, "embedding", None)
+        ]
+        if not items_with_emb:
+            return [(item, 0.5) for item in graph_results[:top_k]]
+
+        cand_vecs = [it.metadata.embedding for it in items_with_emb]
+        sims = _cosine_one_to_many(query_embedding, cand_vecs)
+
+        def get_weight(it: TextualMemoryItem) -> float:
+            level = getattr(it.metadata, self.level_field, None)
+            return self.level_weights.get(level, 1.0)
+
+        weighted = [sim * get_weight(it) for sim, it in zip(sims, items_with_emb, strict=False)]
+        scored_pairs = list(zip(items_with_emb, weighted, strict=False))
+        scored_pairs.sort(key=lambda x: x[1], reverse=True)
+
+        top_items = scored_pairs[:top_k]
+        if len(top_items) < top_k:
+            chosen = {it.id for it, _ in top_items}
+            remain = [(it, -1.0) for it in graph_results if it.id not in chosen]
+            top_items.extend(remain[: top_k - len(top_items)])
+
+        return top_items

memos/reranker/factory.py

@@ -0,0 +1,43 @@
+# memos/reranker/factory.py
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from .cosine_local import CosineLocalReranker
+from .http_bge import HTTPBGEReranker
+from .noop import NoopReranker
+
+
+if TYPE_CHECKING:
+    from memos.configs.reranker import RerankerConfigFactory
+
+    from .base import BaseReranker
+
+
+class RerankerFactory:
+    @staticmethod
+    def from_config(cfg: RerankerConfigFactory | None) -> BaseReranker | None:
+        if not cfg:
+            return None
+
+        backend = (cfg.backend or "").lower()
+        c: dict[str, Any] = cfg.config or {}
+
+        if backend in {"http_bge", "bge"}:
+            return HTTPBGEReranker(
+                reranker_url=c.get("url") or c.get("endpoint") or c.get("reranker_url"),
+                model=c.get("model", "bge-reranker-v2-m3"),
+                timeout=int(c.get("timeout", 10)),
+                headers_extra=c.get("headers_extra"),
+            )
+
+        if backend in {"cosine_local", "cosine"}:
+            return CosineLocalReranker(
+                level_weights=c.get("level_weights"),
+                level_field=c.get("level_field", "background"),
+            )
+
+        if backend in {"noop", "none", "disabled"}:
+            return NoopReranker()
+
+        raise ValueError(f"Unknown reranker backend: {cfg.backend}")

memos/reranker/http_bge.py

@@ -0,0 +1,99 @@
+# memos/reranker/http_bge.py
+from __future__ import annotations
+
+import re
+
+from typing import TYPE_CHECKING
+
+import requests
+
+from .base import BaseReranker
+
+
+if TYPE_CHECKING:
+    from memos.memories.textual.item import TextualMemoryItem
+
+_TAG1 = re.compile(r"^\s*\[[^\]]*\]\s*")
+
+
+class HTTPBGEReranker(BaseReranker):
+    """
+    HTTP-based BGE reranker. Mirrors your old MemoryReranker, but configurable.
+    """
+
+    def __init__(
+        self,
+        reranker_url: str,
+        token: str = "",
+        model: str = "bge-reranker-v2-m3",
+        timeout: int = 10,
+        headers_extra: dict | None = None,
+    ):
+        if not reranker_url:
+            raise ValueError("reranker_url must not be empty")
+        self.reranker_url = reranker_url
+        self.token = token or ""
+        self.model = model
+        self.timeout = timeout
+        self.headers_extra = headers_extra or {}
+
+    def rerank(
+        self,
+        query: str,
+        graph_results: list,
+        top_k: int,
+        **kwargs,
+    ) -> list[tuple[TextualMemoryItem, float]]:
+        if not graph_results:
+            return []
+
+        documents = [
+            (_TAG1.sub("", m) if isinstance((m := getattr(item, "memory", None)), str) else m)
+            for item in graph_results
+        ]
+        documents = [d for d in documents if isinstance(d, str) and d]
+        if not documents:
+            return []
+
+        headers = {"Content-Type": "application/json", **self.headers_extra}
+        payload = {"model": self.model, "query": query, "documents": documents}
+
+        try:
+            resp = requests.post(
+                self.reranker_url, headers=headers, json=payload, timeout=self.timeout
+            )
+            resp.raise_for_status()
+            data = resp.json()
+
+            scored_items: list[tuple[TextualMemoryItem, float]] = []
+
+            if "results" in data:
+                rows = data.get("results", [])
+                for r in rows:
+                    idx = r.get("index")
+                    if isinstance(idx, int) and 0 <= idx < len(graph_results):
+                        score = float(r.get("relevance_score", r.get("score", 0.0)))
+                        scored_items.append((graph_results[idx], score))
+
+                scored_items.sort(key=lambda x: x[1], reverse=True)
+                return scored_items[: min(top_k, len(scored_items))]
+
+            elif "data" in data:
+                rows = data.get("data", [])
+                score_list = [float(r.get("score", 0.0)) for r in rows]
+
+                if len(score_list) < len(graph_results):
+                    score_list += [0.0] * (len(graph_results) - len(score_list))
+                elif len(score_list) > len(graph_results):
+                    score_list = score_list[: len(graph_results)]
+
+                scored_items = list(zip(graph_results, score_list, strict=False))
+                scored_items.sort(key=lambda x: x[1], reverse=True)
+                return scored_items[: min(top_k, len(scored_items))]
+
+            else:
+                return [(item, 0.0) for item in graph_results[:top_k]]
+
+        except Exception as e:
+            print(f"[HTTPBGEReranker] request failed: {e}")
+            return [(item, 0.0) for item in graph_results[:top_k]]

memos/reranker/noop.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from .base import BaseReranker
+
+
+if TYPE_CHECKING:
+    from memos.memories.textual.item import TextualMemoryItem
+
+
+class NoopReranker(BaseReranker):
+    def rerank(
+        self, query: str, graph_results: list, top_k: int, **kwargs
+    ) -> list[tuple[TextualMemoryItem, float]]:
+        return [(item, 0.0) for item in graph_results[:top_k]]