feat: Adding VoyageAI reranker (#31)

fzowl · web-flow · commit c6456bb12d09 · 2025-01-22T16:38:03.000+05:30
* Adding VoyageAI reranker

* Adding VoyageAI reranker

* Refactoring due to the comments

* Correcting the docs
diff --git a/examples/reranker_example.py b/examples/reranker_example.py
@@ -0,0 +1,12 @@
+from haystack import Document
+from haystack_integrations.components.rankers.voyage.voyage_text_reranker import VoyageRanker
+
+ranker = VoyageRanker(model="rerank-2", top_k=2)
+
+docs = [Document(content="Paris"), Document(content="Berlin")]
+query = "What is the capital of germany?"
+output = ranker.run(query=query, documents=docs)
+docs = output["documents"]
+
+for doc in docs:
+    print(f"{doc.content} - {doc.score}")
diff --git a/src/haystack_integrations/components/rankers/voyage/__init__.py b/src/haystack_integrations/components/rankers/voyage/__init__.py
@@ -0,0 +1,3 @@
+from haystack_integrations.components.rankers.voyage.ranker import VoyageRanker
+
+__all__ = ["VoyageRanker"]
diff --git a/src/haystack_integrations/components/rankers/voyage/ranker.py b/src/haystack_integrations/components/rankers/voyage/ranker.py
@@ -0,0 +1,188 @@
+import os
+from typing import Any, Dict, List, Optional
+
+from haystack import component, default_from_dict, default_to_dict, Document, logging
+from haystack.utils import Secret, deserialize_secrets_inplace
+from voyageai import Client
+
+logger = logging.getLogger(__name__)
+
+MAX_NUM_DOCS = 1000
+
+
+@component
+class VoyageRanker:
+    """
+    A component for reranking using Voyage models.
+
+    Usage example:
+    ```python
+    from haystack import Document
+    from haystack_integrations.components.rankers.voyage.ranker import VoyageRanker
+
+    ranker = VoyageRanker(model="rerank-2", top_k=2)
+
+    docs = [Document(content="Paris"), Document(content="Berlin")]
+    query = "What is the capital of germany?"
+    output = ranker.run(query=query, documents=docs)
+    docs = output["documents"]
+    ```
+    """
+
+    def __init__(
+        self,
+        api_key: Secret = Secret.from_env_var("VOYAGE_API_KEY"),
+        model: str = "rerank-2",
+        truncate: Optional[bool] = None,
+        top_k: Optional[int] = None,
+        prefix: str = "",
+        suffix: str = "",
+        timeout: Optional[int] = None,
+        max_retries: Optional[int] = None,
+        meta_fields_to_embed: Optional[List[str]] = None,
+        meta_data_separator: str = "\n",
+    ):
+        """
+        Create an VoyageRanker component.
+
+        :param api_key:
+            The VoyageAI API key. It can be explicitly provided or automatically read from the environment variable
+            VOYAGE_API_KEY (recommended).
+        :param model:
+            The name of the Voyage model to use. Defaults to "voyage-2".
+            For more details on the available models,
+            see [Voyage Rerankers documentation](https://docs.voyageai.com/docs/reranker).
+        :param truncate:
+            Whether to truncate the input texts to fit within the context length.
+            - If `True`, over-length input texts will be truncated to fit within the context length, before vectorized
+              by the reranker model.
+            - If False, an error will be raised if any given text exceeds the context length.
+            - Defaults to `None`, which will truncate the input text before sending it to the reranker model if it
+              slightly exceeds the context window length. If it significantly exceeds the context window length, an
+              error will be raised.
+        :param top_k:
+            The number of most relevant documents to return.
+            If not specified, the reranking results of all documents will be returned.
+        :param prefix:
+            A string to add to the beginning of each text.
+        :param suffix:
+            A string to add to the end of each text.
+        :param timeout:
+            Timeout for VoyageAI Client calls, if not set it is inferred from the `VOYAGE_TIMEOUT` environment variable
+            or set to 30.
+        :param max_retries:
+            Maximum retries to establish contact with VoyageAI if it returns an internal error, if not set it is
+            inferred from the `VOYAGE_MAX_RETRIES` environment variable or set to 5.
+        """
+        self.api_key = api_key
+        self.model = model
+        self.top_k = top_k
+        self.truncate = truncate
+        self.prefix = prefix
+        self.suffix = suffix
+        self.meta_fields_to_embed = meta_fields_to_embed or []
+        self.meta_data_separator = meta_data_separator
+
+        if timeout is None:
+            timeout = int(os.environ.get("VOYAGE_TIMEOUT", 30))
+        if max_retries is None:
+            max_retries = int(os.environ.get("VOYAGE_MAX_RETRIES", 5))
+
+        self.client = Client(api_key=api_key.resolve_value(), max_retries=max_retries, timeout=timeout)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Serializes the component to a dictionary.
+
+        :returns:
+            Dictionary with serialized data.
+        """
+        return default_to_dict(
+            self,
+            model=self.model,
+            top_k=self.top_k,
+            truncate=self.truncate,
+            prefix=self.prefix,
+            suffix=self.suffix,
+            api_key=self.api_key.to_dict(),
+            meta_fields_to_embed=self.meta_fields_to_embed,
+            meta_data_separator=self.meta_data_separator,
+        )
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "VoyageRanker":
+        """
+        Deserializes the component from a dictionary.
+
+        :param data:
+            Dictionary to deserialize from.
+        :returns:
+            Deserialized component.
+        """
+        deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"])
+        return default_from_dict(cls, data)
+
+    def _prepare_input_docs(self, documents: List[Document]) -> List[str]:
+        """
+        Prepare the input by concatenating the document text with the metadata fields specified.
+        :param documents:
+            The list of Document objects.
+
+        :return:
+            A list of strings to be given as input to Voyage AI model.
+        """
+        concatenated_input_list = []
+        for doc in documents:
+            meta_values_to_embed = [
+                str(doc.meta[key]) for key in self.meta_fields_to_embed if key in doc.meta and doc.meta.get(key)
+            ]
+            concatenated_input = self.meta_data_separator.join([*meta_values_to_embed, doc.content or ""])
+            concatenated_input_list.append(concatenated_input)
+
+        return concatenated_input_list
+
+    @component.output_types(documents=List[Document])
+    def run(self, query: str, documents: List[Document], top_k: Optional[int] = None):
+        """
+        Use the Voyage AI Reranker to re-rank the list of documents based on the query.
+
+        :param query:
+            Query string.
+        :param documents:
+            List of Documents.
+        :param top_k:
+            The maximum number of Documents you want the Ranker to return.
+        :returns:
+            A dictionary with the following keys:
+            - `documents`: List of Documents most similar to the given query in descending order of similarity.
+
+        :raises ValueError: If `top_k` is not > 0.
+        """
+        top_k = top_k or self.top_k
+        if top_k is not None and top_k <= 0:
+            msg = f"top_k must be > 0, but got {top_k}"
+            raise ValueError(msg)
+
+        input_docs = self._prepare_input_docs(documents)
+        if len(input_docs) > MAX_NUM_DOCS:
+            logger.warning(
+                f"The Voyage AI reranking endpoint only supports {MAX_NUM_DOCS} documents.\
+                The number of documents has been truncated to {MAX_NUM_DOCS} \
+                from {len(input_docs)}."
+            )
+            input_docs = input_docs[:MAX_NUM_DOCS]
+
+        response = self.client.rerank(
+            model=self.model,
+            query=query,
+            documents=input_docs,
+            top_k=top_k,
+        )
+        indices = [output.index for output in response.results]
+        scores = [output.relevance_score for output in response.results]
+        sorted_docs = []
+        for idx, score in zip(indices, scores):
+            doc = documents[idx]
+            doc.score = score
+            sorted_docs.append(documents[idx])
+        return {"documents": sorted_docs}
diff --git a/tests/test_ranker.py b/tests/test_ranker.py
@@ -0,0 +1,182 @@
+import os
+
+import pytest
+from haystack import Document
+from haystack.utils.auth import Secret
+from voyageai.error import InvalidRequestError
+
+from haystack_integrations.components.rankers.voyage import VoyageRanker
+
+
+class TestVoyageTextReranker:
+    @pytest.mark.unit
+    def test_init_default(self, monkeypatch):
+        monkeypatch.setenv("VOYAGE_API_KEY", "fake-api-key")
+        reranker = VoyageRanker()
+
+        assert reranker.client.api_key == "fake-api-key"
+        assert reranker.model == "rerank-2"
+        assert reranker.truncate is None
+        assert reranker.prefix == ""
+        assert reranker.suffix == ""
+        assert reranker.top_k is None
+        assert reranker.meta_fields_to_embed == []
+        assert reranker.meta_data_separator == "\n"
+
+    @pytest.mark.unit
+    def test_init_with_parameters(self):
+        reranker = VoyageRanker(
+            api_key=Secret.from_token("fake-api-key"),
+            model="model",
+            truncate=True,
+            top_k=10,
+            prefix="prefix",
+            suffix="suffix",
+            meta_fields_to_embed=["meta_field_1", "meta_field_2"],
+            meta_data_separator=",",
+        )
+        assert reranker.client.api_key == "fake-api-key"
+        assert reranker.model == "model"
+        assert reranker.truncate is True
+        assert reranker.top_k == 10
+        assert reranker.prefix == "prefix"
+        assert reranker.suffix == "suffix"
+        assert reranker.meta_fields_to_embed == ["meta_field_1", "meta_field_2"]
+        assert reranker.meta_data_separator == ","
+
+    @pytest.mark.unit
+    def test_init_fail_wo_api_key(self, monkeypatch):
+        monkeypatch.delenv("VOYAGE_API_KEY", raising=False)
+        with pytest.raises(ValueError, match="None of the .* environment variables are set"):
+            VoyageRanker()
+
+    @pytest.mark.unit
+    def test_to_dict(self, monkeypatch):
+        monkeypatch.setenv("VOYAGE_API_KEY", "fake-api-key")
+        component = VoyageRanker()
+        data = component.to_dict()
+        assert data == {
+            "type": "haystack_integrations.components.rankers.voyage.ranker."
+            "VoyageRanker",
+            "init_parameters": {
+                "api_key": {"env_vars": ["VOYAGE_API_KEY"], "strict": True, "type": "env_var"},
+                "model": "rerank-2",
+                "truncate": None,
+                "top_k": None,
+                "prefix": "",
+                "suffix": "",
+                "meta_fields_to_embed": [],
+                "meta_data_separator": "\n"
+            },
+        }
+
+    @pytest.mark.unit
+    def test_from_dict(self, monkeypatch):
+        monkeypatch.setenv("VOYAGE_API_KEY", "fake-api-key")
+        data = {
+            "type": "haystack_integrations.components.rankers.voyage.ranker."
+            "VoyageRanker",
+            "init_parameters": {
+                "api_key": {"env_vars": ["VOYAGE_API_KEY"], "strict": True, "type": "env_var"},
+                "model": "rerank-2",
+                "truncate": None,
+                "top_k": 10,
+                "prefix": "",
+                "suffix": "",
+                "meta_fields_to_embed": None,
+                "meta_data_separator": "\n"
+            },
+        }
+
+        reranker = VoyageRanker.from_dict(data)
+        assert reranker.client.api_key == "fake-api-key"
+        assert reranker.top_k == 10
+        assert reranker.model == "rerank-2"
+        assert reranker.truncate is None
+        assert reranker.prefix == ""
+        assert reranker.suffix == ""
+        assert reranker.meta_fields_to_embed == []
+        assert  reranker.meta_data_separator == '\n'
+
+    @pytest.mark.unit
+    def test_to_dict_with_custom_init_parameters(self, monkeypatch):
+        monkeypatch.setenv("ENV_VAR", "fake-api-key")
+        component = VoyageRanker(
+            api_key=Secret.from_env_var("ENV_VAR", strict=False),
+            model="model",
+            truncate=True,
+            top_k=10,
+            prefix="prefix",
+            suffix="suffix",
+        )
+        data = component.to_dict()
+        assert data == {
+            "type": "haystack_integrations.components.rankers.voyage.ranker."
+            "VoyageRanker",
+            "init_parameters": {
+                "api_key": {"env_vars": ["ENV_VAR"], "strict": False, "type": "env_var"},
+                "model": "model",
+                "truncate": True,
+                "top_k": 10,
+                "prefix": "prefix",
+                "suffix": "suffix",
+                'meta_data_separator': '\n',
+                'meta_fields_to_embed': [],
+
+            },
+        }
+
+    @pytest.mark.unit
+    def test_from_dict_with_custom_init_parameters(self, monkeypatch):
+        monkeypatch.setenv("ENV_VAR", "fake-api-key")
+        data = {
+            "type": "haystack_integrations.components.rankers.voyage.ranker."
+            "VoyageRanker",
+            "init_parameters": {
+                "api_key": {"env_vars": ["ENV_VAR"], "strict": False, "type": "env_var"},
+                "model": "model",
+                "truncate": True,
+                "top_k": 10,
+                "prefix": "prefix",
+                "suffix": "suffix",
+            },
+        }
+
+        reranker = VoyageRanker.from_dict(data)
+        assert reranker.client.api_key == "fake-api-key"
+        assert reranker.model == "model"
+        assert reranker.truncate is True
+        assert reranker.top_k == 10
+        assert reranker.prefix == "prefix"
+        assert reranker.suffix == "suffix"
+
+    @pytest.mark.skipif(os.environ.get("VOYAGE_API_KEY", "") == "", reason="VOYAGE_API_KEY is not set")
+    @pytest.mark.integration
+    def test_run(self):
+        model = "rerank-2"
+
+        documents = [
+            Document(id="abcd", content="Paris is in France"),
+            Document(id="efgh", content="Berlin is in Germany"),
+            Document(id="ijkl", content="Lyon is in France"),
+        ]
+
+        reranker = VoyageRanker(model=model, prefix="prefix ", suffix=" suffix")
+        result = reranker.run(query="The food was delicious", documents=documents, top_k=2)
+
+        assert len(result["documents"]) == 2
+        assert all(isinstance(x, Document) for x in result["documents"])
+
+    @pytest.mark.unit
+    def test_run_wrong_input_format(self):
+        reranker = VoyageRanker(api_key=Secret.from_token("fake-api-key"))
+
+        integer_input = 1
+        documents = [
+            Document(id="abcd", content="Paris is in France"),
+            Document(id="efgh", content="Berlin is in Germany"),
+            Document(id="ijkl", content="Lyon is in France"),
+        ]
+
+        with pytest.raises(InvalidRequestError, match=f"not a valid string"):
+            reranker.run(query=integer_input, documents=documents)

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+from haystack_integrations.components.rankers.voyage.ranker import VoyageRanker`
	`2`	`+`
	`3`	`+__all__ = ["VoyageRanker"]`