langchain-reranker 0.1.0__tar.gz

langchain_reranker-0.1.0/PKG-INFO

@@ -0,0 +1,72 @@
+Metadata-Version: 2.3
+Name: langchain-reranker
+Version: 0.1.0
+Summary: Add init_reranker to the LangChain ecosystem
+Keywords: langchain,reranker,rerank,rag,retrieval
+Author: Ashitaka-Laputa-IV
+Author-email: Ashitaka-Laputa-IV <Ashitaka-Laputa-IV@outlook.com>
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Dist: langchain-core>=1.0
+Requires-Python: >=3.13
+Description-Content-Type: text/markdown
+
+# langchain-reranker
+
+Add `init_reranker` to the LangChain ecosystem, following the same conventions as
+`init_chat_model` and `init_embeddings`.
+
+## Installation
+
+```bash
+pip install langchain-reranker
+```
+
+Provider integration packages are optional and installed on demand, e.g.
+`pip install langchain-cohere`.
+
+## Usage
+
+```python
+from langchain_reranker.rerankers import init_reranker
+
+reranker = init_reranker("cohere:rerank-v3.5")
+
+compressed = reranker.compress_documents(documents, query="what is langchain?")
+```
+
+`init_reranker` returns a `BaseDocumentCompressor`, so it plugs directly into
+existing retrieval pipelines such as a contextual compression retriever.
+
+```python
+reranker = init_reranker(model="jina-reranker-v2-base-multilingual", provider="jina")
+```
+
+## Custom rerankers
+
+Implement `rerank` and inherit `compress_documents` for free.
+
+```python
+from langchain_reranker.rerankers import BaseReranker, RerankResult
+
+
+class MyReranker(BaseReranker):
+    model: str = "my-reranker"
+
+    def rerank(self, query, documents, top_n=None):
+        scores = my_api(query, [doc.page_content for doc in documents])
+        ranked = sorted(zip(documents, scores), key=lambda pair: -pair[1])
+        return [
+            RerankResult(index=i, relevance_score=float(score), document=doc)
+            for i, (doc, score) in enumerate(ranked)
+        ]
+```
+
+## Local cross encoders
+
+```python
+reranker = init_reranker("huggingface:BAAI/bge-reranker-v2-m3")
+```

langchain_reranker-0.1.0/README.md

@@ -0,0 +1,56 @@
+# langchain-reranker
+
+Add `init_reranker` to the LangChain ecosystem, following the same conventions as
+`init_chat_model` and `init_embeddings`.
+
+## Installation
+
+```bash
+pip install langchain-reranker
+```
+
+Provider integration packages are optional and installed on demand, e.g.
+`pip install langchain-cohere`.
+
+## Usage
+
+```python
+from langchain_reranker.rerankers import init_reranker
+
+reranker = init_reranker("cohere:rerank-v3.5")
+
+compressed = reranker.compress_documents(documents, query="what is langchain?")
+```
+
+`init_reranker` returns a `BaseDocumentCompressor`, so it plugs directly into
+existing retrieval pipelines such as a contextual compression retriever.
+
+```python
+reranker = init_reranker(model="jina-reranker-v2-base-multilingual", provider="jina")
+```
+
+## Custom rerankers
+
+Implement `rerank` and inherit `compress_documents` for free.
+
+```python
+from langchain_reranker.rerankers import BaseReranker, RerankResult
+
+
+class MyReranker(BaseReranker):
+    model: str = "my-reranker"
+
+    def rerank(self, query, documents, top_n=None):
+        scores = my_api(query, [doc.page_content for doc in documents])
+        ranked = sorted(zip(documents, scores), key=lambda pair: -pair[1])
+        return [
+            RerankResult(index=i, relevance_score=float(score), document=doc)
+            for i, (doc, score) in enumerate(ranked)
+        ]
+```
+
+## Local cross encoders
+
+```python
+reranker = init_reranker("huggingface:BAAI/bge-reranker-v2-m3")
+```

langchain_reranker-0.1.0/pyproject.toml

