fix: resolve AWS client SigV4 signing, forced SageMaker dep, and missing embed params (#728)

* 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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

* fix: add type annotation for TestCohereAwsBedrockClient.client

Fixes mypy attr-defined error in CI.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

* Update sdk version

* 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 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: jsklan <jsklan.development@gmail.com>

Author: 3 people

pyproject.toml

@@ -4,7 +4,7 @@ dynamic = ["version"]
 
 [tool.poetry]
 name = "cohere"
-version = "5.20.5"
+version = "5.20.6"
 description = ""
 readme = "README.md"
 authors = []

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)

src/cohere/core/client_wrapper.py

@@ -26,12 +26,12 @@ def get_headers(self) -> typing.Dict[str, str]:
         import platform
 
         headers: typing.Dict[str, str] = {
-            "User-Agent": "cohere/5.20.5",
+            "User-Agent": "cohere/5.20.6",
             "X-Fern-Language": "Python",
             "X-Fern-Runtime": f"python/{platform.python_version()}",
             "X-Fern-Platform": f"{platform.system().lower()}/{platform.release()}",
             "X-Fern-SDK-Name": "cohere",
-            "X-Fern-SDK-Version": "5.20.5",
+            "X-Fern-SDK-Version": "5.20.6",
             **(self.get_custom_headers() or {}),
         }
         if self._client_name is not None:

src/cohere/manually_maintained/cohere_aws/client.py

@@ -20,19 +20,29 @@ 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)
+            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:
@@ -50,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
@@ -137,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:
@@ -550,11 +562,15 @@ def embed(
         variant: Optional[str] = None,
         input_type: Optional[str] = None,
         model_id: Optional[str] = None,
-    ) -> Embeddings:
+        output_dimension: Optional[int] = None,
+        embedding_types: Optional[List[str]] = None,
+    ) -> Union[Embeddings, Dict[str, List]]:
         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:
@@ -591,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:
@@ -612,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,
@@ -805,6 +827,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'")
 
@@ -948,6 +971,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. "
@@ -989,6 +1013,7 @@ def summarize(
 
 
     def delete_endpoint(self) -> None:
+        self._require_sagemaker()
         if self._endpoint_name is None:
             raise CohereError("No endpoint connected.")
         try: