From 3d786e033e6eeb66c9188d82371e7751c9a9fa90 Mon Sep 17 00:00:00 2001
From: RerankerGuo <121015044+RerankerGuo@users.noreply.github.com>
Date: Wed, 29 Jul 2026 10:18:27 +0800
Subject: [PATCH 1/2] fix: add numerical stability safeguards to
 CosineLocalReranker (#1365)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Improves _cosine_one_to_many robustness for the reranker used
during contextual search — a path contributing to memory drift
when vague follow-up queries produce all-zero or oddly-shaped
embeddings.

Changes:
- Zero query / candidate vectors now return similarity 0.0
  instead of NaN or divide-by-zero anomalies; this prevents
  undefined similarities from leaking into re-ordered results
- Guard against 2D (batch-shaped) query vectors by flattening
  before the matmul, avoiding silent broadcasting/shape errors
- Return early for empty or malformed candidate matrices
- Wrap final scores with np.nan_to_num to clamp any remaining
  NaN / +Inf / -Inf edge cases to 0.0
- Apply the same NaN/Inf guarding in the pure-Python fallback
  (the _HAS_NUMPY=False branch) for consistency
- Added comprehensive tests in tests/reranker/test_cosine_local.py
  covering: identical/orthogonal vectors, zero vectors, 2D
  query flattening, empty matrix, opposite direction, weights,
  and missing-embedding fallback.

Closes #1365 (memory drift: mismatched reranking scores can
cause old memories to surface when a follow-up search query
produces an edge-case embedding shape.)

Test: python3 -m py_compile src/memos/reranker/cosine_local.py
Test: python3 -m py_compile tests/reranker/test_cosine_local.py
---
 src/memos/reranker/cosine_local.py  |  21 ++++-
 tests/reranker/test_cosine_local.py | 129 ++++++++++++++++++++++++++++
 2 files changed, 147 insertions(+), 3 deletions(-)
 create mode 100644 tests/reranker/test_cosine_local.py

diff --git a/src/memos/reranker/cosine_local.py b/src/memos/reranker/cosine_local.py
index 140074b1d..0b09f2f7a 100644
--- a/src/memos/reranker/cosine_local.py
+++ b/src/memos/reranker/cosine_local.py
@@ -25,6 +25,9 @@
 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).
