feat: add reranker Facktory

CaralHsi · CaralHsi · commit f7992376e33b · 2025-09-08T20:42:00.000+08:00
diff --git a/examples/basic_modules/reranker.py b/examples/basic_modules/reranker.py
@@ -0,0 +1,144 @@
+import os
+import uuid
+
+from dotenv import load_dotenv
+
+from memos import log
+from memos.configs.embedder import EmbedderConfigFactory
+from memos.configs.reranker import RerankerConfigFactory
+from memos.embedders.factory import EmbedderFactory
+from memos.memories.textual.item import TextualMemoryItem, TreeNodeTextualMemoryMetadata
+from memos.reranker.factory import RerankerFactory
+
+
+load_dotenv()
+logger = log.get_logger(__name__)
+
+
+def make_item(text: str) -> TextualMemoryItem:
+    """Build a minimal TextualMemoryItem; embedding will be populated later."""
+    return TextualMemoryItem(
+        id=str(uuid.uuid4()),
+        memory=text,
+        metadata=TreeNodeTextualMemoryMetadata(
+            user_id=None,
+            session_id=None,
+            status="activated",
+            type="fact",
+            memory_time="2024-01-01",
+            source="conversation",
+            confidence=100.0,
+            tags=[],
+            visibility="public",
+            updated_at="2025-01-01T00:00:00",
+            memory_type="LongTermMemory",
+            key="demo_key",
+            sources=["demo://example"],
+            embedding=[],
+            background="demo background...",
+        ),
+    )
+
+
+def show_ranked(title: str, ranked: list[tuple[TextualMemoryItem, float]], top_n: int = 5) -> None:
+    print(f"\n=== {title} ===")
+    for i, (item, score) in enumerate(ranked[:top_n], start=1):
+        preview = (item.memory[:80] + "...") if len(item.memory) > 80 else item.memory
+        print(f"[#{i}] score={score:.6f} | {preview}")
+
+
+def main():
+    # -------------------------------
+    # 1) Build the embedder (real vectors)
+    # -------------------------------
+    embedder_cfg = EmbedderConfigFactory.model_validate(
+        {
+            "backend": "universal_api",
+            "config": {
+                "provider": "openai",  # or "azure"
+                "api_key": os.getenv("OPENAI_API_KEY"),
+                "model_name_or_path": "text-embedding-3-large",
+                "base_url": os.getenv("OPENAI_API_BASE"),  # optional
+            },
+        }
+    )
+    embedder = EmbedderFactory.from_config(embedder_cfg)
+
+    # -------------------------------
+    # 2) Prepare query + documents
+    # -------------------------------
+    query = "What is the capital of France?"
+    items = [
+        make_item("Paris is the capital of France."),
+        make_item("Berlin is the capital of Germany."),
+        make_item("The capital of Brazil is Brasilia."),
+        make_item("Apples and bananas are common fruits."),
+        make_item("The Eiffel Tower is a famous landmark in Paris."),
+    ]
+
+    # -------------------------------
+    # 3) Embed query + docs with real embeddings
+    # -------------------------------
+    texts_to_embed = [query] + [it.memory for it in items]
+    vectors = embedder.embed(texts_to_embed)  # real vectors from your provider/model
+    query_embedding = vectors[0]
+    doc_embeddings = vectors[1:]
+
+    # attach real embeddings back to items
+    for it, emb in zip(items, doc_embeddings, strict=False):
+        it.metadata.embedding = emb
+
+    # -------------------------------
+    # 4) Rerank with cosine_local (uses your real embeddings)
+    # -------------------------------
+    cosine_cfg = RerankerConfigFactory.model_validate(
+        {
+            "backend": "cosine_local",
+            "config": {
+                # structural boosts (optional): uses metadata.background
+                "level_weights": {"topic": 1.0, "concept": 1.0, "fact": 1.0},
+                "level_field": "background",
+            },
+        }
+    )
+    cosine_reranker = RerankerFactory.from_config(cosine_cfg)
+
+    ranked_cosine = cosine_reranker.rerank(
+        query=query,
+        graph_results=items,
+        top_k=10,
+        query_embedding=query_embedding,  # required by cosine_local
+    )
+    show_ranked("CosineLocal Reranker (with real embeddings)", ranked_cosine, top_n=5)
+
+    # -------------------------------
+    # 5) (Optional) Rerank with HTTP BGE (OpenAI-style /query+documents)
+    #     Requires the service URL; no need for embeddings here
+    # -------------------------------
+    bge_url = os.getenv("BGE_RERANKER_URL")  # e.g., "http://xxx.x.xxxxx.xxx:xxxx/v1/rerank"
+    if bge_url:
+        http_cfg = RerankerConfigFactory.model_validate(
+            {
+                "backend": "http_bge",
+                "config": {
+                    "url": bge_url,
+                    "model": os.getenv("BGE_RERANKER_MODEL", "bge-reranker-v2-m3"),
+                    "timeout": int(os.getenv("BGE_RERANKER_TIMEOUT", "10")),
+                    # "headers_extra": {"Authorization": f"Bearer {os.getenv('BGE_RERANKER_TOKEN')}"}
+                },
+            }
+        )
+        http_reranker = RerankerFactory.from_config(http_cfg)
+
+        ranked_http = http_reranker.rerank(
+            query=query,
+            graph_results=items,  # uses item.memory internally as documents
+            top_k=10,
+        )
+        show_ranked("HTTP BGE Reranker (OpenAI-style API)", ranked_http, top_n=5)
+    else:
+        print("\n[Info] Skipped HTTP BGE scenario because BGE_RERANKER_URL is not set.")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/memos/configs/reranker.py b/src/memos/configs/reranker.py
@@ -0,0 +1,18 @@
+# memos/configs/reranker.py
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class RerankerConfigFactory(BaseModel):
+    """
+    {
+      "backend": "http_bge" | "cosine_local" | "noop",
+      "config": { ... backend-specific ... }
+    }
+    """
+
+    backend: str = Field(..., description="Reranker backend id")
+    config: dict[str, Any] = Field(default_factory=dict, description="Backend-specific options")
diff --git a/src/memos/reranker/__init__.py b/src/memos/reranker/__init__.py
@@ -0,0 +1,4 @@
+from .factory import RerankerFactory
+
+
+__all__ = ["RerankerFactory"]
diff --git a/src/memos/reranker/base.py b/src/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
diff --git a/src/memos/reranker/cosine_local.py b/src/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
diff --git a/src/memos/reranker/factory.py b/src/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}")
diff --git a/src/memos/reranker/http_bge.py b/src/memos/reranker/http_bge.py
diff --git a/src/memos/reranker/noop.py b/src/memos/reranker/noop.py

-Original file line number
+Diff line change
@@ @@ -0,0 +1,4 @@ @@
 +from .factory import RerankerFactory
++
++
 +__all__ = ["RerankerFactory"]