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.
Files changed (39) hide show
  1. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/CHANGELOG.md +23 -0
  2. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/PKG-INFO +1 -1
  3. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/pyproject.toml +2 -1
  4. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/embedders/cohere/document_embedder.py +7 -1
  5. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/embedders/cohere/document_image_embedder.py +5 -1
  6. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/embedders/cohere/text_embedder.py +5 -1
  7. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/embedders/cohere/utils.py +26 -17
  8. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/generators/cohere/chat/chat_generator.py +4 -1
  9. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/rankers/cohere/ranker.py +6 -2
  10. cohere_haystack-10.1.0/src/haystack_integrations/utils/cohere/__init__.py +6 -0
  11. cohere_haystack-10.1.0/src/haystack_integrations/utils/cohere/api_base_url.py +29 -0
  12. cohere_haystack-10.1.0/src/haystack_integrations/utils/py.typed +0 -0
  13. cohere_haystack-10.1.0/tests/test_api_base_url.py +75 -0
  14. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/test_document_embedder.py +29 -0
  15. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/.gitignore +0 -0
  16. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/LICENSE.txt +0 -0
  17. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/README.md +0 -0
  18. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/examples/cohere_embedding.py +0 -0
  19. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/examples/cohere_generation.py +0 -0
  20. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/examples/cohere_ranker.py +0 -0
  21. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/pydoc/config_docusaurus.yml +0 -0
  22. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/embedders/cohere/__init__.py +0 -0
  23. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/embedders/cohere/embedding_types.py +0 -0
  24. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
  25. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/generators/cohere/__init__.py +0 -0
  26. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/generators/cohere/chat/__init__.py +0 -0
  27. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/generators/py.typed +0 -0
  28. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/rankers/cohere/__init__.py +0 -0
  29. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
  30. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/__init__.py +0 -0
  31. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/conftest.py +0 -0
  32. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/test_chat_generator.py +0 -0
  33. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/test_chat_generator_async.py +0 -0
  34. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/test_chat_generator_chunks.py +0 -0
  35. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/test_document_image_embedder.py +0 -0
  36. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/test_files/apple.jpg +0 -0
  37. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/test_files/sample_pdf_1.pdf +0 -0
  38. {cohere_haystack-10.0.0 → cohere_haystack-10.1.0}/tests/test_ranker.py +0 -0
  39. {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.0.0
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 {args}"""
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
- response = await cohere_async_client.embed(
41
- texts=texts,
42
- model=model_name,
43
- input_type=input_type,
44
- truncate=truncate,
45
- embedding_types=[embedding_type.value],
46
- )
47
- if response.meta is not None:
48
- metadata = response.meta.model_dump()
49
- for emb_tuple in response.embeddings:
50
- # emb_tuple[0] is a str denoting the embedding type (e.g. "float", "int8", etc.)
51
- if emb_tuple[1] is not None:
52
- # ok we have embeddings for this type, let's take all
53
- # the embeddings (a list of embeddings) and break the loop
54
- all_embeddings.extend(emb_tuple[1])
55
- break
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,6 @@
1
+ # SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+ from .api_base_url import validate_api_base_url
5
+
6
+ __all__ = ["validate_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)
@@ -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"))