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.
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/CHANGELOG.md +27 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/PKG-INFO +2 -1
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/pyproject.toml +12 -2
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/embedders/cohere/document_embedder.py +8 -10
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/embedders/cohere/text_embedder.py +6 -9
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/embedders/cohere/utils.py +4 -2
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/generators/cohere/generator.py +3 -1
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/rankers/cohere/ranker.py +70 -26
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/test_document_embedder.py +6 -6
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/test_ranker.py +78 -4
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/test_text_embedder.py +6 -6
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/.gitignore +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/LICENSE.txt +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/README.md +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/examples/cohere_embedding.py +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/examples/cohere_generation.py +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/examples/cohere_ranker.py +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/pydoc/config_docusaurus.yml +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/embedders/cohere/__init__.py +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/embedders/cohere/document_image_embedder.py +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/embedders/cohere/embedding_types.py +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/generators/cohere/__init__.py +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/generators/cohere/chat/__init__.py +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/generators/cohere/chat/chat_generator.py +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/generators/py.typed +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/rankers/cohere/__init__.py +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/__init__.py +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/conftest.py +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/test_chat_generator.py +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/test_chat_generator_async.py +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/test_chat_generator_chunks.py +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/test_document_image_embedder.py +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/test_files/apple.jpg +0 -0
- {cohere_haystack-8.1.0 → cohere_haystack-9.1.0}/tests/test_files/sample_pdf_1.pdf +0 -0
- {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:
|
|
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-
|
|
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.
|
|
66
|
-
|
|
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".
|
|
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
|
-
"""
|
|
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-
|
|
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.
|
|
58
|
-
|
|
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".
|
|
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
|
-
"""
|
|
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
|
-
"""
|
|
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
|
-
"""
|
|
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
|
|
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,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
|
-
|
|
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
|
-
|
|
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 =
|
|
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
|
-
|
|
162
|
-
|
|
163
|
-
|
|
164
|
-
|
|
165
|
-
|
|
166
|
-
|
|
167
|
-
|
|
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-
|
|
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-
|
|
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-
|
|
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-
|
|
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-
|
|
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-
|
|
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-
|
|
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-
|
|
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-
|
|
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-
|
|
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"]) ==
|
|
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"]) ==
|
|
165
|
+
assert len(result["embedding"]) == 1536
|
|
166
166
|
assert all(isinstance(x, float) for x in result["embedding"])
|
|
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
|