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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (37) hide show
  1. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/CHANGELOG.md +27 -0
  2. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/PKG-INFO +2 -1
  3. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/pyproject.toml +12 -2
  4. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/embedders/cohere/document_embedder.py +8 -10
  5. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/embedders/cohere/text_embedder.py +6 -9
  6. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/embedders/cohere/utils.py +4 -2
  7. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/generators/cohere/generator.py +3 -1
  8. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/rankers/cohere/ranker.py +70 -26
  9. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/test_document_embedder.py +6 -6
  10. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/test_ranker.py +78 -4
  11. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/test_text_embedder.py +6 -6
  12. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/.gitignore +0 -0
  13. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/LICENSE.txt +0 -0
  14. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/README.md +0 -0
  15. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/examples/cohere_embedding.py +0 -0
  16. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/examples/cohere_generation.py +0 -0
  17. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/examples/cohere_ranker.py +0 -0
  18. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/pydoc/config_docusaurus.yml +0 -0
  19. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/embedders/cohere/__init__.py +0 -0
  20. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/embedders/cohere/document_image_embedder.py +0 -0
  21. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/embedders/cohere/embedding_types.py +0 -0
  22. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
  23. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/generators/cohere/__init__.py +0 -0
  24. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/generators/cohere/chat/__init__.py +0 -0
  25. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/generators/cohere/chat/chat_generator.py +0 -0
  26. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/generators/py.typed +0 -0
  27. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/rankers/cohere/__init__.py +0 -0
  28. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
  29. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/__init__.py +0 -0
  30. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/conftest.py +0 -0
  31. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/test_chat_generator.py +0 -0
  32. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/test_chat_generator_async.py +0 -0
  33. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/test_chat_generator_chunks.py +0 -0
  34. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/test_document_image_embedder.py +0 -0
  35. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/test_files/apple.jpg +0 -0
  36. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/test_files/sample_pdf_1.pdf +0 -0
  37. {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/test_generator.py +0 -0
@@ -1,5 +1,32 @@
1
1
  # Changelog
2
2
 
3
+ ## [integrations/cohere-v9.0.0] - 2026-04-07
4
+
5
+ ### 🐛 Bug Fixes
6
+
7
+ - [**breaking**] Cohere Embedders - set embed-v4.0 as the default model (#3111)
8
+
9
+ ### 🧪 Testing
10
+
11
+ - Test compatible integrations with python 3.14; update pyproject (#3001)
12
+ - Track test coverage for all integrations (#3065)
13
+
14
+ ### 🧹 Chores
15
+
16
+ - Enforce ruff docstring rules (D102/D103/D205/D209/D213/D417/D419) in first 10 integrations (#3008)
17
+
18
+
19
+ ## [integrations/cohere-v8.1.0] - 2026-03-19
20
+
21
+ ### 🚀 Features
22
+
23
+ - Add list of models supported by Cohere components (#2993)
24
+
25
+ ### 🧹 Chores
26
+
27
+ - Enable ANN ruff ruleset for cohere integration (#2983)
28
+
29
+
3
30
  ## [integrations/cohere-v8.0.1] - 2026-03-17
4
31
 
5
32
  ### 🐛 Bug Fixes
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: cohere-haystack
3
- Version: 8.1.0
3
+ Version: 9.1.0
4
4
  Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/cohere#readme
5
5
  Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues
6
6
  Project-URL: Source, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/cohere
@@ -14,6 +14,7 @@ Classifier: Programming Language :: Python :: 3.10
14
14
  Classifier: Programming Language :: Python :: 3.11
15
15
  Classifier: Programming Language :: Python :: 3.12
16
16
  Classifier: Programming Language :: Python :: 3.13
17
+ Classifier: Programming Language :: Python :: 3.14
17
18
  Classifier: Programming Language :: Python :: Implementation :: CPython
18
19
  Classifier: Programming Language :: Python :: Implementation :: PyPy
19
20
  Requires-Python: >=3.10
@@ -19,6 +19,7 @@ classifiers = [
19
19
  "Programming Language :: Python :: 3.11",
20
20
  "Programming Language :: Python :: 3.12",
21
21
  "Programming Language :: Python :: 3.13",
22
+ "Programming Language :: Python :: 3.14",
22
23
  "Programming Language :: Python :: Implementation :: CPython",
23
24
  "Programming Language :: Python :: Implementation :: PyPy",
24
25
  ]
@@ -66,7 +67,8 @@ dependencies = [
66
67
  unit = 'pytest -m "not integration" {args:tests}'
67
68
  integration = 'pytest -m "integration" {args:tests}'
68
69
  all = 'pytest {args:tests}'
69
- cov-retry = 'pytest --cov=haystack_integrations --reruns 3 --reruns-delay 30 -x {args:tests}'
70
+ unit-cov-retry = 'pytest --cov=haystack_integrations --reruns 3 --reruns-delay 30 -x -m "not integration" {args:tests}'
71
+ integration-cov-append-retry = 'pytest --cov=haystack_integrations --cov-append --reruns 3 --reruns-delay 30 -x -m "integration" {args:tests}'
70
72
 
71
73
  types = """mypy -p haystack_integrations.components.embedders.cohere \
72
74
  -p haystack_integrations.components.generators.cohere \
@@ -88,6 +90,13 @@ select = [
88
90
  "ARG",
89
91
  "B",
90
92
  "C",
93
+ "D102", # Missing docstring in public method
94
+ "D103", # Missing docstring in public function
95
+ "D205", # 1 blank line required between summary line and description
96
+ "D209", # Closing triple quotes go to new line
97
+ "D213", # summary lines must be positioned on the second physical line of the docstring
98
+ "D417", # Missing argument descriptions in the docstring
99
+ "D419", # Docstring is empty
91
100
  "DTZ",
92
101
  "E",
93
102
  "EM",
@@ -137,11 +146,12 @@ ban-relative-imports = "parents"
137
146
 
138
147
  [tool.ruff.lint.per-file-ignores]
139
148
  # Tests can use magic values, assertions, relative imports, and don't need type annotations
140
- "tests/**/*" = ["PLR2004", "S101", "TID252", "ANN"]
149
+ "tests/**/*" = ["PLR2004", "S101", "TID252", "D", "ANN"]
141
150
 
142
151
  [tool.coverage.run]
143
152
  source = ["haystack_integrations"]
144
153
  branch = true
154
+ relative_files = true
145
155
  parallel = false
146
156
 
147
157
 
@@ -49,7 +49,7 @@ class CohereDocumentEmbedder:
49
49
  def __init__(
50
50
  self,
51
51
  api_key: Secret = Secret.from_env_var(["COHERE_API_KEY", "CO_API_KEY"]),
52
- model: str = "embed-english-v2.0",
52
+ model: str = "embed-v4.0",
53
53
  input_type: str = "search_document",
54
54
  api_base_url: str = "https://api.cohere.com",
55
55
  truncate: str = "END",
@@ -61,16 +61,13 @@ class CohereDocumentEmbedder:
61
61
  embedding_type: EmbeddingTypes | None = None,
62
62
  ) -> None:
63
63
  """
64
+ Initialize the CohereDocumentEmbedder.
65
+
64
66
  :param api_key: the Cohere API key.
65
- :param model: the name of the model to use. Supported Models are:
66
- `"embed-english-v3.0"`, `"embed-english-light-v3.0"`, `"embed-multilingual-v3.0"`,
67
- `"embed-multilingual-light-v3.0"`, `"embed-english-v2.0"`, `"embed-english-light-v2.0"`,
68
- `"embed-multilingual-v2.0"`. This list of all supported models can be found in the
69
- [model documentation](https://docs.cohere.com/docs/models#representation).
67
+ :param model: the name of the model to use.
68
+ Read [Cohere documentation](https://docs.cohere.com/docs/models#embed) for a list of all supported models.
70
69
  :param input_type: specifies the type of input you're giving to the model. Supported values are
71
- "search_document", "search_query", "classification" and "clustering". Not
72
- required for older versions of the embedding models (meaning anything lower than v3), but is required for
73
- more recent versions (meaning anything bigger than v2).
70
+ "search_document", "search_query", "classification" and "clustering".
74
71
  :param api_base_url: the Cohere API Base url.
75
72
  :param truncate: truncate embeddings that are too long from start or end, ("NONE"|"START"|"END").
76
73
  Passing "START" will discard the start of the input. "END" will discard the end of the input. In both
@@ -179,7 +176,8 @@ class CohereDocumentEmbedder:
179
176
 
180
177
  @component.output_types(documents=list[Document], meta=dict[str, Any])
181
178
  def run(self, documents: list[Document]) -> dict[str, list[Document] | dict[str, Any]]:
182
- """Embed a list of `Documents`.
179
+ """
180
+ Embed a list of `Documents`.
183
181
 
184
182
  :param documents: documents to embed.
185
183
  :returns: A dictionary with the following keys:
@@ -45,7 +45,7 @@ class CohereTextEmbedder:
45
45
  def __init__(
46
46
  self,
47
47
  api_key: Secret = Secret.from_env_var(["COHERE_API_KEY", "CO_API_KEY"]),
48
- model: str = "embed-english-v2.0",
48
+ model: str = "embed-v4.0",
49
49
  input_type: str = "search_query",
50
50
  api_base_url: str = "https://api.cohere.com",
51
51
  truncate: str = "END",
@@ -53,16 +53,13 @@ class CohereTextEmbedder:
53
53
  embedding_type: EmbeddingTypes | None = None,
54
54
  ) -> None:
55
55
  """
56
+ Initialize the CohereTextEmbedder.
57
+
56
58
  :param api_key: the Cohere API key.
57
- :param model: the name of the model to use. Supported Models are:
58
- `"embed-english-v3.0"`, `"embed-english-light-v3.0"`, `"embed-multilingual-v3.0"`,
59
- `"embed-multilingual-light-v3.0"`, `"embed-english-v2.0"`, `"embed-english-light-v2.0"`,
60
- `"embed-multilingual-v2.0"`. This list of all supported models can be found in the
61
- [model documentation](https://docs.cohere.com/docs/models#representation).
59
+ :param model: the name of the model to use.
60
+ Read [Cohere documentation](https://docs.cohere.com/docs/models#embed) for a list of all supported models.
62
61
  :param input_type: specifies the type of input you're giving to the model. Supported values are
63
- "search_document", "search_query", "classification" and "clustering". Not
64
- required for older versions of the embedding models (meaning anything lower than v3), but is required for
65
- more recent versions (meaning anything bigger than v2).
62
+ "search_document", "search_query", "classification" and "clustering".
66
63
  :param api_base_url: the Cohere API Base url.
67
64
  :param truncate: truncate embeddings that are too long from start or end, ("NONE"|"START"|"END").
68
65
  Passing "START" will discard the start of the input. "END" will discard the end of the input. In both
@@ -18,7 +18,8 @@ async def get_async_response(
18
18
  truncate: str,
19
19
  embedding_type: EmbeddingTypes | None = None,
20
20
  ) -> tuple[list[list[float]], dict[str, Any]]:
21
- """Embeds a list of texts asynchronously using the Cohere API.
21
+ """
22
+ Embeds a list of texts asynchronously using the Cohere API.
22
23
 
23
24
  :param cohere_async_client: the Cohere `AsyncClient`
24
25
  :param texts: the texts to embed
@@ -66,7 +67,8 @@ def get_response(
66
67
  progress_bar: bool = False,
67
68
  embedding_type: EmbeddingTypes | None = None,
68
69
  ) -> tuple[list[list[float]], dict[str, Any]]:
69
- """Embeds a list of texts using the Cohere API.
70
+ """
71
+ Embeds a list of texts using the Cohere API.
70
72
 
71
73
  :param cohere_client: the Cohere `Client`
72
74
  :param texts: the texts to embed
@@ -15,7 +15,8 @@ logger = logging.getLogger(__name__)
15
15
 
16
16
  @component
17
17
  class CohereGenerator(CohereChatGenerator):
18
- """Generates text using Cohere's models through Cohere's `generate` endpoint.
18
+ """
19
+ Generates text using Cohere's models through Cohere's `generate` endpoint.
19
20
 
20
21
  NOTE: Cohere discontinued the `generate` API, so this generator is a mere wrapper
21
22
  around `CohereChatGenerator` provided for backward compatibility.
@@ -115,6 +116,7 @@ class CohereGenerator(CohereChatGenerator):
115
116
  ) -> dict[str, list[str] | list[dict[str, Any]]]:
116
117
  """
117
118
  Queries the LLM asynchronously with the prompts to produce replies.
119
+
118
120
  :param prompt: the prompt to be sent to the generative model.
119
121
  :returns: A dictionary with the following keys:
120
122
  - `replies`: A list of replies generated by the model.
@@ -4,7 +4,7 @@ from typing import Any
4
4
  from haystack import Document, component, default_from_dict, default_to_dict, logging
5
5
  from haystack.utils import Secret, deserialize_secrets_inplace
6
6
 
7
- import cohere
7
+ from cohere import AsyncClientV2, ClientV2
8
8
 
9
9
  logger = logging.getLogger(__name__)
10
10
 
@@ -63,7 +63,10 @@ class CohereRanker:
63
63
  self.meta_data_separator = meta_data_separator
64
64
  self.max_tokens_per_doc = max_tokens_per_doc
65
65
 
66
- self._cohere_client = cohere.ClientV2(
66
+ self._cohere_client = ClientV2(
67
+ api_key=self.api_key.resolve_value(), base_url=self.api_base_url, client_name="haystack"
68
+ )
69
+ self._cohere_async_client = AsyncClientV2(
67
70
  api_key=self.api_key.resolve_value(), base_url=self.api_base_url, client_name="haystack"
68
71
  )
69
72
 
@@ -103,13 +106,21 @@ class CohereRanker:
103
106
  deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"])
104
107
  return default_from_dict(cls, data)
105
108
 
106
- def _prepare_cohere_input_docs(self, documents: list[Document]) -> list[str]:
109
+ def _prepare_cohere_input_docs(self, documents: list[Document], top_k: int | None = None) -> tuple[list[str], int]:
107
110
  """
108
- Prepare the input by concatenating the document text with the metadata fields specified.
111
+ Validate parameters and prepare the input by concatenating the document text with the metadata fields.
112
+
109
113
  :param documents: The list of Document objects.
114
+ :param top_k: The maximum number of documents to return. Falls back to self.top_k if None.
110
115
 
111
- :return: A list of strings to be given as input to Cohere model.
116
+ :return: A tuple of (list of strings to be given as input to Cohere model, resolved top_k).
117
+ :raises ValueError: If `top_k` is not > 0.
112
118
  """
119
+ top_k = top_k or self.top_k
120
+ if top_k <= 0:
121
+ msg = f"top_k must be > 0, but got {top_k}"
122
+ raise ValueError(msg)
123
+
113
124
  concatenated_input_list = []
114
125
  for doc in documents:
115
126
  meta_values_to_embed = [
@@ -118,7 +129,25 @@ class CohereRanker:
118
129
  concatenated_input = self.meta_data_separator.join([*meta_values_to_embed, doc.content or ""])
119
130
  concatenated_input_list.append(concatenated_input)
120
131
 
121
- return concatenated_input_list
132
+ if len(concatenated_input_list) > MAX_NUM_DOCS_FOR_COHERE_RANKER:
133
+ logger.warning(
134
+ f"The Cohere reranking endpoint only supports {MAX_NUM_DOCS_FOR_COHERE_RANKER} documents.\
135
+ The number of documents has been truncated to {MAX_NUM_DOCS_FOR_COHERE_RANKER} \
136
+ from {len(concatenated_input_list)}."
137
+ )
138
+ concatenated_input_list = concatenated_input_list[:MAX_NUM_DOCS_FOR_COHERE_RANKER]
139
+
140
+ return concatenated_input_list, top_k
141
+
142
+ @staticmethod
143
+ def _build_result(response: Any, documents: list[Document]) -> dict[str, list[Document]]:
144
+ indices = [output.index for output in response.results]
145
+ scores = [output.relevance_score for output in response.results]
146
+ sorted_docs = []
147
+ for idx, score in zip(indices, scores, strict=True):
148
+ doc = documents[idx]
149
+ sorted_docs.append(replace(doc, score=score))
150
+ return {"documents": sorted_docs}
122
151
 
123
152
  @component.output_types(documents=list[Document])
124
153
  def run(self, query: str, documents: list[Document], top_k: int | None = None) -> dict[str, list[Document]]:
@@ -137,19 +166,7 @@ class CohereRanker:
137
166
 
138
167
  :raises ValueError: If `top_k` is not > 0.
139
168
  """
140
- top_k = top_k or self.top_k
141
- if top_k <= 0:
142
- msg = f"top_k must be > 0, but got {top_k}"
143
- raise ValueError(msg)
144
-
145
- cohere_input_docs = self._prepare_cohere_input_docs(documents)
146
- if len(cohere_input_docs) > MAX_NUM_DOCS_FOR_COHERE_RANKER:
147
- logger.warning(
148
- f"The Cohere reranking endpoint only supports {MAX_NUM_DOCS_FOR_COHERE_RANKER} documents.\
149
- The number of documents has been truncated to {MAX_NUM_DOCS_FOR_COHERE_RANKER} \
150
- from {len(cohere_input_docs)}."
151
- )
152
- cohere_input_docs = cohere_input_docs[:MAX_NUM_DOCS_FOR_COHERE_RANKER]
169
+ cohere_input_docs, top_k = self._prepare_cohere_input_docs(documents, top_k)
153
170
 
154
171
  response = self._cohere_client.rerank(
155
172
  model=self.model_name,
@@ -158,10 +175,37 @@ class CohereRanker:
158
175
  max_tokens_per_doc=self.max_tokens_per_doc,
159
176
  top_n=top_k,
160
177
  )
161
- indices = [output.index for output in response.results]
162
- scores = [output.relevance_score for output in response.results]
163
- sorted_docs = []
164
- for idx, score in zip(indices, scores, strict=True):
165
- doc = documents[idx]
166
- sorted_docs.append(replace(doc, score=score))
167
- return {"documents": sorted_docs}
178
+ return self._build_result(response, documents)
179
+
180
+ @component.output_types(documents=list[Document])
181
+ async def run_async(
182
+ self, query: str, documents: list[Document], top_k: int | None = None
183
+ ) -> dict[str, list[Document]]:
184
+ """
185
+ Asynchronously re-rank the list of documents based on the query.
186
+
187
+ This is the asynchronous version of the `run` method. It has the same parameters and return values
188
+ but can be used with `await` in async code.
189
+
190
+ :param query:
191
+ Query string.
192
+ :param documents:
193
+ List of Documents.
194
+ :param top_k:
195
+ The maximum number of Documents you want the Ranker to return.
196
+ :returns:
197
+ A dictionary with the following keys:
198
+ - `documents`: List of Documents most similar to the given query in descending order of similarity.
199
+
200
+ :raises ValueError: If `top_k` is not > 0.
201
+ """
202
+ cohere_input_docs, top_k = self._prepare_cohere_input_docs(documents, top_k)
203
+
204
+ response = await self._cohere_async_client.rerank(
205
+ model=self.model_name,
206
+ query=query,
207
+ documents=cohere_input_docs,
208
+ max_tokens_per_doc=self.max_tokens_per_doc,
209
+ top_n=top_k,
210
+ )
211
+ return self._build_result(response, documents)
@@ -26,7 +26,7 @@ class TestCohereDocumentEmbedder:
26
26
  monkeypatch.setenv("COHERE_API_KEY", "test-api-key")
27
27
  embedder = CohereDocumentEmbedder()
28
28
  assert embedder.api_key == Secret.from_env_var(["COHERE_API_KEY", "CO_API_KEY"])
29
- assert embedder.model == "embed-english-v2.0"
29
+ assert embedder.model == "embed-v4.0"
30
30
  assert embedder.input_type == "search_document"
31
31
  assert embedder.api_base_url == COHERE_API_URL
32
32
  assert embedder.truncate == "END"
@@ -70,7 +70,7 @@ class TestCohereDocumentEmbedder:
70
70
  "type": "haystack_integrations.components.embedders.cohere.document_embedder.CohereDocumentEmbedder",
71
71
  "init_parameters": {
72
72
  "api_key": {"env_vars": ["COHERE_API_KEY", "CO_API_KEY"], "strict": True, "type": "env_var"},
73
- "model": "embed-english-v2.0",
73
+ "model": "embed-v4.0",
74
74
  "input_type": "search_document",
75
75
  "api_base_url": COHERE_API_URL,
76
76
  "truncate": "END",
@@ -122,7 +122,7 @@ class TestCohereDocumentEmbedder:
122
122
  "type": "haystack_integrations.components.embedders.cohere.document_embedder.CohereDocumentEmbedder",
123
123
  "init_parameters": {
124
124
  "api_key": {"env_vars": ["COHERE_API_KEY", "CO_API_KEY"], "strict": True, "type": "env_var"},
125
- "model": "embed-english-v2.0",
125
+ "model": "embed-v4.0",
126
126
  "input_type": "search_document",
127
127
  "api_base_url": COHERE_API_URL,
128
128
  "truncate": "END",
@@ -137,7 +137,7 @@ class TestCohereDocumentEmbedder:
137
137
  }
138
138
  embedder = CohereDocumentEmbedder.from_dict(component_dict)
139
139
  assert embedder.api_key == Secret.from_env_var(["COHERE_API_KEY", "CO_API_KEY"])
140
- assert embedder.model == "embed-english-v2.0"
140
+ assert embedder.model == "embed-v4.0"
141
141
  assert embedder.input_type == "search_document"
142
142
  assert embedder.api_base_url == COHERE_API_URL
143
143
  assert embedder.truncate == "END"
@@ -253,7 +253,7 @@ class TestCohereDocumentEmbedder:
253
253
  )
254
254
  @pytest.mark.integration
255
255
  def test_live_run(self):
256
- embedder = CohereDocumentEmbedder(model="embed-english-v2.0", embedding_type=EmbeddingTypes.FLOAT)
256
+ embedder = CohereDocumentEmbedder(model="embed-v4.0", embedding_type=EmbeddingTypes.FLOAT)
257
257
 
258
258
  docs = [
259
259
  Document(content="I love cheese", meta={"topic": "Cuisine"}),
@@ -276,7 +276,7 @@ class TestCohereDocumentEmbedder:
276
276
  )
277
277
  @pytest.mark.integration
278
278
  async def test_live_run_async(self):
279
- embedder = CohereDocumentEmbedder(model="embed-english-v2.0", embedding_type=EmbeddingTypes.FLOAT)
279
+ embedder = CohereDocumentEmbedder(model="embed-v4.0", embedding_type=EmbeddingTypes.FLOAT)
280
280
 
281
281
  docs = [
282
282
  Document(content="I love cheese", meta={"topic": "Cuisine"}),
@@ -35,6 +35,24 @@ def mock_ranker_response():
35
35
  yield mock_ranker_response
36
36
 
37
37
 
38
+ @pytest.fixture
39
+ def mock_async_ranker_response():
40
+ with patch("cohere.AsyncClientV2.rerank", autospec=True) as mock_ranker_response:
41
+ mock_response = Mock()
42
+
43
+ mock_ranker_res_obj1 = Mock()
44
+ mock_ranker_res_obj1.index = 2
45
+ mock_ranker_res_obj1.relevance_score = 0.98
46
+
47
+ mock_ranker_res_obj2 = Mock()
48
+ mock_ranker_res_obj2.index = 1
49
+ mock_ranker_res_obj2.relevance_score = 0.95
50
+
51
+ mock_response.results = [mock_ranker_res_obj1, mock_ranker_res_obj2]
52
+ mock_ranker_response.return_value = mock_response
53
+ yield mock_ranker_response
54
+
55
+
38
56
  class TestCohereRanker:
39
57
  def test_init_default(self, monkeypatch):
40
58
  monkeypatch.setenv("CO_API_KEY", "test-api-key")
@@ -167,7 +185,7 @@ class TestCohereRanker:
167
185
  for i in range(5)
168
186
  ]
169
187
 
170
- texts = component._prepare_cohere_input_docs(documents=documents)
188
+ texts, top_k = component._prepare_cohere_input_docs(documents=documents)
171
189
 
172
190
  assert texts == [
173
191
  "meta_value_1 0\nmeta_value_2 5\ndocument number 0",
@@ -176,6 +194,7 @@ class TestCohereRanker:
176
194
  "meta_value_1 3\nmeta_value_2 8\ndocument number 3",
177
195
  "meta_value_1 4\nmeta_value_2 9\ndocument number 4",
178
196
  ]
197
+ assert top_k == 10
179
198
 
180
199
  def test_prepare_cohere_input_docs_custom_separator(self, monkeypatch):
181
200
  monkeypatch.setenv("CO_API_KEY", "test-api-key")
@@ -192,7 +211,7 @@ class TestCohereRanker:
192
211
  for i in range(5)
193
212
  ]
194
213
 
195
- texts = component._prepare_cohere_input_docs(documents=documents)
214
+ texts, top_k = component._prepare_cohere_input_docs(documents=documents)
196
215
 
197
216
  assert texts == [
198
217
  "meta_value_1 0 meta_value_2 5 document number 0",
@@ -201,13 +220,14 @@ class TestCohereRanker:
201
220
  "meta_value_1 3 meta_value_2 8 document number 3",
202
221
  "meta_value_1 4 meta_value_2 9 document number 4",
203
222
  ]
223
+ assert top_k == 10
204
224
 
205
225
  def test_prepare_cohere_input_docs_no_meta_data(self, monkeypatch):
206
226
  monkeypatch.setenv("CO_API_KEY", "test-api-key")
207
227
  component = CohereRanker(meta_fields_to_embed=["meta_field_1", "meta_field_2"], meta_data_separator=" ")
208
228
  documents = [Document(content=f"document number {i}") for i in range(5)]
209
229
 
210
- texts = component._prepare_cohere_input_docs(documents=documents)
230
+ texts, top_k = component._prepare_cohere_input_docs(documents=documents)
211
231
 
212
232
  assert texts == [
213
233
  "document number 0",
@@ -216,15 +236,17 @@ class TestCohereRanker:
216
236
  "document number 3",
217
237
  "document number 4",
218
238
  ]
239
+ assert top_k == 10
219
240
 
220
241
  def test_prepare_cohere_input_docs_no_docs(self, monkeypatch):
221
242
  monkeypatch.setenv("CO_API_KEY", "test-api-key")
222
243
  component = CohereRanker(meta_fields_to_embed=["meta_field_1", "meta_field_2"], meta_data_separator=" ")
223
244
  documents = []
224
245
 
225
- texts = component._prepare_cohere_input_docs(documents=documents)
246
+ texts, top_k = component._prepare_cohere_input_docs(documents=documents)
226
247
 
227
248
  assert texts == []
249
+ assert top_k == 10
228
250
 
229
251
  def test_run_negative_topk_in_init(self, monkeypatch):
230
252
  monkeypatch.setenv("CO_API_KEY", "test-api-key")
@@ -316,6 +338,34 @@ class TestCohereRanker:
316
338
  for doc in reranked_docs:
317
339
  assert doc.score is not None
318
340
 
341
+ @pytest.mark.asyncio
342
+ async def test_run_async(self, monkeypatch, mock_async_ranker_response): # noqa: ARG002
343
+ monkeypatch.setenv("CO_API_KEY", "test-api-key")
344
+ ranker = CohereRanker(top_k=2)
345
+ query = "test"
346
+ documents = [
347
+ Document(id="abcd", content="doc1", meta={"meta_field": "meta_value_1"}),
348
+ Document(id="efgh", content="doc2", meta={"meta_field": "meta_value_2"}),
349
+ Document(id="ijkl", content="doc3", meta={"meta_field": "meta_value_3"}),
350
+ ]
351
+ ranker_results = await ranker.run_async(query, documents, 2)
352
+
353
+ assert isinstance(ranker_results, dict)
354
+ reranked_docs = ranker_results["documents"]
355
+ assert reranked_docs == [
356
+ Document(id="ijkl", content="doc3", meta={"meta_field": "meta_value_3"}, score=0.98),
357
+ Document(id="efgh", content="doc2", meta={"meta_field": "meta_value_2"}, score=0.95),
358
+ ]
359
+
360
+ @pytest.mark.asyncio
361
+ async def test_run_async_negative_topk(self, monkeypatch):
362
+ monkeypatch.setenv("CO_API_KEY", "test-api-key")
363
+ ranker = CohereRanker(top_k=-2)
364
+ query = "test"
365
+ documents = [Document(content="doc1"), Document(content="doc2"), Document(content="doc3")]
366
+ with pytest.raises(ValueError, match=r"top_k must be > 0, but got *"):
367
+ await ranker.run_async(query, documents)
368
+
319
369
  @pytest.mark.skipif(
320
370
  not os.environ.get("COHERE_API_KEY", None) and not os.environ.get("CO_API_KEY", None),
321
371
  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:
363
413
  assert len(ranker_result["documents"]) == 3
364
414
  assert all(isinstance(doc, Document) for doc in ranker_result["documents"])
365
415
  assert set(result_documents_contents) == set(expected_documents_content)
416
+
417
+ @pytest.mark.skipif(
418
+ not os.environ.get("COHERE_API_KEY", None) and not os.environ.get("CO_API_KEY", None),
419
+ reason="Export an env var called COHERE_API_KEY/CO_API_KEY containing the Cohere API key to run this test.",
420
+ )
421
+ @pytest.mark.integration
422
+ @pytest.mark.asyncio
423
+ async def test_live_run_async(self):
424
+ component = CohereRanker()
425
+ documents = [
426
+ Document(id="abcd", content="Paris is in France"),
427
+ Document(id="efgh", content="Berlin is in Germany"),
428
+ Document(id="ijkl", content="Lyon is in France"),
429
+ ]
430
+
431
+ ranker_result = await component.run_async("Cities in France", documents, 2)
432
+ expected_documents_content = ["Paris is in France", "Lyon is in France"]
433
+ result_documents_contents = [doc.content for doc in ranker_result["documents"]]
434
+
435
+ assert isinstance(ranker_result, dict)
436
+ assert isinstance(ranker_result["documents"], list)
437
+ assert len(ranker_result["documents"]) == 2
438
+ assert all(isinstance(doc, Document) for doc in ranker_result["documents"])
439
+ assert set(result_documents_contents) == set(expected_documents_content)
@@ -28,7 +28,7 @@ class TestCohereTextEmbedder:
28
28
  embedder = CohereTextEmbedder()
29
29
 
30
30
  assert embedder.api_key == Secret.from_env_var(["COHERE_API_KEY", "CO_API_KEY"])
31
- assert embedder.model == "embed-english-v2.0"
31
+ assert embedder.model == "embed-v4.0"
32
32
  assert embedder.input_type == "search_query"
33
33
  assert embedder.api_base_url == COHERE_API_URL
34
34
  assert embedder.truncate == "END"
@@ -65,7 +65,7 @@ class TestCohereTextEmbedder:
65
65
  "type": "haystack_integrations.components.embedders.cohere.text_embedder.CohereTextEmbedder",
66
66
  "init_parameters": {
67
67
  "api_key": {"env_vars": ["COHERE_API_KEY", "CO_API_KEY"], "strict": True, "type": "env_var"},
68
- "model": "embed-english-v2.0",
68
+ "model": "embed-v4.0",
69
69
  "input_type": "search_query",
70
70
  "api_base_url": COHERE_API_URL,
71
71
  "truncate": "END",
@@ -108,7 +108,7 @@ class TestCohereTextEmbedder:
108
108
  "type": "haystack_integrations.components.embedders.cohere.text_embedder.CohereTextEmbedder",
109
109
  "init_parameters": {
110
110
  "api_key": {"env_vars": ["COHERE_API_KEY", "CO_API_KEY"], "strict": True, "type": "env_var"},
111
- "model": "embed-english-v2.0",
111
+ "model": "embed-v4.0",
112
112
  "input_type": "search_query",
113
113
  "api_base_url": COHERE_API_URL,
114
114
  "truncate": "END",
@@ -120,7 +120,7 @@ class TestCohereTextEmbedder:
120
120
 
121
121
  embedder = CohereTextEmbedder.from_dict(component_dict)
122
122
  assert embedder.api_key == Secret.from_env_var(["COHERE_API_KEY", "CO_API_KEY"])
123
- assert embedder.model == "embed-english-v2.0"
123
+ assert embedder.model == "embed-v4.0"
124
124
  assert embedder.input_type == "search_query"
125
125
  assert embedder.api_base_url == COHERE_API_URL
126
126
  assert embedder.truncate == "END"
@@ -148,7 +148,7 @@ class TestCohereTextEmbedder:
148
148
  text = "The food was delicious"
149
149
  result = embedder.run(text=text)
150
150
 
151
- assert len(result["embedding"]) == 4096
151
+ assert len(result["embedding"]) == 1536
152
152
  assert all(isinstance(x, float) for x in result["embedding"])
153
153
 
154
154
  @pytest.mark.asyncio
@@ -162,5 +162,5 @@ class TestCohereTextEmbedder:
162
162
  text = "The food was delicious"
163
163
  result = await embedder.run_async(text=text)
164
164
 
165
- assert len(result["embedding"]) == 4096
165
+ assert len(result["embedding"]) == 1536
166
166
  assert all(isinstance(x, float) for x in result["embedding"])