From e542ef0363da3ec702f8616a583d9154ae35ab42 Mon Sep 17 00:00:00 2001 From: fern-support <126544928+fern-support@users.noreply.github.com> Date: Wed, 11 Feb 2026 12:28:26 -0500 Subject: [PATCH 1/8] fix: resolve AWS client SigV4 signing, forced SageMaker dep, and missing embed params - Fix SigV4 host header mismatch: update copied headers dict with correct host after URL rewrite, so AWSRequest signs with the Bedrock/SageMaker host instead of stale api.cohere.com - Add mode parameter to cohere_aws.Client to conditionally initialize boto3 clients (bedrock-runtime/bedrock vs sagemaker-runtime/sagemaker), avoiding forced SageMaker dependency for Bedrock users - Add output_dimension and embedding_types params to embed() for Embed v4 Closes #721 Co-Authored-By: Claude Opus 4.6 --- src/cohere/aws_client.py | 1 + .../manually_maintained/cohere_aws/client.py | 20 ++++++++++++++----- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/src/cohere/aws_client.py b/src/cohere/aws_client.py index 8aea9d15c..12a168276 100644 --- a/src/cohere/aws_client.py +++ b/src/cohere/aws_client.py @@ -239,6 +239,7 @@ def _event_hook(request: httpx.Request) -> None: ) request.url = URL(url) request.headers["host"] = request.url.host + headers["host"] = request.url.host if endpoint == "rerank": body["api_version"] = get_api_version(version=api_version) diff --git a/src/cohere/manually_maintained/cohere_aws/client.py b/src/cohere/manually_maintained/cohere_aws/client.py index 065a5fd04..03d3592ee 100644 --- a/src/cohere/manually_maintained/cohere_aws/client.py +++ b/src/cohere/manually_maintained/cohere_aws/client.py @@ -20,17 +20,23 @@ class Client: def __init__( self, aws_region: typing.Optional[str] = None, + mode: Mode = Mode.SAGEMAKER, ): """ By default we assume region configured in AWS CLI (`aws configure get region`). You can change the region with `aws configure set region us-west-2` or override it with `region_name` parameter. """ - self._client = lazy_boto3().client("sagemaker-runtime", region_name=aws_region) - self._service_client = lazy_boto3().client("sagemaker", region_name=aws_region) + self.mode = mode if os.environ.get('AWS_DEFAULT_REGION') is None: os.environ['AWS_DEFAULT_REGION'] = aws_region - self._sess = lazy_sagemaker().Session(sagemaker_client=self._service_client) - self.mode = Mode.SAGEMAKER + + if self.mode == Mode.SAGEMAKER: + self._client = lazy_boto3().client("sagemaker-runtime", region_name=aws_region) + self._service_client = lazy_boto3().client("sagemaker", region_name=aws_region) + self._sess = lazy_sagemaker().Session(sagemaker_client=self._service_client) + elif self.mode == Mode.BEDROCK: + self._client = lazy_boto3().client("bedrock-runtime", region_name=aws_region) + self._service_client = lazy_boto3().client("bedrock", region_name=aws_region) @@ -550,11 +556,15 @@ def embed( variant: Optional[str] = None, input_type: Optional[str] = None, model_id: Optional[str] = None, + output_dimension: Optional[int] = None, + embedding_types: Optional[List[str]] = None, ) -> Embeddings: json_params = { 'texts': texts, 'truncate': truncate, - "input_type": input_type + "input_type": input_type, + "output_dimension": output_dimension, + "embedding_types": embedding_types, } for key, value in list(json_params.items()): if value is None: From 065e0106dd3164db03232d828d4f32cf6df24f68 Mon Sep 17 00:00:00 2001 From: fern-support <126544928+fern-support@users.noreply.github.com> Date: Wed, 11 Feb 2026 12:40:22 -0500 Subject: [PATCH 2/8] test: add integration tests for AWS client fixes Add skipped integration tests (gated by TEST_AWS) covering: - BedrockClientV2 embed with SigV4 signing (validates host header fix) - cohere_aws.Client in Bedrock mode (validates mode param fix) - embed() with output_dimension and embedding_types (validates v4 params) Co-Authored-By: Claude Opus 4.6 --- tests/test_bedrock_client.py | 108 +++++++++++++++++++++++++++++++++++ 1 file changed, 108 insertions(+) diff --git a/tests/test_bedrock_client.py b/tests/test_bedrock_client.py index d588ca38c..d5edc9774 100644 --- a/tests/test_bedrock_client.py +++ b/tests/test_bedrock_client.py @@ -10,6 +10,17 @@ aws_region = os.getenv("AWS_REGION") endpoint_type = os.getenv("ENDPOINT_TYPE") + +def _setup_boto3_env(): + """Bridge custom test env vars to standard boto3 credential env vars.""" + if aws_access_key: + os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key + if aws_secret_key: + os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_key + if aws_session_token: + os.environ["AWS_SESSION_TOKEN"] = aws_session_token + + @unittest.skipIf(None == os.getenv("TEST_AWS"), "tests skipped because TEST_AWS is not set") class TestClient(unittest.TestCase): platform: str = "bedrock" @@ -109,3 +120,100 @@ def test_chat_stream(self) -> None: self.assertIsNotNone(event.response.text) self.assertSetEqual(response_types, {"text-generation", "stream-end"}) + + +@unittest.skipIf(None == os.getenv("TEST_AWS"), "tests skipped because TEST_AWS is not set") +class TestBedrockClientV2(unittest.TestCase): + """Integration tests for BedrockClientV2 (httpx-based). + + Fix 1 validation: If these pass, SigV4 signing uses the correct host header, + since the request would fail with a signature mismatch otherwise. + """ + + client: cohere.ClientV2 = cohere.BedrockClientV2( + aws_access_key=aws_access_key, + aws_secret_key=aws_secret_key, + aws_session_token=aws_session_token, + aws_region=aws_region, + ) + + def test_embed(self) -> None: + response = self.client.embed( + model="cohere.embed-multilingual-v3", + texts=["I love Cohere!"], + input_type="search_document", + embedding_types=["float"], + ) + self.assertIsNotNone(response) + + def test_embed_with_output_dimension(self) -> None: + response = self.client.embed( + model="cohere.embed-english-v3", + texts=["I love Cohere!"], + input_type="search_document", + embedding_types=["float"], + output_dimension=256, + ) + self.assertIsNotNone(response) + + +@unittest.skipIf(None == os.getenv("TEST_AWS"), "tests skipped because TEST_AWS is not set") +class TestCohereAwsBedrockClient(unittest.TestCase): + """Integration tests for cohere_aws.Client in Bedrock mode (boto3-based). + + Validates: + - Fix 2: Client can be initialized with mode=BEDROCK without importing sagemaker + - Fix 3: embed() accepts output_dimension and embedding_types + """ + + @classmethod + def setUpClass(cls) -> None: + _setup_boto3_env() + from cohere.manually_maintained.cohere_aws.client import Client + from cohere.manually_maintained.cohere_aws.mode import Mode + cls.client = Client(aws_region=aws_region, mode=Mode.BEDROCK) + + def test_client_is_bedrock_mode(self) -> None: + from cohere.manually_maintained.cohere_aws.mode import Mode + self.assertEqual(self.client.mode, Mode.BEDROCK) + + def test_embed(self) -> None: + response = self.client.embed( + texts=["I love Cohere!"], + input_type="search_document", + model_id="cohere.embed-multilingual-v3", + ) + self.assertIsNotNone(response) + self.assertIsNotNone(response.embeddings) + self.assertGreater(len(response.embeddings), 0) + + def test_embed_with_embedding_types(self) -> None: + response = self.client.embed( + texts=["I love Cohere!"], + input_type="search_document", + model_id="cohere.embed-multilingual-v3", + embedding_types=["float"], + ) + self.assertIsNotNone(response) + self.assertIsNotNone(response.embeddings) + + def test_embed_with_output_dimension(self) -> None: + response = self.client.embed( + texts=["I love Cohere!"], + input_type="search_document", + model_id="cohere.embed-english-v3", + output_dimension=256, + embedding_types=["float"], + ) + self.assertIsNotNone(response) + self.assertIsNotNone(response.embeddings) + + def test_embed_without_new_params(self) -> None: + """Backwards compat: embed() still works without the new v4 params.""" + response = self.client.embed( + texts=["I love Cohere!"], + input_type="search_document", + model_id="cohere.embed-multilingual-v3", + ) + self.assertIsNotNone(response) + self.assertIsNotNone(response.embeddings) From b20b617a0b87fef2174af692234750c0221859fc Mon Sep 17 00:00:00 2001 From: fern-support <126544928+fern-support@users.noreply.github.com> Date: Wed, 11 Feb 2026 12:54:42 -0500 Subject: [PATCH 3/8] fix: guard SageMaker-only methods in Bedrock mode Address review feedback: In Bedrock mode, `self._sess` was never set, so SageMaker-only methods would throw confusing AttributeErrors. Now: - Initialize `_sess=None` and `_endpoint_name=None` in Bedrock mode - Add `_require_sagemaker()` guard to connect_to_endpoint, create_endpoint, export_finetune, summarize, and delete_endpoint Co-Authored-By: Claude Opus 4.6 --- src/cohere/manually_maintained/cohere_aws/client.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/src/cohere/manually_maintained/cohere_aws/client.py b/src/cohere/manually_maintained/cohere_aws/client.py index 03d3592ee..ee9b94ce8 100644 --- a/src/cohere/manually_maintained/cohere_aws/client.py +++ b/src/cohere/manually_maintained/cohere_aws/client.py @@ -37,8 +37,12 @@ def __init__( elif self.mode == Mode.BEDROCK: self._client = lazy_boto3().client("bedrock-runtime", region_name=aws_region) self._service_client = lazy_boto3().client("bedrock", region_name=aws_region) + self._sess = None + self._endpoint_name = None - + def _require_sagemaker(self) -> None: + if self.mode != Mode.SAGEMAKER: + raise CohereError("This method is only supported in SageMaker mode.") def _does_endpoint_exist(self, endpoint_name: str) -> bool: try: @@ -56,6 +60,7 @@ def connect_to_endpoint(self, endpoint_name: str) -> None: Raises: CohereError: Connection to the endpoint failed. """ + self._require_sagemaker() if not self._does_endpoint_exist(endpoint_name): raise CohereError(f"Endpoint {endpoint_name} does not exist.") self._endpoint_name = endpoint_name @@ -143,6 +148,7 @@ def create_endpoint( will be used to get the role. This should work when one uses the client inside SageMaker. If this errors out, the default role "ServiceRoleSagemaker" will be used, which generally works outside of SageMaker. """ + self._require_sagemaker() # First, check if endpoint already exists if self._does_endpoint_exist(endpoint_name): if recreate: @@ -815,6 +821,7 @@ def export_finetune( This should work when one uses the client inside SageMaker. If this errors out, the default role "ServiceRoleSagemaker" will be used, which generally works outside SageMaker. """ + self._require_sagemaker() if name == "model": raise ValueError("name cannot be 'model'") @@ -958,6 +965,7 @@ def summarize( additional_command: Optional[str] = "", variant: Optional[str] = None ) -> Summary: + self._require_sagemaker() if self._endpoint_name is None: raise CohereError("No endpoint connected. " @@ -999,6 +1007,7 @@ def summarize( def delete_endpoint(self) -> None: + self._require_sagemaker() if self._endpoint_name is None: raise CohereError("No endpoint connected.") try: From a04c0822de630e26c39eed2bbcbb170c45856d34 Mon Sep 17 00:00:00 2001 From: fern-support <126544928+fern-support@users.noreply.github.com> Date: Wed, 11 Feb 2026 12:59:43 -0500 Subject: [PATCH 4/8] fix: add type annotation for TestCohereAwsBedrockClient.client Fixes mypy attr-defined error in CI. Co-Authored-By: Claude Opus 4.6 --- tests/test_bedrock_client.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_bedrock_client.py b/tests/test_bedrock_client.py index d5edc9774..819399c8c 100644 --- a/tests/test_bedrock_client.py +++ b/tests/test_bedrock_client.py @@ -165,6 +165,7 @@ class TestCohereAwsBedrockClient(unittest.TestCase): - Fix 2: Client can be initialized with mode=BEDROCK without importing sagemaker - Fix 3: embed() accepts output_dimension and embedding_types """ + client: typing.Any = None @classmethod def setUpClass(cls) -> None: From bae43fa4da0ec5ee7c69f80162242f6f81f2522e Mon Sep 17 00:00:00 2001 From: fern-support <126544928+fern-support@users.noreply.github.com> Date: Wed, 11 Feb 2026 13:01:48 -0500 Subject: [PATCH 5/8] test: add mocked unit tests for AWS client fixes These run in CI without AWS credentials, covering: - SigV4 signing uses correct host header after URL rewrite - Mode-conditional boto3 client initialization (sagemaker vs bedrock) - Default mode is SAGEMAKER for backwards compat - embed() accepts, passes, and strips output_dimension/embedding_types Co-Authored-By: Claude Opus 4.6 --- tests/test_aws_client_unit.py | 208 ++++++++++++++++++++++++++++++++++ 1 file changed, 208 insertions(+) create mode 100644 tests/test_aws_client_unit.py diff --git a/tests/test_aws_client_unit.py b/tests/test_aws_client_unit.py new file mode 100644 index 000000000..34d17c99f --- /dev/null +++ b/tests/test_aws_client_unit.py @@ -0,0 +1,208 @@ +""" +Unit tests (mocked, no AWS credentials needed) for AWS client fixes. + +Covers: +- Fix 1: SigV4 signing uses the correct host header after URL rewrite +- Fix 2: cohere_aws.Client conditionally initializes based on mode +- Fix 3: embed() accepts and passes output_dimension and embedding_types +""" + +import inspect +import json +import os +import unittest +from unittest.mock import MagicMock, patch + +import httpx + +from cohere.manually_maintained.cohere_aws.mode import Mode + + +class TestSigV4HostHeader(unittest.TestCase): + """Fix 1: The headers dict passed to AWSRequest for SigV4 signing must + contain the rewritten Bedrock/SageMaker host, not the stale api.cohere.com.""" + + def test_sigv4_signs_with_correct_host(self) -> None: + captured_aws_request_kwargs: dict = {} + + mock_aws_request_cls = MagicMock() + + def capture_aws_request(**kwargs): # type: ignore + captured_aws_request_kwargs.update(kwargs) + mock_req = MagicMock() + mock_req.prepare.return_value = MagicMock( + headers={"host": "bedrock-runtime.us-east-1.amazonaws.com"} + ) + return mock_req + + mock_aws_request_cls.side_effect = capture_aws_request + + mock_botocore = MagicMock() + mock_botocore.awsrequest.AWSRequest = mock_aws_request_cls + mock_botocore.auth.SigV4Auth.return_value = MagicMock() + + mock_boto3 = MagicMock() + mock_session = MagicMock() + mock_session.region_name = "us-east-1" + mock_session.get_credentials.return_value = MagicMock() + mock_boto3.Session.return_value = mock_session + + with patch("cohere.aws_client.lazy_botocore", return_value=mock_botocore), \ + patch("cohere.aws_client.lazy_boto3", return_value=mock_boto3): + + from cohere.aws_client import map_request_to_bedrock + + hook = map_request_to_bedrock(service="bedrock", aws_region="us-east-1") + + request = httpx.Request( + method="POST", + url="https://api.cohere.com/v1/chat", + headers={"connection": "keep-alive"}, + json={"model": "cohere.command-r-plus-v1:0", "message": "hello"}, + ) + + self.assertEqual(request.url.host, "api.cohere.com") + + hook(request) + + self.assertIn("bedrock-runtime.us-east-1.amazonaws.com", str(request.url)) + + signed_headers = captured_aws_request_kwargs["headers"] + self.assertEqual( + signed_headers["host"], + "bedrock-runtime.us-east-1.amazonaws.com", + ) + + +class TestModeConditionalInit(unittest.TestCase): + """Fix 2: cohere_aws.Client should initialize different boto3 clients + depending on mode, and default to SAGEMAKER for backwards compat.""" + + def test_sagemaker_mode_creates_sagemaker_clients(self) -> None: + mock_boto3 = MagicMock() + mock_sagemaker = MagicMock() + + with patch("cohere.manually_maintained.cohere_aws.client.lazy_boto3", return_value=mock_boto3), \ + patch("cohere.manually_maintained.cohere_aws.client.lazy_sagemaker", return_value=mock_sagemaker), \ + patch.dict(os.environ, {"AWS_DEFAULT_REGION": "us-east-1"}): + + from cohere.manually_maintained.cohere_aws.client import Client + + client = Client(aws_region="us-east-1") + + self.assertEqual(client.mode, Mode.SAGEMAKER) + + service_names = [c[0][0] for c in mock_boto3.client.call_args_list] + self.assertIn("sagemaker-runtime", service_names) + self.assertIn("sagemaker", service_names) + self.assertNotIn("bedrock-runtime", service_names) + self.assertNotIn("bedrock", service_names) + + mock_sagemaker.Session.assert_called_once() + + def test_bedrock_mode_creates_bedrock_clients(self) -> None: + mock_boto3 = MagicMock() + mock_sagemaker = MagicMock() + + with patch("cohere.manually_maintained.cohere_aws.client.lazy_boto3", return_value=mock_boto3), \ + patch("cohere.manually_maintained.cohere_aws.client.lazy_sagemaker", return_value=mock_sagemaker), \ + patch.dict(os.environ, {"AWS_DEFAULT_REGION": "us-west-2"}): + + from cohere.manually_maintained.cohere_aws.client import Client + + client = Client(aws_region="us-west-2", mode=Mode.BEDROCK) + + self.assertEqual(client.mode, Mode.BEDROCK) + + service_names = [c[0][0] for c in mock_boto3.client.call_args_list] + self.assertIn("bedrock-runtime", service_names) + self.assertIn("bedrock", service_names) + self.assertNotIn("sagemaker-runtime", service_names) + self.assertNotIn("sagemaker", service_names) + + mock_sagemaker.Session.assert_not_called() + + def test_default_mode_is_sagemaker(self) -> None: + from cohere.manually_maintained.cohere_aws.client import Client + + sig = inspect.signature(Client.__init__) + self.assertEqual(sig.parameters["mode"].default, Mode.SAGEMAKER) + + +class TestEmbedV4Params(unittest.TestCase): + """Fix 3: embed() should accept output_dimension and embedding_types, + pass them through to the request body, and strip them when None.""" + + @staticmethod + def _make_bedrock_client(): # type: ignore + mock_boto3 = MagicMock() + mock_botocore = MagicMock() + captured_body: dict = {} + + def fake_invoke_model(**kwargs): # type: ignore + captured_body.update(json.loads(kwargs["body"])) + mock_body = MagicMock() + mock_body.read.return_value = json.dumps({"embeddings": [[0.1, 0.2]]}).encode() + return {"body": mock_body} + + mock_bedrock_client = MagicMock() + mock_bedrock_client.invoke_model.side_effect = fake_invoke_model + + def fake_boto3_client(service_name, **kwargs): # type: ignore + if service_name == "bedrock-runtime": + return mock_bedrock_client + return MagicMock() + + mock_boto3.client.side_effect = fake_boto3_client + return mock_boto3, mock_botocore, captured_body + + def test_embed_accepts_new_params(self) -> None: + from cohere.manually_maintained.cohere_aws.client import Client + + sig = inspect.signature(Client.embed) + self.assertIn("output_dimension", sig.parameters) + self.assertIn("embedding_types", sig.parameters) + self.assertIsNone(sig.parameters["output_dimension"].default) + self.assertIsNone(sig.parameters["embedding_types"].default) + + def test_embed_passes_params_to_bedrock(self) -> None: + mock_boto3, mock_botocore, captured_body = self._make_bedrock_client() + + with patch("cohere.manually_maintained.cohere_aws.client.lazy_boto3", return_value=mock_boto3), \ + patch("cohere.manually_maintained.cohere_aws.client.lazy_botocore", return_value=mock_botocore), \ + patch("cohere.manually_maintained.cohere_aws.client.lazy_sagemaker", return_value=MagicMock()), \ + patch.dict(os.environ, {"AWS_DEFAULT_REGION": "us-east-1"}): + + from cohere.manually_maintained.cohere_aws.client import Client + + client = Client(aws_region="us-east-1", mode=Mode.BEDROCK) + client.embed( + texts=["hello world"], + input_type="search_document", + model_id="cohere.embed-english-v3", + output_dimension=256, + embedding_types=["float", "int8"], + ) + + self.assertEqual(captured_body["output_dimension"], 256) + self.assertEqual(captured_body["embedding_types"], ["float", "int8"]) + + def test_embed_omits_none_params(self) -> None: + mock_boto3, mock_botocore, captured_body = self._make_bedrock_client() + + with patch("cohere.manually_maintained.cohere_aws.client.lazy_boto3", return_value=mock_boto3), \ + patch("cohere.manually_maintained.cohere_aws.client.lazy_botocore", return_value=mock_botocore), \ + patch("cohere.manually_maintained.cohere_aws.client.lazy_sagemaker", return_value=MagicMock()), \ + patch.dict(os.environ, {"AWS_DEFAULT_REGION": "us-east-1"}): + + from cohere.manually_maintained.cohere_aws.client import Client + + client = Client(aws_region="us-east-1", mode=Mode.BEDROCK) + client.embed( + texts=["hello world"], + input_type="search_document", + model_id="cohere.embed-english-v3", + ) + + self.assertNotIn("output_dimension", captured_body) + self.assertNotIn("embedding_types", captured_body) From 5f77e85f228758657d8a36870d213a854205b44c Mon Sep 17 00:00:00 2001 From: jsklan Date: Tue, 17 Feb 2026 14:55:26 -0500 Subject: [PATCH 6/8] fix: handle embedding_types dict response in AWS client embed methods When embedding_types is specified, the Cohere API returns embeddings as a dict (e.g. {"float": [[...]], "int8": [[...]]}) instead of a flat list. Both _bedrock_embed and _sagemaker_embed now detect the dict format and return it directly instead of wrapping it in Embeddings, which would silently produce wrong results for len() and iteration. Co-Authored-By: Claude Opus 4.6 --- .../manually_maintained/cohere_aws/client.py | 12 +++-- tests/test_aws_client_unit.py | 44 +++++++++++++++++++ 2 files changed, 53 insertions(+), 3 deletions(-) diff --git a/src/cohere/manually_maintained/cohere_aws/client.py b/src/cohere/manually_maintained/cohere_aws/client.py index ee9b94ce8..84d0da0ae 100644 --- a/src/cohere/manually_maintained/cohere_aws/client.py +++ b/src/cohere/manually_maintained/cohere_aws/client.py @@ -564,7 +564,7 @@ def embed( model_id: Optional[str] = None, output_dimension: Optional[int] = None, embedding_types: Optional[List[str]] = None, - ) -> Embeddings: + ) -> Union[Embeddings, Dict[str, List]]: json_params = { 'texts': texts, 'truncate': truncate, @@ -607,7 +607,10 @@ def _sagemaker_embed(self, json_params: Dict[str, Any], variant: str): # ValidationError, e.g. when variant is bad raise CohereError(str(e)) - return Embeddings(response['embeddings']) + embeddings = response['embeddings'] + if isinstance(embeddings, dict): + return embeddings + return Embeddings(embeddings) def _bedrock_embed(self, json_params: Dict[str, Any], model_id: str): if not model_id: @@ -628,7 +631,10 @@ def _bedrock_embed(self, json_params: Dict[str, Any], model_id: str): # ValidationError, e.g. when variant is bad raise CohereError(str(e)) - return Embeddings(response['embeddings']) + embeddings = response['embeddings'] + if isinstance(embeddings, dict): + return embeddings + return Embeddings(embeddings) def rerank(self, diff --git a/tests/test_aws_client_unit.py b/tests/test_aws_client_unit.py index 34d17c99f..94e584922 100644 --- a/tests/test_aws_client_unit.py +++ b/tests/test_aws_client_unit.py @@ -206,3 +206,47 @@ def test_embed_omits_none_params(self) -> None: self.assertNotIn("output_dimension", captured_body) self.assertNotIn("embedding_types", captured_body) + + def test_embed_with_embedding_types_returns_dict(self) -> None: + """When embedding_types is specified, the API returns embeddings as a dict. + The client should return that dict rather than wrapping it in Embeddings.""" + mock_boto3 = MagicMock() + mock_botocore = MagicMock() + + by_type_embeddings = {"float": [[0.1, 0.2]], "int8": [[1, 2]]} + + def fake_invoke_model(**kwargs): # type: ignore + mock_body = MagicMock() + mock_body.read.return_value = json.dumps({ + "embeddings": by_type_embeddings, + "response_type": "embeddings_by_type", + }).encode() + return {"body": mock_body} + + mock_bedrock_client = MagicMock() + mock_bedrock_client.invoke_model.side_effect = fake_invoke_model + + def fake_boto3_client(service_name, **kwargs): # type: ignore + if service_name == "bedrock-runtime": + return mock_bedrock_client + return MagicMock() + + mock_boto3.client.side_effect = fake_boto3_client + + with patch("cohere.manually_maintained.cohere_aws.client.lazy_boto3", return_value=mock_boto3), \ + patch("cohere.manually_maintained.cohere_aws.client.lazy_botocore", return_value=mock_botocore), \ + patch("cohere.manually_maintained.cohere_aws.client.lazy_sagemaker", return_value=MagicMock()), \ + patch.dict(os.environ, {"AWS_DEFAULT_REGION": "us-east-1"}): + + from cohere.manually_maintained.cohere_aws.client import Client + + client = Client(aws_region="us-east-1", mode=Mode.BEDROCK) + result = client.embed( + texts=["hello world"], + input_type="search_document", + model_id="cohere.embed-english-v3", + embedding_types=["float", "int8"], + ) + + self.assertIsInstance(result, dict) + self.assertEqual(result, by_type_embeddings) From 9050f96decfc176a0df7f4a6b6cd79081506a37e Mon Sep 17 00:00:00 2001 From: fern-support <126544928+fern-support@users.noreply.github.com> Date: Wed, 18 Feb 2026 09:26:31 -0500 Subject: [PATCH 7/8] Update sdk version --- pyproject.toml | 2 +- src/cohere/core/client_wrapper.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 04a187b5c..2eb0b0b51 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ name = "cohere" [tool.poetry] name = "cohere" -version = "5.20.4" +version = "5.20.6" description = "" readme = "README.md" authors = [] diff --git a/src/cohere/core/client_wrapper.py b/src/cohere/core/client_wrapper.py index 8a23c07f0..ac7e71d6b 100644 --- a/src/cohere/core/client_wrapper.py +++ b/src/cohere/core/client_wrapper.py @@ -24,10 +24,10 @@ def __init__( def get_headers(self) -> typing.Dict[str, str]: headers: typing.Dict[str, str] = { - "User-Agent": "cohere/5.20.4", + "User-Agent": "cohere/5.20.6", "X-Fern-Language": "Python", "X-Fern-SDK-Name": "cohere", - "X-Fern-SDK-Version": "5.20.4", + "X-Fern-SDK-Version": "5.20.6", **(self.get_custom_headers() or {}), } if self._client_name is not None: From 54d22f2a0b3ef0780381bcfd230a95f758bd6f4c Mon Sep 17 00:00:00 2001 From: fern-support <126544928+fern-support@users.noreply.github.com> Date: Wed, 18 Feb 2026 10:25:22 -0500 Subject: [PATCH 8/8] fix: correct assertions for dict response in embed integration tests When embedding_types is passed, _bedrock_embed returns a raw dict instead of an Embeddings object. Update test assertions to check for dict type and key presence instead of accessing .embeddings attribute. Co-Authored-By: Claude Opus 4.6 --- tests/test_bedrock_client.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/test_bedrock_client.py b/tests/test_bedrock_client.py index 819399c8c..f16eec3b5 100644 --- a/tests/test_bedrock_client.py +++ b/tests/test_bedrock_client.py @@ -196,7 +196,9 @@ def test_embed_with_embedding_types(self) -> None: embedding_types=["float"], ) self.assertIsNotNone(response) - self.assertIsNotNone(response.embeddings) + # When embedding_types is passed, the response is a raw dict + self.assertIsInstance(response, dict) + self.assertIn("float", response) def test_embed_with_output_dimension(self) -> None: response = self.client.embed( @@ -207,7 +209,9 @@ def test_embed_with_output_dimension(self) -> None: embedding_types=["float"], ) self.assertIsNotNone(response) - self.assertIsNotNone(response.embeddings) + # When embedding_types is passed, the response is a raw dict + self.assertIsInstance(response, dict) + self.assertIn("float", response) def test_embed_without_new_params(self) -> None: """Backwards compat: embed() still works without the new v4 params."""