+    Returns a list of floats in [-1, 1], with numerical-stability safeguards:
+    zero vectors produce similarity 0 instead of NaN/Inf, and any remaining edge cases
+    (NaN, Inf) are clamped to 0.
     """
     if not _HAS_NUMPY:
 
@@ -38,15 +41,27 @@ def norm(a):  # lowercase per N806
         sims = []
         for v in m:
             vn = norm(v) or 1e-10
-            sims.append(dot(q, v) / (qn * vn))
+            score = dot(q, v) / (qn * vn)
+            if score != score or score in (float("inf"), float("-inf")):
+                score = 0.0
+            sims.append(score)
         return sims
 
     qv = _np.asarray(q, dtype=float)  # lowercase
+    if qv.ndim > 1:
+        qv = qv.reshape(-1)
     mv = _np.asarray(m, dtype=float)  # lowercase
-    qn = _np.linalg.norm(qv) or 1e-10
+    if mv.ndim != 2 or mv.shape[1] == 0:
+        return [0.0] * len(m)
+    qn = _np.linalg.norm(qv)
     mn = _np.linalg.norm(mv, axis=1)  # lowercase
+    if qn < 1e-10:
+        return [0.0] * len(m)
+    denom = mn * qn + 1e-10
     dots = mv @ qv
-    return (dots / (mn * qn + 1e-10)).tolist()
+    scores = dots / denom
+    scores = _np.nan_to_num(scores, nan=0.0, posinf=0.0, neginf=0.0)
+    return scores.tolist()
 
 
 class CosineLocalReranker(BaseReranker):
diff --git a/tests/reranker/test_cosine_local.py b/tests/reranker/test_cosine_local.py
new file mode 100644
index 000000000..090eb4e3e
--- /dev/null
+++ b/tests/reranker/test_cosine_local.py
@@ -0,0 +1,129 @@
+"""Tests for CosineLocalReranker numerical stability."""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+from memos.reranker.cosine_local import CosineLocalReranker, _cosine_one_to_many
+
+
+class MemoryStub:
+    def __init__(self, memory_id, embedding, background="fact"):
+        self.id = memory_id
+        self.memory = f"memory-{memory_id}"
+        self.metadata = SimpleNamespace(embedding=embedding, background=background)
+
+
+def _make_items(embeddings, backgrounds=None):
+    backgrounds = backgrounds or ["fact"] * len(embeddings)
+    return [MemoryStub(i, emb, bg) for i, (emb, bg) in enumerate(zip(embeddings, backgrounds, strict=False))]
+
+
+class TestCosineOneToManyNumerical:
+    def test_basic_identical_vector_returns_one(self):
+        q = [1.0, 0.0, 0.0]
+        m = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]
+        result = _cosine_one_to_many(q, m)
+        assert len(result) == 2
+        assert abs(result[0] - 1.0) < 1e-6
+        assert abs(result[1]) < 1e-6
+
+    def test_orthogonal_returns_zero(self):
+        q = [1.0, 0.0]
+        m = [[0.0, 1.0]]
+        result = _cosine_one_to_many(q, m)
+        assert abs(result[0]) < 1e-6
+
+    def test_opposite_direction(self):
+        q = [1.0, 0.0]
+        m = [[-1.0, 0.0]]
+        result = _cosine_one_to_many(q, m)
+        assert abs(result[0] - (-1.0)) < 1e-6
+
+    def test_empty_matrix_returns_empty(self):
+        q = [1.0, 2.0, 3.0]
+        m = []
+        result = _cosine_one_to_many(q, m)
+        assert result == []
+
+    def test_zero_query_vector(self):
+        q = [0.0, 0.0, 0.0]
+        m = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]
+        result = _cosine_one_to_many(q, m)
+        assert len(result) == 2
+        for r in result:
+            assert r == 0.0
+
+    def test_zero_candidate_vectors(self):
+        q = [1.0, 0.0, 0.0]
+        m = [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]
+        result = _cosine_one_to_many(q, m)
+        assert len(result) == 2
+        assert result[0] == 0.0
+        assert abs(result[1] - 1.0) < 1e-6
+
+    def test_2d_query_vector_flattened(self):
+        q = [[1.0, 0.0, 0.0]]
+        m = [[1.0, 0.0, 0.0]]
+        result = _cosine_one_to_many(q, m)
+        assert len(result) == 1
+        assert abs(result[0] - 1.0) < 1e-6
+
+    def test_no_nan_in_output_normalized_vectors(self):
+        q = [1.0]
+        m = [[0.0]]
+        result = _cosine_one_to_many(q, m)
+        assert not any(r != r for r in result)
+        assert result[0] == 0.0
+
+    def test_multiple_equal_length_vectors(self):
+        q = [1.0, 2.0]
+        m = [[3.0, 4.0], [1.0, 2.0], [0.5, 1.0]]
+        result = _cosine_one_to_many(q, m)
+        assert len(result) == 3
+        assert all(-1.0 - 1e-6 <= r <= 1.0 + 1e-6 for r in result)
+
+
+class TestCosineLocalReranker:
+    def test_empty_graph_results(self):
+        reranker = CosineLocalReranker()
+        assert reranker.rerank("query", [], top_k=5) == []
+
+    def test_rerank_preserves_top_items(self):
+        embeddings = [
+            [1.0, 0.0, 0.0],
+            [0.9, 0.1, 0.0],
+            [0.0, 1.0, 0.0],
+            [0.0, 0.9, 0.1],
+        ]
+        items = _make_items(embeddings)
+        query_emb = [1.0, 0.0, 0.0]
+        reranker = CosineLocalReranker()
+        result = reranker.rerank("any", items, top_k=2, query_embedding=query_emb)
+        ids = [it.id for it, _ in result]
+        assert len(result) == 2
+        assert ids[0] == 0 or ids[0] == 1
+        assert ids[1] == 0 or ids[1] == 1
+
+    def test_rerank_without_embeddings_falls_back(self):
+        items = [MemoryStub(i, None) for i in range(3)]
+        query_emb = [1.0, 0.0, 0.0]
+        reranker = CosineLocalReranker()
+        result = reranker.rerank("any", items, top_k=5, query_embedding=query_emb)
+        assert len(result) == 3
+
+    def test_level_weights_applied(self):
+        embeddings = [[1.0, 0.0], [1.0, 0.0]]
+        backgrounds = ["topic", "fact"]
+        items = _make_items(embeddings, backgrounds)
+        query_emb = [1.0, 0.0]
+        reranker = CosineLocalReranker(
+            level_weights={"topic": 2.0, "concept": 1.0, "fact": 0.5},
+            level_field="background",
+        )
+        result = reranker.rerank("any", items, top_k=2, query_embedding=query_emb)
+        ids = [it.id for it, _ in result]
+        scores = {it.id: sc for it, sc in result}
+        assert ids[0] == 0
+        assert scores[0] > scores[1]

From a96a4211584887947c83c2844e27240e3ad4ccfd Mon Sep 17 00:00:00 2001
From: RerankerGuo <121015044+RerankerGuo@users.noreply.github.com>
Date: Fri, 31 Jul 2026 10:07:07 +0800
Subject: [PATCH 2/2] style(test): ruff format + PLR0124 fix for cosine
 reranker tests

---
 tests/reranker/test_cosine_local.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/tests/reranker/test_cosine_local.py b/tests/reranker/test_cosine_local.py
index 090eb4e3e..f4c6a157a 100644
--- a/tests/reranker/test_cosine_local.py
+++ b/tests/reranker/test_cosine_local.py
@@ -1,9 +1,8 @@
 """Tests for CosineLocalReranker numerical stability."""
 
-from types import SimpleNamespace
-from unittest.mock import patch
+import math
 
-import pytest
+from types import SimpleNamespace
 
 from memos.reranker.cosine_local import CosineLocalReranker, _cosine_one_to_many
 
@@ -17,7 +16,10 @@ def __init__(self, memory_id, embedding, background="fact"):
 
 def _make_items(embeddings, backgrounds=None):
     backgrounds = backgrounds or ["fact"] * len(embeddings)
-    return [MemoryStub(i, emb, bg) for i, (emb, bg) in enumerate(zip(embeddings, backgrounds, strict=False))]
+    return [
+        MemoryStub(i, emb, bg)
+        for i, (emb, bg) in enumerate(zip(embeddings, backgrounds, strict=False))
+    ]
 
 
 class TestCosineOneToManyNumerical:
@@ -74,7 +76,7 @@ def test_no_nan_in_output_normalized_vectors(self):
         q = [1.0]
         m = [[0.0]]
         result = _cosine_one_to_many(q, m)
-        assert not any(r != r for r in result)
+        assert not any(math.isnan(r) for r in result)
         assert result[0] == 0.0
 
     def test_multiple_equal_length_vectors(self):