__init__.py 1.64 KB
"""
Rerank backends - pluggable implementations of the rerank protocol.

Each backend implements score_with_meta(query, docs, normalize) -> (scores, meta).
Service loads one backend via get_rerank_backend(name, config) from config.
"""

from __future__ import annotations

from typing import Any, Dict, List, Protocol, Tuple


class RerankBackendProtocol(Protocol):
    """Protocol for reranker backends (service-internal)."""

    def score_with_meta(
        self,
        query: str,
        docs: List[str],
        normalize: bool = True,
    ) -> Tuple[List[float], Dict[str, Any]]:
        """
        Input:
          query: search query string
          docs: list of documents, scores must align 1:1 with docs
          normalize: whether to normalize scores (e.g. sigmoid)
        Output:
          scores: list same length as docs, same order
          meta: at least input_docs, usable_docs, unique_docs, elapsed_ms
        """
        ...


def get_rerank_backend(name: str, config: Dict[str, Any]) -> RerankBackendProtocol:
    """
    Factory: return a reranker backend instance for the given name and config.
    Config is the corresponding block from services.rerank.backends.<name>.
    """
    name = (name or "bge").strip().lower()
    if name == "bge":
        from reranker.backends.bge import BGERerankerBackend
        return BGERerankerBackend(config)
    if name == "qwen3_vllm":
        from reranker.backends.qwen3_vllm import Qwen3VLLMRerankerBackend
        return Qwen3VLLMRerankerBackend(config)
    raise ValueError(f"Unknown rerank backend: {name!r}. Supported: bge, qwen3_vllm")


__all__ = ["RerankBackendProtocol", "get_rerank_backend"]