cohere-haystack 9.0.0__tar.gz → 9.1.0__tar.gz

{cohere_haystack-9.0.0 → cohere_haystack-9.1.0}/CHANGELOG.md

@@ -1,5 +1,21 @@
 # Changelog
 
+## [integrations/cohere-v9.0.0] - 2026-04-07
+
+### 🐛 Bug Fixes
+
+- [**breaking**] Cohere Embedders - set embed-v4.0 as the default model (#3111)
+
+### 🧪 Testing
+
+- Test compatible integrations with python 3.14; update pyproject (#3001)
+- Track test coverage for all integrations (#3065)
+
+### 🧹 Chores
+
+- Enforce ruff docstring rules (D102/D103/D205/D209/D213/D417/D419) in first 10 integrations (#3008)
+
+
 ## [integrations/cohere-v8.1.0] - 2026-03-19
 
 ### 🚀 Features

{cohere_haystack-9.0.0 → cohere_haystack-9.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cohere-haystack
-Version: 9.0.0
+Version: 9.1.0
 Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/cohere#readme
 Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues
 Project-URL: Source, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/cohere

{cohere_haystack-9.0.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/rankers/cohere/ranker.py

@@ -4,7 +4,7 @@ from typing import Any
 from haystack import Document, component, default_from_dict, default_to_dict, logging
 from haystack.utils import Secret, deserialize_secrets_inplace
 
-import cohere
+from cohere import AsyncClientV2, ClientV2
 
 logger = logging.getLogger(__name__)
 
@@ -63,7 +63,10 @@ class CohereRanker:
         self.meta_data_separator = meta_data_separator
         self.max_tokens_per_doc = max_tokens_per_doc
 
-        self._cohere_client = cohere.ClientV2(
+        self._cohere_client = ClientV2(
+            api_key=self.api_key.resolve_value(), base_url=self.api_base_url, client_name="haystack"
+        )
+        self._cohere_async_client = AsyncClientV2(
             api_key=self.api_key.resolve_value(), base_url=self.api_base_url, client_name="haystack"
         )
 
@@ -103,14 +106,21 @@ class CohereRanker:
         deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"])
         return default_from_dict(cls, data)
 
-    def _prepare_cohere_input_docs(self, documents: list[Document]) -> list[str]:
+    def _prepare_cohere_input_docs(self, documents: list[Document], top_k: int | None = None) -> tuple[list[str], int]:
         """
-        Prepare the input by concatenating the document text with the metadata fields specified.
+        Validate parameters and prepare the input by concatenating the document text with the metadata fields.
 
         :param documents: The list of Document objects.
+        :param top_k: The maximum number of documents to return. Falls back to self.top_k if None.
 
-        :return: A list of strings to be given as input to Cohere model.
+        :return: A tuple of (list of strings to be given as input to Cohere model, resolved top_k).
+        :raises ValueError: If `top_k` is not > 0.
         """
+        top_k = top_k or self.top_k
+        if top_k <= 0:
+            msg = f"top_k must be > 0, but got {top_k}"
+            raise ValueError(msg)
+
         concatenated_input_list = []
         for doc in documents:
             meta_values_to_embed = [
@@ -119,7 +129,25 @@ class CohereRanker:
             concatenated_input = self.meta_data_separator.join([*meta_values_to_embed, doc.content or ""])
             concatenated_input_list.append(concatenated_input)
 
-        return concatenated_input_list
+        if len(concatenated_input_list) > MAX_NUM_DOCS_FOR_COHERE_RANKER:
+            logger.warning(
+                f"The Cohere reranking endpoint only supports {MAX_NUM_DOCS_FOR_COHERE_RANKER} documents.\
+                The number of documents has been truncated to {MAX_NUM_DOCS_FOR_COHERE_RANKER} \
+                from {len(concatenated_input_list)}."
+            )
+            concatenated_input_list = concatenated_input_list[:MAX_NUM_DOCS_FOR_COHERE_RANKER]
+
+        return concatenated_input_list, top_k
+
+    @staticmethod
+    def _build_result(response: Any, documents: list[Document]) -> dict[str, list[Document]]:
+        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, strict=True):
+            doc = documents[idx]
+            sorted_docs.append(replace(doc, score=score))
+        return {"documents": sorted_docs}
 
     @component.output_types(documents=list[Document])
     def run(self, query: str, documents: list[Document], top_k: int | None = None) -> dict[str, list[Document]]:
@@ -138,19 +166,7 @@ class CohereRanker:
 
         :raises ValueError: If `top_k` is not > 0.
         """
-        top_k = top_k or self.top_k
-        if top_k <= 0:
-            msg = f"top_k must be > 0, but got {top_k}"
-            raise ValueError(msg)
-
-        cohere_input_docs = self._prepare_cohere_input_docs(documents)
-        if len(cohere_input_docs) > MAX_NUM_DOCS_FOR_COHERE_RANKER:
-            logger.warning(
-                f"The Cohere reranking endpoint only supports {MAX_NUM_DOCS_FOR_COHERE_RANKER} documents.\
-                The number of documents has been truncated to {MAX_NUM_DOCS_FOR_COHERE_RANKER} \
-                from {len(cohere_input_docs)}."
-            )
-            cohere_input_docs = cohere_input_docs[:MAX_NUM_DOCS_FOR_COHERE_RANKER]
+        cohere_input_docs, top_k = self._prepare_cohere_input_docs(documents, top_k)
 
         response = self._cohere_client.rerank(
             model=self.model_name,
@@ -159,10 +175,37 @@ class CohereRanker:
             max_tokens_per_doc=self.max_tokens_per_doc,
             top_n=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, strict=True):
-            doc = documents[idx]
-            sorted_docs.append(replace(doc, score=score))
-        return {"documents": sorted_docs}
+        return self._build_result(response, documents)
+
+    @component.output_types(documents=list[Document])
+    async def run_async(
+        self, query: str, documents: list[Document], top_k: int | None = None
+    ) -> dict[str, list[Document]]:
+        """
+        Asynchronously re-rank the list of documents based on the query.
+
+        This is the asynchronous version of the `run` method. It has the same parameters and return values
+        but can be used with `await` in async code.
+
+        :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.
+        """
+        cohere_input_docs, top_k = self._prepare_cohere_input_docs(documents, top_k)
+
+        response = await self._cohere_async_client.rerank(
+            model=self.model_name,
+            query=query,
+            documents=cohere_input_docs,
+            max_tokens_per_doc=self.max_tokens_per_doc,
+            top_n=top_k,
+        )
+        return self._build_result(response, documents)

{cohere_haystack-9.0.0 → cohere_haystack-9.1.0}/tests/test_ranker.py

@@ -35,6 +35,24 @@ def mock_ranker_response():
         yield mock_ranker_response
 
 
+@pytest.fixture
+def mock_async_ranker_response():
+    with patch("cohere.AsyncClientV2.rerank", autospec=True) as mock_ranker_response:
+        mock_response = Mock()
+
+        mock_ranker_res_obj1 = Mock()
+        mock_ranker_res_obj1.index = 2
+        mock_ranker_res_obj1.relevance_score = 0.98
+
+        mock_ranker_res_obj2 = Mock()
+        mock_ranker_res_obj2.index = 1
+        mock_ranker_res_obj2.relevance_score = 0.95
+
+        mock_response.results = [mock_ranker_res_obj1, mock_ranker_res_obj2]
+        mock_ranker_response.return_value = mock_response
+        yield mock_ranker_response
+
+
 class TestCohereRanker:
     def test_init_default(self, monkeypatch):
         monkeypatch.setenv("CO_API_KEY", "test-api-key")
@@ -167,7 +185,7 @@ class TestCohereRanker:
             for i in range(5)
         ]
 
-        texts = component._prepare_cohere_input_docs(documents=documents)
+        texts, top_k = component._prepare_cohere_input_docs(documents=documents)
 
         assert texts == [
             "meta_value_1 0\nmeta_value_2 5\ndocument number 0",
@@ -176,6 +194,7 @@ class TestCohereRanker:
             "meta_value_1 3\nmeta_value_2 8\ndocument number 3",
             "meta_value_1 4\nmeta_value_2 9\ndocument number 4",
         ]
+        assert top_k == 10
 
     def test_prepare_cohere_input_docs_custom_separator(self, monkeypatch):
         monkeypatch.setenv("CO_API_KEY", "test-api-key")
@@ -192,7 +211,7 @@ class TestCohereRanker:
             for i in range(5)
         ]
 
-        texts = component._prepare_cohere_input_docs(documents=documents)
+        texts, top_k = component._prepare_cohere_input_docs(documents=documents)
 
         assert texts == [
             "meta_value_1 0 meta_value_2 5 document number 0",
@@ -201,13 +220,14 @@ class TestCohereRanker:
             "meta_value_1 3 meta_value_2 8 document number 3",
             "meta_value_1 4 meta_value_2 9 document number 4",
         ]
+        assert top_k == 10
 
     def test_prepare_cohere_input_docs_no_meta_data(self, monkeypatch):
         monkeypatch.setenv("CO_API_KEY", "test-api-key")
         component = CohereRanker(meta_fields_to_embed=["meta_field_1", "meta_field_2"], meta_data_separator=" ")
         documents = [Document(content=f"document number {i}") for i in range(5)]
 
-        texts = component._prepare_cohere_input_docs(documents=documents)
+        texts, top_k = component._prepare_cohere_input_docs(documents=documents)
 
         assert texts == [
             "document number 0",
@@ -216,15 +236,17 @@ class TestCohereRanker:
             "document number 3",
             "document number 4",
         ]
+        assert top_k == 10
 
     def test_prepare_cohere_input_docs_no_docs(self, monkeypatch):
         monkeypatch.setenv("CO_API_KEY", "test-api-key")
         component = CohereRanker(meta_fields_to_embed=["meta_field_1", "meta_field_2"], meta_data_separator=" ")
         documents = []
 
-        texts = component._prepare_cohere_input_docs(documents=documents)
+        texts, top_k = component._prepare_cohere_input_docs(documents=documents)
 
         assert texts == []
+        assert top_k == 10
 
     def test_run_negative_topk_in_init(self, monkeypatch):
         monkeypatch.setenv("CO_API_KEY", "test-api-key")
@@ -316,6 +338,34 @@ class TestCohereRanker:
         for doc in reranked_docs:
             assert doc.score is not None
 
+    @pytest.mark.asyncio
+    async def test_run_async(self, monkeypatch, mock_async_ranker_response):  # noqa: ARG002
+        monkeypatch.setenv("CO_API_KEY", "test-api-key")
+        ranker = CohereRanker(top_k=2)
+        query = "test"
+        documents = [
+            Document(id="abcd", content="doc1", meta={"meta_field": "meta_value_1"}),
+            Document(id="efgh", content="doc2", meta={"meta_field": "meta_value_2"}),
+            Document(id="ijkl", content="doc3", meta={"meta_field": "meta_value_3"}),
+        ]
+        ranker_results = await ranker.run_async(query, documents, 2)
+
+        assert isinstance(ranker_results, dict)
+        reranked_docs = ranker_results["documents"]
+        assert reranked_docs == [
+            Document(id="ijkl", content="doc3", meta={"meta_field": "meta_value_3"}, score=0.98),
+            Document(id="efgh", content="doc2", meta={"meta_field": "meta_value_2"}, score=0.95),
+        ]
+
+    @pytest.mark.asyncio
+    async def test_run_async_negative_topk(self, monkeypatch):
+        monkeypatch.setenv("CO_API_KEY", "test-api-key")
+        ranker = CohereRanker(top_k=-2)
+        query = "test"
+        documents = [Document(content="doc1"), Document(content="doc2"), Document(content="doc3")]
+        with pytest.raises(ValueError, match=r"top_k must be > 0, but got *"):
+            await ranker.run_async(query, documents)
+
     @pytest.mark.skipif(
         not os.environ.get("COHERE_API_KEY", None) and not os.environ.get("CO_API_KEY", None),
         reason="Export an env var called COHERE_API_KEY/CO_API_KEY containing the Cohere API key to run this test.",
@@ -363,3 +413,27 @@ class TestCohereRanker:
         assert len(ranker_result["documents"]) == 3
         assert all(isinstance(doc, Document) for doc in ranker_result["documents"])
         assert set(result_documents_contents) == set(expected_documents_content)
+
+    @pytest.mark.skipif(
+        not os.environ.get("COHERE_API_KEY", None) and not os.environ.get("CO_API_KEY", None),
+        reason="Export an env var called COHERE_API_KEY/CO_API_KEY containing the Cohere API key to run this test.",
+    )
+    @pytest.mark.integration
+    @pytest.mark.asyncio
+    async def test_live_run_async(self):
+        component = CohereRanker()
+        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"),
+        ]
+
+        ranker_result = await component.run_async("Cities in France", documents, 2)
+        expected_documents_content = ["Paris is in France", "Lyon is in France"]
+        result_documents_contents = [doc.content for doc in ranker_result["documents"]]
+
+        assert isinstance(ranker_result, dict)
+        assert isinstance(ranker_result["documents"], list)
+        assert len(ranker_result["documents"]) == 2
+        assert all(isinstance(doc, Document) for doc in ranker_result["documents"])
+        assert set(result_documents_contents) == set(expected_documents_content)