fix(reranker): export all five rerankers from package root (#5636)
This commit is contained in:
@@ -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"]
|
||||
__all__ = [
|
||||
"BaseReranker",
|
||||
"CohereReranker",
|
||||
"HuggingFaceReranker",
|
||||
"LLMReranker",
|
||||
"SentenceTransformerReranker",
|
||||
"ZeroEntropyReranker",
|
||||
]
|
||||
@@ -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__)
|
||||
Reference in New Issue
Block a user