From b90e3c0b76bd86b22ce4eaa86f027737e0271aa7 Mon Sep 17 00:00:00 2001 From: Yash Singh <123385188+yashs33244@users.noreply.github.com> Date: Mon, 15 Jun 2026 16:10:00 +0530 Subject: [PATCH] fix(reranker): respect config.top_k in Cohere and ZeroEntropy fallback paths (#5560) --- mem0/reranker/cohere_reranker.py | 3 +- mem0/reranker/zero_entropy_reranker.py | 3 +- .../rerankers/test_reranker_fallback_topk.py | 109 ++++++++++++++++++ 3 files changed, 113 insertions(+), 2 deletions(-) create mode 100644 tests/rerankers/test_reranker_fallback_topk.py diff --git a/mem0/reranker/cohere_reranker.py b/mem0/reranker/cohere_reranker.py index 665d85ba5..8de2d4ac9 100644 --- a/mem0/reranker/cohere_reranker.py +++ b/mem0/reranker/cohere_reranker.py @@ -82,4 +82,5 @@ class CohereReranker(BaseReranker): # Fallback to original order if reranking fails for doc in documents: doc['rerank_score'] = 0.0 - return documents[:top_k] if top_k else documents \ No newline at end of file + final_top_k = top_k or self.config.top_k + return documents[:final_top_k] if final_top_k else documents \ No newline at end of file diff --git a/mem0/reranker/zero_entropy_reranker.py b/mem0/reranker/zero_entropy_reranker.py index 680228db3..df5762306 100644 --- a/mem0/reranker/zero_entropy_reranker.py +++ b/mem0/reranker/zero_entropy_reranker.py @@ -93,4 +93,5 @@ class ZeroEntropyReranker(BaseReranker): # Fallback to original order if reranking fails for doc in documents: doc['rerank_score'] = 0.0 - return documents[:top_k] if top_k else documents \ No newline at end of file + final_top_k = top_k or self.config.top_k + return documents[:final_top_k] if final_top_k else documents \ No newline at end of file diff --git a/tests/rerankers/test_reranker_fallback_topk.py b/tests/rerankers/test_reranker_fallback_topk.py new file mode 100644 index 000000000..a2767e8ae --- /dev/null +++ b/tests/rerankers/test_reranker_fallback_topk.py @@ -0,0 +1,109 @@ +"""Regression tests for the reranker fallback path honoring ``config.top_k``. + +When the underlying rerank call fails, the reranker falls back to returning the +documents in their original order. That fallback must still respect the +configured ``top_k`` limit, exactly like the success path does. The HuggingFace +and SentenceTransformer rerankers already behave this way; these tests pin the +same contract for the Cohere and ZeroEntropy rerankers. +""" + +import sys +from types import ModuleType +from unittest.mock import MagicMock + +import pytest + +from mem0.configs.rerankers.cohere import CohereRerankerConfig +from mem0.configs.rerankers.zero_entropy import ZeroEntropyRerankerConfig + + +@pytest.fixture +def mock_cohere(monkeypatch): + """Provide a fake ``cohere`` module so CohereReranker imports/constructs.""" + fake_cohere = ModuleType("cohere") + fake_client = MagicMock() + fake_cohere.Client = MagicMock(return_value=fake_client) + monkeypatch.setitem(sys.modules, "cohere", fake_cohere) + + import mem0.reranker.cohere_reranker as cohere_reranker + + monkeypatch.setattr(cohere_reranker, "cohere", fake_cohere, raising=False) + monkeypatch.setattr(cohere_reranker, "COHERE_AVAILABLE", True, raising=False) + return cohere_reranker, fake_client + + +@pytest.fixture +def mock_zero_entropy(monkeypatch): + """Provide a fake ``zeroentropy`` module so ZeroEntropyReranker imports.""" + fake_module = ModuleType("zeroentropy") + fake_client = MagicMock() + fake_module.ZeroEntropy = MagicMock(return_value=fake_client) + monkeypatch.setitem(sys.modules, "zeroentropy", fake_module) + + import mem0.reranker.zero_entropy_reranker as zero_entropy_reranker + + monkeypatch.setattr(zero_entropy_reranker, "ZeroEntropy", fake_module.ZeroEntropy, raising=False) + monkeypatch.setattr(zero_entropy_reranker, "ZERO_ENTROPY_AVAILABLE", True, raising=False) + return zero_entropy_reranker, fake_client + + +def _docs(n): + return [{"memory": f"doc{i}"} for i in range(n)] + + +class TestCohereFallbackTopK: + def test_fallback_respects_config_top_k(self, mock_cohere): + module, fake_client = mock_cohere + fake_client.rerank.side_effect = RuntimeError("API error") + + reranker = module.CohereReranker(CohereRerankerConfig(api_key="test-key", top_k=2)) + result = reranker.rerank("query", _docs(5)) + + assert len(result) == 2 + + def test_fallback_per_call_top_k_overrides_config(self, mock_cohere): + module, fake_client = mock_cohere + fake_client.rerank.side_effect = RuntimeError("API error") + + reranker = module.CohereReranker(CohereRerankerConfig(api_key="test-key", top_k=4)) + result = reranker.rerank("query", _docs(5), top_k=1) + + assert len(result) == 1 + + def test_fallback_returns_all_when_no_top_k(self, mock_cohere): + module, fake_client = mock_cohere + fake_client.rerank.side_effect = RuntimeError("API error") + + reranker = module.CohereReranker(CohereRerankerConfig(api_key="test-key")) + result = reranker.rerank("query", _docs(5)) + + assert len(result) == 5 + + +class TestZeroEntropyFallbackTopK: + def test_fallback_respects_config_top_k(self, mock_zero_entropy): + module, fake_client = mock_zero_entropy + fake_client.models.rerank.side_effect = RuntimeError("API error") + + reranker = module.ZeroEntropyReranker(ZeroEntropyRerankerConfig(api_key="test-key", top_k=2)) + result = reranker.rerank("query", _docs(5)) + + assert len(result) == 2 + + def test_fallback_per_call_top_k_overrides_config(self, mock_zero_entropy): + module, fake_client = mock_zero_entropy + fake_client.models.rerank.side_effect = RuntimeError("API error") + + reranker = module.ZeroEntropyReranker(ZeroEntropyRerankerConfig(api_key="test-key", top_k=4)) + result = reranker.rerank("query", _docs(5), top_k=1) + + assert len(result) == 1 + + def test_fallback_returns_all_when_no_top_k(self, mock_zero_entropy): + module, fake_client = mock_zero_entropy + fake_client.models.rerank.side_effect = RuntimeError("API error") + + reranker = module.ZeroEntropyReranker(ZeroEntropyRerankerConfig(api_key="test-key")) + result = reranker.rerank("query", _docs(5)) + + assert len(result) == 5