cohere-haystack 9.0.0__tar.gz → 9.2.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.
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/CHANGELOG.md +23 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/PKG-INFO +3 -2
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/pyproject.toml +1 -1
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/src/haystack_integrations/components/generators/cohere/chat/chat_generator.py +7 -3
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/src/haystack_integrations/components/generators/cohere/generator.py +7 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/src/haystack_integrations/components/rankers/cohere/ranker.py +69 -26
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/tests/test_chat_generator.py +54 -1
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/tests/test_ranker.py +78 -4
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/.gitignore +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/LICENSE.txt +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/README.md +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/examples/cohere_embedding.py +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/examples/cohere_generation.py +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/examples/cohere_ranker.py +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/pydoc/config_docusaurus.yml +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/src/haystack_integrations/components/embedders/cohere/__init__.py +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/src/haystack_integrations/components/embedders/cohere/document_embedder.py +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/src/haystack_integrations/components/embedders/cohere/document_image_embedder.py +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/src/haystack_integrations/components/embedders/cohere/embedding_types.py +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/src/haystack_integrations/components/embedders/cohere/text_embedder.py +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/src/haystack_integrations/components/embedders/cohere/utils.py +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/src/haystack_integrations/components/generators/cohere/__init__.py +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/src/haystack_integrations/components/generators/cohere/chat/__init__.py +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/src/haystack_integrations/components/generators/py.typed +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/src/haystack_integrations/components/rankers/cohere/__init__.py +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/tests/__init__.py +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/tests/conftest.py +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/tests/test_chat_generator_async.py +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/tests/test_chat_generator_chunks.py +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/tests/test_document_embedder.py +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/tests/test_document_image_embedder.py +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/tests/test_files/apple.jpg +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/tests/test_files/sample_pdf_1.pdf +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/tests/test_generator.py +0 -0
- {cohere_haystack-9.0.0 → cohere_haystack-9.2.0}/tests/test_text_embedder.py +0 -0
|
@@ -1,5 +1,28 @@
|
|
|
1
1
|
# Changelog
|
|
2
2
|
|
|
3
|
+
## [integrations/cohere-v9.1.0] - 2026-04-13
|
|
4
|
+
|
|
5
|
+
### 🚀 Features
|
|
6
|
+
|
|
7
|
+
- *(cohere)* Add async support to CohereRanker (#3143)
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
## [integrations/cohere-v9.0.0] - 2026-04-07
|
|
11
|
+
|
|
12
|
+
### 🐛 Bug Fixes
|
|
13
|
+
|
|
14
|
+
- [**breaking**] Cohere Embedders - set embed-v4.0 as the default model (#3111)
|
|
15
|
+
|
|
16
|
+
### 🧪 Testing
|
|
17
|
+
|
|
18
|
+
- Test compatible integrations with python 3.14; update pyproject (#3001)
|
|
19
|
+
- Track test coverage for all integrations (#3065)
|
|
20
|
+
|
|
21
|
+
### 🧹 Chores
|
|
22
|
+
|
|
23
|
+
- Enforce ruff docstring rules (D102/D103/D205/D209/D213/D417/D419) in first 10 integrations (#3008)
|
|
24
|
+
|
|
25
|
+
|
|
3
26
|
## [integrations/cohere-v8.1.0] - 2026-03-19
|
|
4
27
|
|
|
5
28
|
### 🚀 Features
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: cohere-haystack
|
|
3
|
-
Version: 9.
|
|
3
|
+
Version: 9.2.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
|
|
@@ -19,7 +19,8 @@ Classifier: Programming Language :: Python :: Implementation :: CPython
|
|
|
19
19
|
Classifier: Programming Language :: Python :: Implementation :: PyPy
|
|
20
20
|
Requires-Python: >=3.10
|
|
21
21
|
Requires-Dist: cohere>=5.17.0
|
|
22
|
-
Requires-Dist: haystack-ai>=2.
|
|
22
|
+
Requires-Dist: haystack-ai>=2.30.0
|
|
23
|
+
Requires-Dist: pydantic>=2.12.0
|
|
23
24
|
Description-Content-Type: text/markdown
|
|
24
25
|
|
|
25
26
|
# cohere-haystack
|
|
@@ -23,7 +23,7 @@ classifiers = [
|
|
|
23
23
|
"Programming Language :: Python :: Implementation :: CPython",
|
|
24
24
|
"Programming Language :: Python :: Implementation :: PyPy",
|
|
25
25
|
]
|
|
26
|
-
dependencies = ["haystack-ai>=2.
|
|
26
|
+
dependencies = ["haystack-ai>=2.30.0", "cohere>=5.17.0", "pydantic>=2.12.0"]
|
|
27
27
|
|
|
28
28
|
[project.urls]
|
|
29
29
|
Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/cohere#readme"
|
|
@@ -3,7 +3,7 @@ from collections.abc import AsyncIterator, Iterator
|
|
|
3
3
|
from typing import Any, ClassVar, Literal, get_args
|
|
4
4
|
|
|
5
5
|
from haystack import component, default_from_dict, default_to_dict, logging
|
|
6
|
-
from haystack.components.generators.utils import _convert_streaming_chunks_to_chat_message
|
|
6
|
+
from haystack.components.generators.utils import _convert_streaming_chunks_to_chat_message, _normalize_messages
|
|
7
7
|
from haystack.dataclasses import ChatMessage, ComponentInfo, ImageContent, TextContent, ToolCall
|
|
8
8
|
from haystack.dataclasses.streaming_chunk import (
|
|
9
9
|
AsyncStreamingCallbackT,
|
|
@@ -616,7 +616,7 @@ class CohereChatGenerator:
|
|
|
616
616
|
@component.output_types(replies=list[ChatMessage])
|
|
617
617
|
def run(
|
|
618
618
|
self,
|
|
619
|
-
messages: list[ChatMessage],
|
|
619
|
+
messages: list[ChatMessage] | str,
|
|
620
620
|
generation_kwargs: dict[str, Any] | None = None,
|
|
621
621
|
tools: ToolsType | None = None,
|
|
622
622
|
streaming_callback: StreamingCallbackT | None = None,
|
|
@@ -625,6 +625,7 @@ class CohereChatGenerator:
|
|
|
625
625
|
Invoke the chat endpoint based on the provided messages and generation parameters.
|
|
626
626
|
|
|
627
627
|
:param messages: list of `ChatMessage` instances representing the input messages.
|
|
628
|
+
If a string is provided, it is converted to a list containing a ChatMessage with user role.
|
|
628
629
|
:param generation_kwargs: additional keyword arguments for chat generation. These parameters will
|
|
629
630
|
potentially override the parameters passed in the __init__ method.
|
|
630
631
|
For more details on the parameters supported by the Cohere API, refer to the
|
|
@@ -637,6 +638,7 @@ class CohereChatGenerator:
|
|
|
637
638
|
:returns: A dictionary with the following keys:
|
|
638
639
|
- `replies`: a list of `ChatMessage` instances representing the generated responses.
|
|
639
640
|
"""
|
|
641
|
+
messages = _normalize_messages(messages)
|
|
640
642
|
|
|
641
643
|
# update generation kwargs by merging with the generation kwargs passed to the run method
|
|
642
644
|
generation_kwargs = {**self.generation_kwargs, **(generation_kwargs or {})}
|
|
@@ -682,7 +684,7 @@ class CohereChatGenerator:
|
|
|
682
684
|
@component.output_types(replies=list[ChatMessage])
|
|
683
685
|
async def run_async(
|
|
684
686
|
self,
|
|
685
|
-
messages: list[ChatMessage],
|
|
687
|
+
messages: list[ChatMessage] | str,
|
|
686
688
|
generation_kwargs: dict[str, Any] | None = None,
|
|
687
689
|
tools: ToolsType | None = None,
|
|
688
690
|
streaming_callback: StreamingCallbackT | None = None,
|
|
@@ -691,6 +693,7 @@ class CohereChatGenerator:
|
|
|
691
693
|
Asynchronously invoke the chat endpoint based on the provided messages and generation parameters.
|
|
692
694
|
|
|
693
695
|
:param messages: list of `ChatMessage` instances representing the input messages.
|
|
696
|
+
If a string is provided, it is converted to a list containing a ChatMessage with user role.
|
|
694
697
|
:param generation_kwargs: additional keyword arguments for chat generation. These parameters will
|
|
695
698
|
potentially override the parameters passed in the __init__ method.
|
|
696
699
|
For more details on the parameters supported by the Cohere API, refer to the
|
|
@@ -701,6 +704,7 @@ class CohereChatGenerator:
|
|
|
701
704
|
:returns: A dictionary with the following keys:
|
|
702
705
|
- `replies`: a list of `ChatMessage` instances representing the generated responses.
|
|
703
706
|
"""
|
|
707
|
+
messages = _normalize_messages(messages)
|
|
704
708
|
|
|
705
709
|
# update generation kwargs by merging with the generation kwargs passed to the run method
|
|
706
710
|
generation_kwargs = {**self.generation_kwargs, **(generation_kwargs or {})}
|
|
@@ -1,6 +1,7 @@
|
|
|
1
1
|
# SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
|
|
2
2
|
#
|
|
3
3
|
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
import warnings
|
|
4
5
|
from collections.abc import Callable
|
|
5
6
|
from typing import Any, ClassVar
|
|
6
7
|
|
|
@@ -70,6 +71,12 @@ class CohereGenerator(CohereChatGenerator):
|
|
|
70
71
|
You can check them in model's documentation.
|
|
71
72
|
"""
|
|
72
73
|
|
|
74
|
+
warnings.warn(
|
|
75
|
+
"The `CohereGenerator` component is deprecated and will be removed in a future version. "
|
|
76
|
+
"Use `CohereChatGenerator` instead, which now also supports string inputs.",
|
|
77
|
+
FutureWarning,
|
|
78
|
+
stacklevel=2,
|
|
79
|
+
)
|
|
73
80
|
# from_dict deserialization, where `generation_kwargs` is in **kwargs
|
|
74
81
|
if "generation_kwargs" in kwargs:
|
|
75
82
|
generation_kwargs = kwargs.pop("generation_kwargs")
|
|
@@ -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
|
|
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 =
|
|
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,14 +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
|
-
|
|
111
|
+
Validate parameters and prepare the input by concatenating the document text with the metadata fields.
|
|
109
112
|
|
|
110
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.
|
|
111
115
|
|
|
112
|
-
: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.
|
|
113
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
|
+
|
|
114
124
|
concatenated_input_list = []
|
|
115
125
|
for doc in documents:
|
|
116
126
|
meta_values_to_embed = [
|
|
@@ -119,7 +129,25 @@ class CohereRanker:
|
|
|
119
129
|
concatenated_input = self.meta_data_separator.join([*meta_values_to_embed, doc.content or ""])
|
|
120
130
|
concatenated_input_list.append(concatenated_input)
|
|
121
131
|
|
|
122
|
-
|
|
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}
|
|
123
151
|
|
|
124
152
|
@component.output_types(documents=list[Document])
|
|
125
153
|
def run(self, query: str, documents: list[Document], top_k: int | None = None) -> dict[str, list[Document]]:
|
|
@@ -138,19 +166,7 @@ class CohereRanker:
|
|
|
138
166
|
|
|
139
167
|
:raises ValueError: If `top_k` is not > 0.
|
|
140
168
|
"""
|
|
141
|
-
top_k =
|
|
142
|
-
if top_k <= 0:
|
|
143
|
-
msg = f"top_k must be > 0, but got {top_k}"
|
|
144
|
-
raise ValueError(msg)
|
|
145
|
-
|
|
146
|
-
cohere_input_docs = self._prepare_cohere_input_docs(documents)
|
|
147
|
-
if len(cohere_input_docs) > MAX_NUM_DOCS_FOR_COHERE_RANKER:
|
|
148
|
-
logger.warning(
|
|
149
|
-
f"The Cohere reranking endpoint only supports {MAX_NUM_DOCS_FOR_COHERE_RANKER} documents.\
|
|
150
|
-
The number of documents has been truncated to {MAX_NUM_DOCS_FOR_COHERE_RANKER} \
|
|
151
|
-
from {len(cohere_input_docs)}."
|
|
152
|
-
)
|
|
153
|
-
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)
|
|
154
170
|
|
|
155
171
|
response = self._cohere_client.rerank(
|
|
156
172
|
model=self.model_name,
|
|
@@ -159,10 +175,37 @@ class CohereRanker:
|
|
|
159
175
|
max_tokens_per_doc=self.max_tokens_per_doc,
|
|
160
176
|
top_n=top_k,
|
|
161
177
|
)
|
|
162
|
-
|
|
163
|
-
|
|
164
|
-
|
|
165
|
-
|
|
166
|
-
|
|
167
|
-
|
|
168
|
-
|
|
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)
|
|
@@ -1,5 +1,5 @@
|
|
|
1
1
|
import os
|
|
2
|
-
from unittest.mock import MagicMock
|
|
2
|
+
from unittest.mock import AsyncMock, MagicMock
|
|
3
3
|
|
|
4
4
|
import pytest
|
|
5
5
|
from cohere.core import ApiError
|
|
@@ -498,6 +498,59 @@ class TestCohereChatGenerator:
|
|
|
498
498
|
assert multimodal_msg["content"][0]["type"] == "text"
|
|
499
499
|
assert multimodal_msg["content"][1]["type"] == "image_url"
|
|
500
500
|
|
|
501
|
+
def test_run_with_string_input(self):
|
|
502
|
+
"""Test that a string input is converted to a user ChatMessage before sending to the API."""
|
|
503
|
+
generator = CohereChatGenerator(api_key=Secret.from_token("test-api-key"))
|
|
504
|
+
|
|
505
|
+
mock_response = MagicMock()
|
|
506
|
+
mock_response.message.content = [MagicMock()]
|
|
507
|
+
mock_response.message.content[0].text = "Paris"
|
|
508
|
+
mock_response.message.content[0].type = "text"
|
|
509
|
+
mock_response.message.tool_calls = None
|
|
510
|
+
mock_response.finish_reason = "COMPLETE"
|
|
511
|
+
mock_response.usage = None
|
|
512
|
+
|
|
513
|
+
generator.client.chat = MagicMock(return_value=mock_response)
|
|
514
|
+
|
|
515
|
+
result = generator.run("What's the capital of France?")
|
|
516
|
+
|
|
517
|
+
generator.client.chat.assert_called_once()
|
|
518
|
+
call_args = generator.client.chat.call_args
|
|
519
|
+
sent_messages = call_args[1]["messages"]
|
|
520
|
+
|
|
521
|
+
assert len(sent_messages) == 1
|
|
522
|
+
assert sent_messages[0]["role"] == "user"
|
|
523
|
+
assert isinstance(result["replies"], list)
|
|
524
|
+
assert len(result["replies"]) == 1
|
|
525
|
+
assert isinstance(result["replies"][0], ChatMessage)
|
|
526
|
+
|
|
527
|
+
@pytest.mark.asyncio
|
|
528
|
+
async def test_run_async_with_string_input(self):
|
|
529
|
+
"""Test that a string input is converted to a user ChatMessage before sending to the async API."""
|
|
530
|
+
generator = CohereChatGenerator(api_key=Secret.from_token("test-api-key"))
|
|
531
|
+
|
|
532
|
+
mock_response = MagicMock()
|
|
533
|
+
mock_response.message.content = [MagicMock()]
|
|
534
|
+
mock_response.message.content[0].text = "Paris"
|
|
535
|
+
mock_response.message.content[0].type = "text"
|
|
536
|
+
mock_response.message.tool_calls = None
|
|
537
|
+
mock_response.finish_reason = "COMPLETE"
|
|
538
|
+
mock_response.usage = None
|
|
539
|
+
|
|
540
|
+
generator.async_client.chat = AsyncMock(return_value=mock_response)
|
|
541
|
+
|
|
542
|
+
result = await generator.run_async("What's the capital of France?")
|
|
543
|
+
|
|
544
|
+
generator.async_client.chat.assert_called_once()
|
|
545
|
+
call_args = generator.async_client.chat.call_args
|
|
546
|
+
sent_messages = call_args[1]["messages"]
|
|
547
|
+
|
|
548
|
+
assert len(sent_messages) == 1
|
|
549
|
+
assert sent_messages[0]["role"] == "user"
|
|
550
|
+
assert isinstance(result["replies"], list)
|
|
551
|
+
assert len(result["replies"]) == 1
|
|
552
|
+
assert isinstance(result["replies"][0], ChatMessage)
|
|
553
|
+
|
|
501
554
|
|
|
502
555
|
@pytest.mark.skipif(
|
|
503
556
|
not os.environ.get("COHERE_API_KEY", None) and not os.environ.get("CO_API_KEY", None),
|
|
@@ -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)
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|