@@ -0,0 +1,32 @@
+[project]
+name = "langchain-reranker"
+version = "0.1.0"
+description = "Add init_reranker to the LangChain ecosystem"
+readme = "README.md"
+requires-python = ">=3.13"
+keywords = [
+    "langchain",
+    "reranker",
+    "rerank",
+    "rag",
+    "retrieval",
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = ["langchain-core>=1.0"]
+
+[[project.authors]]
+name = "Ashitaka-Laputa-IV"
+email = "Ashitaka-Laputa-IV@outlook.com"
+
+[dependency-groups]
+dev = ["pytest>=8"]
+
+[build-system]
+requires = ["uv_build>=0.12.23,<0.13.0"]
+build-backend = "uv_build"

langchain_reranker-0.1.0/pyproject.toml.orig

@@ -0,0 +1,29 @@
+[project]
+name = "langchain-reranker"
+version = "0.1.0"
+description = "Add init_reranker to the LangChain ecosystem"
+readme = "README.md"
+authors = [
+    { name = "Ashitaka-Laputa-IV", email = "Ashitaka-Laputa-IV@outlook.com" }
+]
+requires-python = ">=3.13"
+keywords = ["langchain", "reranker", "rerank", "rag", "retrieval"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "langchain-core>=1.0",
+]
+
+[dependency-groups]
+dev = [
+    "pytest>=8",
+]
+
+[build-system]
+requires = ["uv_build>=0.12.23,<0.13.0"]
+build-backend = "uv_build"

langchain_reranker-0.1.0/src/langchain_reranker/__init__.py

@@ -0,0 +1,3 @@
+"""Main entrypoint into LangChain rerankers."""
+
+__version__ = "0.1.0"

langchain_reranker-0.1.0/src/langchain_reranker/rerankers/__init__.py

@@ -0,0 +1,15 @@
+"""Entrypoint to using rerankers in LangChain."""
+
+from langchain_reranker.rerankers.base import (
+    BaseReranker,
+    CrossEncoderReranker,
+    RerankResult,
+    init_reranker,
+)
+
+__all__ = [
+    "BaseReranker",
+    "CrossEncoderReranker",
+    "RerankResult",
+    "init_reranker",
+]

langchain_reranker-0.1.0/src/langchain_reranker/rerankers/base.py

@@ -0,0 +1,419 @@
+"""Factory function and base class for rerankers."""
+
+from __future__ import annotations
+
+import functools
+import importlib
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Sequence
+from typing import Any, ClassVar, override
+
+from langchain_core.callbacks import Callbacks
+from langchain_core.cross_encoders import BaseCrossEncoder
+from langchain_core.documents import Document
+from langchain_core.documents.compressor import BaseDocumentCompressor
+from pydantic import BaseModel, ConfigDict
+
+
+class RerankResult(BaseModel):
+    """A single reranked document together with its relevance score."""
+
+    index: int
+    """Index of the document in the input sequence."""
+
+    relevance_score: float
+    """Relevance score assigned by the reranker. Higher means more relevant."""
+
+    document: Document
+    """The reranked document."""
+
+
+class BaseReranker(BaseDocumentCompressor, ABC):
+    """Base class for rerankers.
+
+    Subclasses implement [`rerank`][langchain_reranker.rerankers.BaseReranker.rerank].
+    The inherited `compress_documents` method is implemented in terms of `rerank`, so a
+    reranker can be used anywhere a
+    [`BaseDocumentCompressor`][langchain_core.documents.compressor.BaseDocumentCompressor]
+    is expected (for example inside a contextual compression retriever).
+    """
+
+    model: str = ""
+    """Name of the reranking model."""
+
+    top_n: int | None = None
+    """Number of documents to return. If `None`, all documents are returned."""
+
+    @abstractmethod
+    def rerank(
+        self,
+        query: str,
+        documents: Sequence[Document],
+        top_n: int | None = None,
+    ) -> list[RerankResult]:
+        """Rerank `documents` by relevance to `query`.
+
+        Args:
+            query: The query to rank the documents against.
+            documents: The documents to rerank.
+            top_n: Number of results to return. Defaults to `self.top_n`.
+
+        Returns:
+            The reranked results, ordered from most to least relevant.
+        """
+
+    @override
+    def compress_documents(
+        self,
+        documents: Sequence[Document],
+        query: str,
+        callbacks: Callbacks | None = None,
+    ) -> Sequence[Document]:
+        """Rerank documents and attach the relevance score to their metadata."""
+        _ = callbacks
+        return [
+            result.document.model_copy(
+                update={
+                    "metadata": {
+                        **result.document.metadata,
+                        "relevance_score": result.relevance_score,
+                    }
+                }
+            )
+            for result in self.rerank(query=query, documents=documents, top_n=self.top_n)
+        ]
+
+
+class CrossEncoderReranker(BaseReranker):
+    """Reranker backed by a cross encoder.
+
+    Example:
+        ```python
+        from langchain_community.cross_encoders import HuggingFaceCrossEncoder
+        from langchain_reranker.rerankers import CrossEncoderReranker
+
+        cross_encoder = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-base")
+        reranker = CrossEncoderReranker(cross_encoder=cross_encoder, top_n=3)
+        reranker.rerank(query="...", documents=documents)
+        ```
+    """
+
+    cross_encoder: BaseCrossEncoder
+    """Cross encoder used to score query-document pairs."""
+
+    model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True)
+
+    @override
+    def rerank(
+        self,
+        query: str,
+        documents: Sequence[Document],
+        top_n: int | None = None,
+    ) -> list[RerankResult]:
+        """Score each document against the query and sort by relevance."""
+        scores = self.cross_encoder.score(
+            [(query, document.page_content) for document in documents]
+        )
+        results = [
+            RerankResult(
+                index=index,
+                relevance_score=float(score),
+                document=document,
+            )
+            for index, (document, score) in enumerate(
+                zip(documents, scores, strict=False)
+            )
+        ]
+        results.sort(key=lambda result: result.relevance_score, reverse=True)
+        limit = self.top_n if top_n is None else top_n
+        if limit is not None:
+            results = results[:limit]
+        return results
+
+
+def _call(
+    cls: type[Any],
+    model: str,
+    **kwargs: Any,
+) -> BaseDocumentCompressor:
+    return cls(model=model, **kwargs)
+
+
+def _call_cross_encoder(
+    cls: type[Any],
+    model: str,
+    *,
+    top_n: int | None = None,
+    **kwargs: Any,
+) -> BaseReranker:
+    return CrossEncoderReranker(
+        cross_encoder=cls(model_name=model, **kwargs),
+        model=model,
+        top_n=top_n,
+    )
+
+
+_BUILTIN_PROVIDERS: dict[str, tuple[str, str, Callable[..., BaseDocumentCompressor]]] = {
+    "bedrock": ("langchain_aws", "AmazonBedrockRerank", _call),
+    "cohere": ("langchain_cohere", "CohereRerank", _call),
+    "dashscope": (
+        "langchain_community.document_compressors",
+        "DashScopeRerank",
+        _call,
+    ),
+    "flashrank": (
+        "langchain_community.document_compressors",
+        "FlashrankRerank",
+        _call,
+    ),
+    "huggingface": (
+        "langchain_huggingface",
+        "HuggingFaceCrossEncoder",
+        _call_cross_encoder,
+    ),
+    "infinity": (
+        "langchain_community.document_compressors",
+        "InfinityRerank",
+        _call,
+    ),
+    "jina": ("langchain_community.document_compressors", "JinaRerank", _call),
+    "local": (
+        "langchain_community.cross_encoders",
+        "HuggingFaceCrossEncoder",
+        _call_cross_encoder,
+    ),
+    "openvino": (
+        "langchain_community.document_compressors",
+        "OpenVINOReranker",
+        _call,
+    ),
+    "rankllm": (
+        "langchain_community.document_compressors",
+        "RankLLMRerank",
+        _call,
+    ),
+    "volcengine": (
+        "langchain_community.document_compressors",
+        "VolcengineRerank",
+        _call,
+    ),
+    "voyageai": ("langchain_voyageai", "VoyageAIRerank", _call),
+}
+"""Registry mapping provider names to their import configuration.
+
+Each entry maps a provider key to a tuple of:
+
+- `module_path`: The Python module path containing the reranker class.
+- `class_name`: The name of the reranker class to import.
+- `creator_func`: A callable that instantiates the class with provided kwargs.
+
+!!! note
+
+    This dict is not exhaustive of all providers supported by LangChain, but is
+    meant to cover the most popular ones and serve as a template for adding more
+    providers in the future. If a provider is not in this dict, it can still be
+    used with `init_reranker` as long as its integration package is installed, but
+    the provider key will not be inferred from the model name and must be
+    specified explicitly via the `provider` parameter.
+"""
+
+
+@functools.lru_cache(maxsize=len(_BUILTIN_PROVIDERS))
+def _get_reranker_class_creator(
+    provider: str,
+) -> Callable[..., BaseDocumentCompressor]:
+    """Return a factory that creates a reranker for the given provider.
+
+    This function is cached to avoid repeated module imports.
+
+    Args:
+        provider: The name of the provider (e.g., `'cohere'`, `'jina'`). Must be a
+            key in `_BUILTIN_PROVIDERS`.
+
+    Returns:
+        A callable that accepts model kwargs and returns a
+            `BaseDocumentCompressor` instance for the specified provider.
+
+    Raises:
+        ValueError: If the provider is not in `_BUILTIN_PROVIDERS`.
+        ImportError: If the provider's integration package is not installed.
+    """
+    if provider not in _BUILTIN_PROVIDERS:
+        msg = (
+            f"Provider '{provider}' is not supported.\n"
+            f"Supported providers and their required packages:\n"
+            f"{_get_provider_list()}"
+        )
+        raise ValueError(msg)
+
+    module_name, class_name, creator_func = _BUILTIN_PROVIDERS[provider]
+    try:
+        module = importlib.import_module(module_name)
+    except ImportError as e:
+        pkg = module_name.split(".", maxsplit=1)[0].replace("_", "-")
+        msg = f"Could not import {pkg} python package. Please install it with `pip install {pkg}`"
+        raise ImportError(msg) from e
+
+    cls = getattr(module, class_name)
+    return functools.partial(creator_func, cls=cls)
+
+
+def _get_provider_list() -> str:
+    """Get formatted list of providers and their packages."""
+    return "\n".join(
+        f"  - {p}: {pkg[0].split('.', maxsplit=1)[0].replace('_', '-')}"
+        for p, pkg in _BUILTIN_PROVIDERS.items()
+    )
+
+
+def _parse_model_string(model_name: str) -> tuple[str, str]:
+    """Parse a model string into provider and model name components.
+
+    The model string should be in the format `'provider:model-name'`, where
+    provider is one of the supported providers.
+
+    Args:
+        model_name: A model string in the format `'provider:model-name'`.
+
+    Returns:
+        A tuple of `(provider, model_name)`.
+
+    Raises:
+        ValueError: If the model string is not in the correct format or the
+            provider is unsupported.
+    """
+    if ":" not in model_name:
+        msg = (
+            f"Invalid model format '{model_name}'.\n"
+            f"Model name must be in format 'provider:model-name'\n"
+            f"Example valid model strings:\n"
+            f"  - cohere:rerank-v3.5\n"
+            f"  - jina:jina-reranker-v2-base-multilingual\n"
+            f"  - huggingface:BAAI/bge-reranker-v2-m3\n"
+            f"Supported providers: {_BUILTIN_PROVIDERS.keys()}"
+        )
+        raise ValueError(msg)
+
+    provider, model = model_name.split(":", 1)
+    provider = provider.lower().strip()
+    model = model.strip()
+
+    if provider not in _BUILTIN_PROVIDERS:
+        msg = (
+            f"Provider '{provider}' is not supported.\n"
+            f"Supported providers and their required packages:\n"
+            f"{_get_provider_list()}"
+        )
+        raise ValueError(msg)
+    if not model:
+        msg = "Model name cannot be empty"
+        raise ValueError(msg)
+    return provider, model
+
+
+def _infer_model_and_provider(
+    model: str,
+    *,
+    provider: str | None = None,
+) -> tuple[str, str]:
+    if not model.strip():
+        msg = "Model name cannot be empty"
+        raise ValueError(msg)
+    if provider is None and ":" in model:
+        provider, model_name = _parse_model_string(model)
+    else:
+        model_name = model
+
+    if not provider:
+        msg = (
+            "Must specify either:\n"
+            "1. A model string in format 'provider:model-name'\n"
+            "   Example: 'cohere:rerank-v3.5'\n"
+            "2. Or explicitly set provider from: "
+            f"{_BUILTIN_PROVIDERS.keys()}"
+        )
+        raise ValueError(msg)
+
+    if provider not in _BUILTIN_PROVIDERS:
+        msg = (
+            f"Provider '{provider}' is not supported.\n"
+            f"Supported providers and their required packages:\n"
+            f"{_get_provider_list()}"
+        )
+        raise ValueError(msg)
+    return provider, model_name
+
+
+def init_reranker(
+    model: str,
+    *,
+    provider: str | None = None,
+    **kwargs: Any,
+) -> BaseDocumentCompressor:
+    """Initialize a reranker from a model name and optional provider.
+
+    !!! note
+
+        Requires the integration package for the chosen provider to be installed.
+
+        See the `provider` parameter below for specific package names
+        (e.g., `pip install langchain-cohere`).
+
+    Args:
+        model: The name of the model, e.g. `'cohere:rerank-v3.5'`.
+
+            You can also specify model and provider in a single argument using
+            the `'{provider}:{model}'` format, e.g. `'cohere:rerank-v3.5'`.
+        provider: The provider if not specified as part of the model arg.
+
+            Supported `provider` values and the corresponding integration package
+            are:
+
+            - `bedrock`      -> `langchain-aws`
+            - `cohere`       -> `langchain-cohere`
+            - `dashscope`    -> `langchain-community`
+            - `flashrank`    -> `langchain-community`
+            - `huggingface`  -> `langchain-huggingface`
+            - `infinity`     -> `langchain-community`
+            - `jina`         -> `langchain-community`
+            - `local`        -> `langchain-community` (sentence-transformers)
+            - `openvino`     -> `langchain-community`
+            - `rankllm`      -> `langchain-community`
+            - `volcengine`   -> `langchain-community`
+            - `voyageai`     -> `langchain-voyageai`
+
+        **kwargs: Additional model-specific parameters passed to the reranker.
+
+    Returns:
+        A `BaseDocumentCompressor` that reranks documents for a query.
+
+    Raises:
+        ValueError: If the provider is not supported or cannot be determined.
+        ImportError: If the required provider package is not installed.
+
+    ???+ example
+
+        ```python
+        # pip install langchain-reranker langchain-cohere
+
+        # Using a model string
+        reranker = init_reranker("cohere:rerank-v3.5")
+        reranker.compress_documents(documents, query="...")
+
+        # Using an explicit provider
+        reranker = init_reranker(model="jina-reranker-v2-base-multilingual", provider="jina")
+
+        # With additional parameters
+        reranker = init_reranker("cohere:rerank-v3.5", top_n=5)
+        ```
+    """
+    if not model:
+        providers = ", ".join(_BUILTIN_PROVIDERS.keys())
+        msg = f"Must specify model name. Supported providers are: {providers}"
+        raise ValueError(msg)
+
+    provider, model_name = _infer_model_and_provider(model, provider=provider)
+    return _get_reranker_class_creator(provider)(model=model_name, **kwargs)
+
+
+__all__ = ["BaseReranker", "CrossEncoderReranker", "RerankResult", "init_reranker"]