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 .base import BaseReranker
|
||||||
from .cohere_reranker import CohereReranker
|
from .cohere_reranker import CohereReranker
|
||||||
|
from .huggingface_reranker import HuggingFaceReranker
|
||||||
|
from .llm_reranker import LLMReranker
|
||||||
from .sentence_transformer_reranker import SentenceTransformerReranker
|
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