add a decorator for registering custom rerankers for use as a library.

Author: Zhe Yu

src/vectorcode/subcommands/query/reranker/__init__.py

@@ -1,5 +1,6 @@
 import logging
 import sys
+from typing import Type
 
 from vectorcode.cli_utils import Config
 
@@ -11,11 +12,47 @@
 
 logger = logging.getLogger(name=__name__)
 
+__supported_rerankers: dict[str, Type[RerankerBase]] = {
+    "CrossEncoderReranker": CrossEncoderReranker,
+    "NaiveReranker": NaiveReranker,
+}
+
+
+def add_reranker(cls):
+    """
+    This is a class decorator that allows you to add a custom reranker that can be
+    recognised by the `get_reranker` function.
+
+    Your reranker should inherit `RerankerBase` and be decorated by `add_reranker`:
+    ```python
+    @add_reranker
+    class CustomReranker(RerankerBase):
+        # override the methods according to your need.
+    ```
+    """
+    if issubclass(cls, RerankerBase):
+        if __supported_rerankers.get(cls.__name__):
+            error_message = f"{cls.__name__} has been registered."
+            logger.error(error_message)
+            raise AttributeError(error_message)
+        __supported_rerankers[cls.__name__] = cls
+        return cls
+    else:
+        error_message = f'{cls} should be a subclass of "RerankerBase"'
+        logger.error(error_message)
+        raise TypeError(error_message)
+
 
 def get_reranker(configs: Config) -> RerankerBase:
-    if configs.reranker and hasattr(sys.modules[__name__], configs.reranker):
-        # dynamic dispatch
-        return getattr(sys.modules[__name__], configs.reranker)(configs)
+    if configs.reranker:
+        if hasattr(sys.modules[__name__], configs.reranker):
+            # dynamic dispatch for built-in rerankers
+            return getattr(sys.modules[__name__], configs.reranker)(configs)
+
+        elif issubclass(
+            __supported_rerankers.get(configs.reranker, type(None)), RerankerBase
+        ):
+            return __supported_rerankers[configs.reranker](configs)
 
     # TODO: replace the following with an Exception before the release of 0.6.0.
     logger.warning(

tests/subcommands/query/test_reranker.py

@@ -244,3 +244,66 @@ def test_get_reranker():
         reranker.configs.reranker_params.get("model_name_or_path")
         == "cross-encoder/ms-marco-MiniLM-L-6-v2"
     ), "configs.reranker_params should fallback to default params."
+
+
+def test_supported_rerankers_initialization():
+    """Test that __supported_rerankers contains the expected default rerankers"""
+    from vectorcode.subcommands.query.reranker import __supported_rerankers
+
+    assert "CrossEncoderReranker" in __supported_rerankers
+    assert "NaiveReranker" in __supported_rerankers
+    assert len(__supported_rerankers) == 2
+
+
+def test_add_reranker_success():
+    """Test successful registration of a new reranker"""
+    from vectorcode.subcommands.query.reranker import (
+        RerankerBase,
+        __supported_rerankers,
+        add_reranker,
+    )
+
+    original_count = len(__supported_rerankers)
+
+    @add_reranker
+    class TestReranker(RerankerBase):
+        def rerank(self, results, query_chunks):
+            return []
+
+    assert len(__supported_rerankers) == original_count + 1
+    assert "TestReranker" in __supported_rerankers
+    assert isinstance(get_reranker(Config(reranker="TestReranker")), TestReranker)
+    __supported_rerankers.pop("TestReranker")
+
+
+def test_add_reranker_duplicate():
+    """Test duplicate reranker registration raises error"""
+    from vectorcode.subcommands.query.reranker import (
+        RerankerBase,
+        __supported_rerankers,
+        add_reranker,
+    )
+
+    # First registration should succeed
+    @add_reranker
+    class TestReranker(RerankerBase):
+        def rerank(self, results, query_chunks):
+            return []
+
+    # Second registration should fail
+    with pytest.raises(AttributeError):
+        add_reranker(TestReranker)
+    __supported_rerankers.pop("TestReranker")
+
+
+def test_add_reranker_invalid_baseclass():
+    """Test that non-RerankerBase classes can't be registered"""
+    from vectorcode.subcommands.query.reranker import (
+        add_reranker,
+    )
+
+    with pytest.raises(TypeError):
+
+        @add_reranker
+        class InvalidReranker:
+            pass