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

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]]