cohere-haystack 10.0.0__tar.gz → 10.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-10.0.0 → cohere_haystack-10.1.0}/CHANGELOG.md +23 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/PKG-INFO +1 -1
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/pyproject.toml +2 -1
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/embedders/cohere/document_embedder.py +7 -1
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/embedders/cohere/document_image_embedder.py +5 -1
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/embedders/cohere/text_embedder.py +5 -1
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/embedders/cohere/utils.py +26 -17
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/generators/cohere/chat/chat_generator.py +4 -1
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/rankers/cohere/ranker.py +6 -2
- cohere_haystack-10.1.0/src/haystack_integrations/utils/cohere/__init__.py +6 -0
- cohere_haystack-10.1.0/src/haystack_integrations/utils/cohere/api_base_url.py +29 -0
- cohere_haystack-10.1.0/src/haystack_integrations/utils/py.typed +0 -0
- cohere_haystack-10.1.0/tests/test_api_base_url.py +75 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/test_document_embedder.py +29 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/.gitignore +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/LICENSE.txt +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/README.md +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/examples/cohere_embedding.py +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/examples/cohere_generation.py +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/examples/cohere_ranker.py +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/pydoc/config_docusaurus.yml +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/embedders/cohere/__init__.py +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/embedders/cohere/embedding_types.py +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/generators/cohere/__init__.py +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/generators/cohere/chat/__init__.py +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/generators/py.typed +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/rankers/cohere/__init__.py +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/__init__.py +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/conftest.py +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/test_chat_generator.py +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/test_chat_generator_async.py +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/test_chat_generator_chunks.py +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/test_document_image_embedder.py +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/test_files/apple.jpg +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/test_files/sample_pdf_1.pdf +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/test_ranker.py +0 -0
- {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/test_text_embedder.py +0 -0
|
@@ -1,5 +1,28 @@
|
|
|
1
1
|
# Changelog
|
|
2
2
|
|
|
3
|
+
## [integrations/cohere-v10.0.1] - 2026-09-10
|
|
4
|
+
|
|
5
|
+
### 🐛 Bug Fixes
|
|
6
|
+
|
|
7
|
+
- Batch CohereDocumentEmbedder.run_async like the sync path (#3926)
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
## [integrations/cohere-v10.0.0] - 2026-09-04
|
|
11
|
+
|
|
12
|
+
### 🚜 Refactor
|
|
13
|
+
|
|
14
|
+
- [**breaking**] Cohere - create clients in warm up methods (#3911)
|
|
15
|
+
|
|
16
|
+
### ⚙️ CI
|
|
17
|
+
|
|
18
|
+
- Improve changelog generation; fix existing changelogs (#3883)
|
|
19
|
+
|
|
20
|
+
### 🧹 Chores
|
|
21
|
+
|
|
22
|
+
- Clarify how generation_kwargs passed in run are handled (#3805)
|
|
23
|
+
- [**breaking**] Remove deprecated `CohereGenerator` (use `CohereChatGenerator` instead) (#3885)
|
|
24
|
+
|
|
25
|
+
|
|
3
26
|
## [integrations/cohere-v9.5.0] - 2026-08-11
|
|
4
27
|
|
|
5
28
|
### 🐛 Bug Fixes
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: cohere-haystack
|
|
3
|
-
Version: 10.
|
|
3
|
+
Version: 10.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
|
|
@@ -72,7 +72,8 @@ integration-cov-append-retry = 'pytest --cov=haystack_integrations --cov-append
|
|
|
72
72
|
|
|
73
73
|
types = """mypy -p haystack_integrations.components.embedders.cohere \
|
|
74
74
|
-p haystack_integrations.components.generators.cohere \
|
|
75
|
-
-p haystack_integrations.components.rankers.cohere
|
|
75
|
+
-p haystack_integrations.components.rankers.cohere \
|
|
76
|
+
-p haystack_integrations.utils.cohere {args}"""
|
|
76
77
|
|
|
77
78
|
[tool.mypy]
|
|
78
79
|
install_types = true
|
|
@@ -8,6 +8,7 @@ from haystack import Document, component, default_from_dict, default_to_dict
|
|
|
8
8
|
from haystack.utils import Secret, deserialize_secrets_inplace
|
|
9
9
|
|
|
10
10
|
from cohere import AsyncClientV2, ClientV2
|
|
11
|
+
from haystack_integrations.utils.cohere import validate_api_base_url
|
|
11
12
|
|
|
12
13
|
from .embedding_types import EmbeddingTypes
|
|
13
14
|
from .utils import get_async_response, get_response
|
|
@@ -68,7 +69,7 @@ class CohereDocumentEmbedder:
|
|
|
68
69
|
Read [Cohere documentation](https://docs.cohere.com/docs/models#embed) for a list of all supported models.
|
|
69
70
|
:param input_type: specifies the type of input you're giving to the model. Supported values are
|
|
70
71
|
"search_document", "search_query", "classification" and "clustering".
|
|
71
|
-
:param api_base_url: the Cohere API Base url.
|
|
72
|
+
:param api_base_url: the Cohere API Base url. The Cohere client appends the endpoint path to it.
|
|
72
73
|
:param truncate: truncate embeddings that are too long from start or end, ("NONE"|"START"|"END").
|
|
73
74
|
Passing "START" will discard the start of the input. "END" will discard the end of the input. In both
|
|
74
75
|
cases, input is discarded until the remaining input is exactly the maximum input token length for the model.
|
|
@@ -81,7 +82,10 @@ class CohereDocumentEmbedder:
|
|
|
81
82
|
:param embedding_separator: separator used to concatenate the meta fields to the Document text.
|
|
82
83
|
:param embedding_type: the type of embeddings to return. Defaults to float embeddings.
|
|
83
84
|
Note that int8, uint8, binary, and ubinary are only valid for v3 models.
|
|
85
|
+
|
|
86
|
+
:raises ValueError: If `api_base_url` is a full endpoint URL rather than a base URL.
|
|
84
87
|
"""
|
|
88
|
+
validate_api_base_url(api_base_url)
|
|
85
89
|
|
|
86
90
|
self.api_key = api_key
|
|
87
91
|
self.model = model
|
|
@@ -249,6 +253,8 @@ class CohereDocumentEmbedder:
|
|
|
249
253
|
model_name=self.model,
|
|
250
254
|
input_type=self.input_type,
|
|
251
255
|
truncate=self.truncate,
|
|
256
|
+
batch_size=self.batch_size,
|
|
257
|
+
progress_bar=self.progress_bar,
|
|
252
258
|
embedding_type=self.embedding_type,
|
|
253
259
|
)
|
|
254
260
|
|
|
@@ -17,6 +17,7 @@ from haystack.utils.auth import Secret, deserialize_secrets_inplace
|
|
|
17
17
|
from tqdm import tqdm
|
|
18
18
|
|
|
19
19
|
from cohere import AsyncClientV2, ClientV2
|
|
20
|
+
from haystack_integrations.utils.cohere import validate_api_base_url
|
|
20
21
|
|
|
21
22
|
from .embedding_types import EmbeddingTypes
|
|
22
23
|
|
|
@@ -101,7 +102,7 @@ class CohereDocumentImageEmbedder:
|
|
|
101
102
|
The Cohere model to use for calculating embeddings.
|
|
102
103
|
Read [Cohere documentation](https://docs.cohere.com/docs/models#embed) for a list of all supported models.
|
|
103
104
|
:param api_base_url:
|
|
104
|
-
The Cohere API base URL.
|
|
105
|
+
The Cohere API base URL. The Cohere client appends the endpoint path to it.
|
|
105
106
|
:param timeout:
|
|
106
107
|
Request timeout in seconds.
|
|
107
108
|
:param embedding_dimension:
|
|
@@ -114,7 +115,10 @@ class CohereDocumentImageEmbedder:
|
|
|
114
115
|
:param progress_bar:
|
|
115
116
|
Whether to show a progress bar or not. Can be helpful to disable in production deployments
|
|
116
117
|
to keep the logs clean.
|
|
118
|
+
|
|
119
|
+
:raises ValueError: If `api_base_url` is a full endpoint URL rather than a base URL.
|
|
117
120
|
"""
|
|
121
|
+
validate_api_base_url(api_base_url)
|
|
118
122
|
|
|
119
123
|
self.file_path_meta_field = file_path_meta_field
|
|
120
124
|
self.root_path = root_path or ""
|
|
@@ -7,6 +7,7 @@ from haystack import component, default_from_dict, default_to_dict
|
|
|
7
7
|
from haystack.utils import Secret, deserialize_secrets_inplace
|
|
8
8
|
|
|
9
9
|
from cohere import AsyncClientV2, ClientV2
|
|
10
|
+
from haystack_integrations.utils.cohere import validate_api_base_url
|
|
10
11
|
|
|
11
12
|
from .embedding_types import EmbeddingTypes
|
|
12
13
|
from .utils import get_async_response, get_response
|
|
@@ -60,7 +61,7 @@ class CohereTextEmbedder:
|
|
|
60
61
|
Read [Cohere documentation](https://docs.cohere.com/docs/models#embed) for a list of all supported models.
|
|
61
62
|
:param input_type: specifies the type of input you're giving to the model. Supported values are
|
|
62
63
|
"search_document", "search_query", "classification" and "clustering".
|
|
63
|
-
:param api_base_url: the Cohere API Base url.
|
|
64
|
+
:param api_base_url: the Cohere API Base url. The Cohere client appends the endpoint path to it.
|
|
64
65
|
:param truncate: truncate embeddings that are too long from start or end, ("NONE"|"START"|"END").
|
|
65
66
|
Passing "START" will discard the start of the input. "END" will discard the end of the input. In both
|
|
66
67
|
cases, input is discarded until the remaining input is exactly the maximum input token length for the model.
|
|
@@ -68,7 +69,10 @@ class CohereTextEmbedder:
|
|
|
68
69
|
:param timeout: request timeout in seconds.
|
|
69
70
|
:param embedding_type: the type of embeddings to return. Defaults to float embeddings.
|
|
70
71
|
Note that int8, uint8, binary, and ubinary are only valid for v3 models.
|
|
72
|
+
|
|
73
|
+
:raises ValueError: If `api_base_url` is a full endpoint URL rather than a base URL.
|
|
71
74
|
"""
|
|
75
|
+
validate_api_base_url(api_base_url)
|
|
72
76
|
|
|
73
77
|
self.api_key = api_key
|
|
74
78
|
self.model = model
|
|
@@ -16,6 +16,8 @@ async def get_async_response(
|
|
|
16
16
|
model_name: str,
|
|
17
17
|
input_type: str,
|
|
18
18
|
truncate: str,
|
|
19
|
+
batch_size: int = 32,
|
|
20
|
+
progress_bar: bool = False,
|
|
19
21
|
embedding_type: EmbeddingTypes | None = None,
|
|
20
22
|
) -> tuple[list[list[float]], dict[str, Any]]:
|
|
21
23
|
"""
|
|
@@ -27,6 +29,9 @@ async def get_async_response(
|
|
|
27
29
|
:param input_type: one of "classification", "clustering", "search_document", "search_query".
|
|
28
30
|
The type of input text provided to embed.
|
|
29
31
|
:param truncate: one of "NONE", "START", "END". How the API handles text longer than the maximum token length.
|
|
32
|
+
:param batch_size: the batch size to use. The Cohere embed endpoint caps the number of texts per call, so the
|
|
33
|
+
texts are sent in batches just like the synchronous path.
|
|
34
|
+
:param progress_bar: if `True`, show a progress bar
|
|
30
35
|
:param embedding_type: the type of embeddings to return. Defaults to float embeddings.
|
|
31
36
|
|
|
32
37
|
:returns: A tuple of the embeddings and metadata.
|
|
@@ -35,24 +40,28 @@ async def get_async_response(
|
|
|
35
40
|
"""
|
|
36
41
|
all_embeddings: list[list[float]] = []
|
|
37
42
|
metadata: dict[str, Any] = {}
|
|
38
|
-
|
|
39
43
|
embedding_type = embedding_type or EmbeddingTypes.FLOAT
|
|
40
|
-
|
|
41
|
-
|
|
42
|
-
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
|
|
46
|
-
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
|
|
55
|
-
|
|
44
|
+
|
|
45
|
+
for i in tqdm(
|
|
46
|
+
range(0, len(texts), batch_size),
|
|
47
|
+
disable=not progress_bar,
|
|
48
|
+
desc="Calculating embeddings",
|
|
49
|
+
):
|
|
50
|
+
batch = texts[i : i + batch_size]
|
|
51
|
+
response = await cohere_async_client.embed(
|
|
52
|
+
texts=batch,
|
|
53
|
+
model=model_name,
|
|
54
|
+
input_type=input_type,
|
|
55
|
+
truncate=truncate,
|
|
56
|
+
embedding_types=[embedding_type.value],
|
|
57
|
+
)
|
|
58
|
+
for emb_tuple in response.embeddings:
|
|
59
|
+
# emb_tuple[0] is a str denoting the embedding type (e.g. "float", "int8", etc.)
|
|
60
|
+
if emb_tuple[1] is not None:
|
|
61
|
+
# ok we have embeddings for this type, let's take all the embeddings (a list of embeddings)
|
|
62
|
+
all_embeddings.extend(emb_tuple[1])
|
|
63
|
+
if response.meta is not None:
|
|
64
|
+
metadata = response.meta.model_dump()
|
|
56
65
|
|
|
57
66
|
return all_embeddings, metadata
|
|
58
67
|
|
|
@@ -29,6 +29,7 @@ from httpx import AsyncHTTPTransport, HTTPTransport
|
|
|
29
29
|
from httpx import Client as HTTPXClient
|
|
30
30
|
|
|
31
31
|
from cohere import AsyncClientV2, ChatResponse, ClientV2, StreamedChatResponseV2
|
|
32
|
+
from haystack_integrations.utils.cohere import validate_api_base_url
|
|
32
33
|
|
|
33
34
|
logger = logging.getLogger(__name__)
|
|
34
35
|
|
|
@@ -544,7 +545,7 @@ class CohereChatGenerator:
|
|
|
544
545
|
:param streaming_callback: A callback function that is called when a new token is received from the stream.
|
|
545
546
|
The callback function accepts [StreamingChunk](https://docs.haystack.deepset.ai/docs/data-classes#streamingchunk)
|
|
546
547
|
as an argument.
|
|
547
|
-
:param api_base_url: The base URL of the Cohere API.
|
|
548
|
+
:param api_base_url: The base URL of the Cohere API. The Cohere client appends the endpoint path to it.
|
|
548
549
|
:param generation_kwargs: Other parameters to use for the model during generation. For a list of parameters,
|
|
549
550
|
see [Cohere Chat endpoint](https://docs.cohere.com/reference/chat).
|
|
550
551
|
Some of the parameters are:
|
|
@@ -564,11 +565,13 @@ class CohereChatGenerator:
|
|
|
564
565
|
Maximum number of retries to attempt for failed requests. If not set, it defaults to the default set by
|
|
565
566
|
the Cohere client.
|
|
566
567
|
|
|
568
|
+
:raises ValueError: If `api_base_url` is a full endpoint URL rather than a base URL.
|
|
567
569
|
"""
|
|
568
570
|
_check_duplicate_tool_names(flatten_tools_or_toolsets(tools))
|
|
569
571
|
|
|
570
572
|
if not api_base_url:
|
|
571
573
|
api_base_url = "https://api.cohere.com"
|
|
574
|
+
validate_api_base_url(api_base_url)
|
|
572
575
|
|
|
573
576
|
self.api_key = api_key
|
|
574
577
|
self.model = model
|
|
@@ -5,6 +5,7 @@ from haystack import Document, component, default_from_dict, default_to_dict, lo
|
|
|
5
5
|
from haystack.utils import Secret, deserialize_secrets_inplace
|
|
6
6
|
|
|
7
7
|
from cohere import AsyncClientV2, ClientV2
|
|
8
|
+
from haystack_integrations.utils.cohere import validate_api_base_url
|
|
8
9
|
|
|
9
10
|
logger = logging.getLogger(__name__)
|
|
10
11
|
|
|
@@ -48,19 +49,22 @@ class CohereRanker:
|
|
|
48
49
|
:param model: Cohere model name. Check the list of supported models in the [Cohere documentation](https://docs.cohere.com/docs/models).
|
|
49
50
|
:param top_k: The maximum number of documents to return.
|
|
50
51
|
:param api_key: Cohere API key.
|
|
51
|
-
:param api_base_url: the base URL of the Cohere API.
|
|
52
|
+
:param api_base_url: the base URL of the Cohere API. The Cohere client appends the endpoint path to it.
|
|
52
53
|
:param meta_fields_to_embed: List of meta fields that should be concatenated
|
|
53
54
|
with the document content for reranking.
|
|
54
55
|
:param meta_data_separator: Separator used to concatenate the meta fields
|
|
55
56
|
to the Document content.
|
|
56
57
|
:param max_tokens_per_doc: The maximum number of tokens to embed for each document defaults to 4096.
|
|
57
58
|
|
|
58
|
-
:raises ValueError: If `top_k` is not > 0
|
|
59
|
+
:raises ValueError: If `top_k` is not > 0, or if `api_base_url` is a full endpoint URL rather than a
|
|
60
|
+
base URL.
|
|
59
61
|
"""
|
|
60
62
|
if top_k <= 0:
|
|
61
63
|
msg = f"top_k must be > 0, but got {top_k}"
|
|
62
64
|
raise ValueError(msg)
|
|
63
65
|
|
|
66
|
+
validate_api_base_url(api_base_url)
|
|
67
|
+
|
|
64
68
|
self.model_name = model
|
|
65
69
|
self.api_key = api_key
|
|
66
70
|
self.api_base_url = api_base_url
|
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
|
|
2
|
+
#
|
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
import re
|
|
5
|
+
|
|
6
|
+
# Cohere clients build a request URL by appending the endpoint path to the configured base URL, so a base URL that
|
|
7
|
+
# already ends with one of these paths produces a doubled path and a generic API error at request time.
|
|
8
|
+
_COHERE_ENDPOINT_PATH = re.compile(r"/v\d+/(?:rerank|embed|chat)$")
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def validate_api_base_url(api_base_url: str) -> None:
|
|
12
|
+
"""
|
|
13
|
+
Check that `api_base_url` is a base URL rather than a full endpoint URL.
|
|
14
|
+
|
|
15
|
+
:param api_base_url: The base URL to validate.
|
|
16
|
+
:raises ValueError: If `api_base_url` ends with a Cohere endpoint path, which the Cohere client appends itself.
|
|
17
|
+
"""
|
|
18
|
+
trimmed = api_base_url.rstrip("/")
|
|
19
|
+
match = _COHERE_ENDPOINT_PATH.search(trimmed)
|
|
20
|
+
if match is None:
|
|
21
|
+
return
|
|
22
|
+
|
|
23
|
+
base = trimmed[: match.start()]
|
|
24
|
+
hint = f" and pass '{base}' instead" if base else ""
|
|
25
|
+
msg = (
|
|
26
|
+
f"'api_base_url' must be a base URL, not a full endpoint URL, but got {api_base_url!r}. "
|
|
27
|
+
f"The Cohere client appends the endpoint path itself, so remove the trailing '{match.group()}'{hint}."
|
|
28
|
+
)
|
|
29
|
+
raise ValueError(msg)
|
|
File without changes
|
|
@@ -0,0 +1,75 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
|
|
2
|
+
#
|
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
import pytest
|
|
5
|
+
|
|
6
|
+
from haystack_integrations.components.embedders.cohere import (
|
|
7
|
+
CohereDocumentEmbedder,
|
|
8
|
+
CohereDocumentImageEmbedder,
|
|
9
|
+
CohereTextEmbedder,
|
|
10
|
+
)
|
|
11
|
+
from haystack_integrations.components.generators.cohere import CohereChatGenerator
|
|
12
|
+
from haystack_integrations.components.rankers.cohere import CohereRanker
|
|
13
|
+
from haystack_integrations.utils.cohere import validate_api_base_url
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
class TestValidateApiBaseUrl:
|
|
17
|
+
@pytest.mark.parametrize(
|
|
18
|
+
"api_base_url",
|
|
19
|
+
[
|
|
20
|
+
"https://api.cohere.com",
|
|
21
|
+
"https://api.cohere.com/",
|
|
22
|
+
"https://resource.services.ai.azure.com/models",
|
|
23
|
+
# only a trailing endpoint path is rejected, not one in the middle of the base path
|
|
24
|
+
"https://gateway.example.com/v2/rerank/proxy",
|
|
25
|
+
],
|
|
26
|
+
)
|
|
27
|
+
def test_accepts_base_urls(self, api_base_url):
|
|
28
|
+
validate_api_base_url(api_base_url)
|
|
29
|
+
|
|
30
|
+
@pytest.mark.parametrize(
|
|
31
|
+
"api_base_url",
|
|
32
|
+
[
|
|
33
|
+
"https://api.cohere.com/v1/rerank",
|
|
34
|
+
"https://api.cohere.com/v2/embed",
|
|
35
|
+
"https://api.cohere.com/v2/chat",
|
|
36
|
+
# a future API version is rejected without needing a new entry in the pattern
|
|
37
|
+
"https://api.cohere.com/v3/rerank",
|
|
38
|
+
],
|
|
39
|
+
)
|
|
40
|
+
def test_rejects_endpoint_urls(self, api_base_url):
|
|
41
|
+
with pytest.raises(ValueError, match="must be a base URL"):
|
|
42
|
+
validate_api_base_url(api_base_url)
|
|
43
|
+
|
|
44
|
+
@pytest.mark.parametrize("suffix", ["", "/", "//"])
|
|
45
|
+
def test_error_names_the_endpoint_path_and_the_corrected_url(self, suffix):
|
|
46
|
+
url = f"https://resource.services.ai.azure.com/models/v2/rerank{suffix}"
|
|
47
|
+
with pytest.raises(ValueError) as exc_info:
|
|
48
|
+
validate_api_base_url(url)
|
|
49
|
+
|
|
50
|
+
message = str(exc_info.value)
|
|
51
|
+
assert "remove the trailing '/v2/rerank'" in message
|
|
52
|
+
assert "pass 'https://resource.services.ai.azure.com/models' instead" in message
|
|
53
|
+
|
|
54
|
+
def test_omits_the_hint_when_nothing_would_remain(self):
|
|
55
|
+
with pytest.raises(ValueError) as exc_info:
|
|
56
|
+
validate_api_base_url("/v2/rerank")
|
|
57
|
+
|
|
58
|
+
assert "instead" not in str(exc_info.value)
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
@pytest.mark.parametrize(
|
|
62
|
+
"component_cls",
|
|
63
|
+
[
|
|
64
|
+
CohereRanker,
|
|
65
|
+
CohereTextEmbedder,
|
|
66
|
+
CohereDocumentEmbedder,
|
|
67
|
+
CohereDocumentImageEmbedder,
|
|
68
|
+
CohereChatGenerator,
|
|
69
|
+
],
|
|
70
|
+
)
|
|
71
|
+
def test_components_reject_full_endpoint_urls(component_cls, monkeypatch):
|
|
72
|
+
monkeypatch.setenv("COHERE_API_KEY", "test-api-key")
|
|
73
|
+
|
|
74
|
+
with pytest.raises(ValueError, match="must be a base URL"):
|
|
75
|
+
component_cls(api_base_url="https://resource.services.ai.azure.com/models/v2/rerank")
|
|
@@ -247,6 +247,35 @@ class TestRun:
|
|
|
247
247
|
assert doc_with_embedding.meta == doc.meta
|
|
248
248
|
assert doc_with_embedding.embedding == embedding
|
|
249
249
|
|
|
250
|
+
@pytest.mark.asyncio
|
|
251
|
+
@patch("haystack_integrations.components.embedders.cohere.document_embedder.AsyncClientV2")
|
|
252
|
+
async def test_run_async_sends_texts_in_batches(self, mock_async_client_cls):
|
|
253
|
+
"""
|
|
254
|
+
run_async must honour batch_size like run does. The Cohere embed endpoint caps the number of
|
|
255
|
+
texts per call, so sending every text in a single request breaks on larger inputs.
|
|
256
|
+
"""
|
|
257
|
+
batch_sizes: list[int] = []
|
|
258
|
+
|
|
259
|
+
class _Response:
|
|
260
|
+
def __init__(self, n: int) -> None:
|
|
261
|
+
self.embeddings = [("float", [[0.1, 0.2, 0.3]] * n)]
|
|
262
|
+
self.meta = None
|
|
263
|
+
|
|
264
|
+
async def fake_embed(*, texts, **_kwargs):
|
|
265
|
+
batch_sizes.append(len(texts))
|
|
266
|
+
return _Response(len(texts))
|
|
267
|
+
|
|
268
|
+
mock_async_client_cls.return_value.embed = fake_embed
|
|
269
|
+
|
|
270
|
+
embedder = CohereDocumentEmbedder(api_key=Secret.from_token("test-api-key"), batch_size=2)
|
|
271
|
+
docs = [Document(content=f"doc {i}") for i in range(5)]
|
|
272
|
+
|
|
273
|
+
result = await embedder.run_async(docs)
|
|
274
|
+
|
|
275
|
+
assert batch_sizes == [2, 2, 1]
|
|
276
|
+
assert len(result["documents"]) == 5
|
|
277
|
+
assert all(doc.embedding == [0.1, 0.2, 0.3] for doc in result["documents"])
|
|
278
|
+
|
|
250
279
|
@patch("haystack_integrations.components.embedders.cohere.document_embedder.get_response")
|
|
251
280
|
def test_run_does_not_modify_original_documents(self, mock_get_response):
|
|
252
281
|
embedder = CohereDocumentEmbedder(api_key=Secret.from_token("test-api-key"))
|
|
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
|