diff --git a/mem0/reranker/__init__.py b/mem0/reranker/__init__.py index 33919b788..78823e9ba 100644 --- a/mem0/reranker/__init__.py +++ b/mem0/reranker/__init__.py @@ -4,6 +4,16 @@ Reranker implementations for mem0 search functionality. from .base import BaseReranker from .cohere_reranker import CohereReranker +from .huggingface_reranker import HuggingFaceReranker +from .llm_reranker import LLMReranker from .sentence_transformer_reranker import SentenceTransformerReranker +from .zero_entropy_reranker import ZeroEntropyReranker -__all__ = ["BaseReranker", "CohereReranker", "SentenceTransformerReranker"] \ No newline at end of file +__all__ = [ + "BaseReranker", + "CohereReranker", + "HuggingFaceReranker", + "LLMReranker", + "SentenceTransformerReranker", + "ZeroEntropyReranker", +] \ No newline at end of file diff --git a/tests/rerankers/test_reranker_public_exports.py b/tests/rerankers/test_reranker_public_exports.py new file mode 100644 index 000000000..b854f65be --- /dev/null +++ b/tests/rerankers/test_reranker_public_exports.py @@ -0,0 +1,41 @@ +"""Regression test pinning the public exports of ``mem0.reranker``. + +All five rerankers are first-class providers in ``RerankerFactory``, so all five +classes must be importable from the package root. A regression once dropped the +LLM, HuggingFace, and ZeroEntropy rerankers from ``__init__`` while keeping them +in the factory, so ``from mem0.reranker import LLMReranker`` raised ImportError. +""" + +import mem0.reranker as reranker_pkg + + +def test_all_rerankers_are_importable_from_package_root(): + from mem0.reranker import ( + BaseReranker, + CohereReranker, + HuggingFaceReranker, + LLMReranker, + SentenceTransformerReranker, + ZeroEntropyReranker, + ) + + assert { + BaseReranker, + CohereReranker, + HuggingFaceReranker, + LLMReranker, + SentenceTransformerReranker, + ZeroEntropyReranker, + } + + +def test_all_exported_names_are_present_in_dunder_all(): + expected = { + "BaseReranker", + "CohereReranker", + "HuggingFaceReranker", + "LLMReranker", + "SentenceTransformerReranker", + "ZeroEntropyReranker", + } + assert expected <= set(reranker_pkg.__all__)