fix: resolve AWS client SigV4 signing, forced SageMaker dep, and missing embed params#728
Merged
jasonozuzu-cohere merged 9 commits intomainfrom Feb 18, 2026
Merged
Conversation
…ing 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>
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>
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>
Fixes mypy attr-defined error in CI. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
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>
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>
jasonozuzu-cohere
approved these changes
Feb 18, 2026
jasonozuzu-cohere
approved these changes
Feb 18, 2026
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>
jasonozuzu-cohere
approved these changes
Feb 18, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
headersdict passed toAWSRequestfor signing contained the staleapi.cohere.comhost instead of the rewritten Bedrock/SageMaker host, causing signature mismatch errors. Fixed by syncingheaders["host"]after the URL rewrite.cohere_aws.Client.__init__unconditionally created SageMaker boto3 clients and asagemaker.Session, even for Bedrock users. Added amodeparameter (defaultMode.SAGEMAKERfor backwards compat) that conditionally initializes the correct service clients (bedrock-runtime/bedrockvssagemaker-runtime/sagemaker). SageMaker-only methods (connect_to_endpoint,create_endpoint,export_finetune,summarize,delete_endpoint) now raise a clear error if called in Bedrock mode.output_dimensionandembedding_typesto theembed()method signature, included injson_params(stripped whenNoneby existing cleanup loop).Closes #721
Test plan
test_bedrock_client.py(gated byTEST_AWSenv var) covering BedrockClientV2 signing,cohere_aws.Clientin Bedrock mode, and embed with v4 params.fernignore(safe from Fern regeneration)TEST_AWS=1+ AWS credential env vars)🤖 Generated with Claude Code
Note
Medium Risk
Touches AWS request signing and changes Bedrock/SageMaker client initialization and
embed()return types, which could break integrations if assumptions about headers, dependencies, or response shape differ.Overview
Fixes AWS SigV4 signing for Bedrock/SageMaker by ensuring the rewritten request host is also used in the header set passed to
botocore’sAWSRequest, preventing signature mismatches.Updates the manually maintained
cohere_aws.Clientto supportmode-based initialization (defaulting toSAGEMAKER) so Bedrock users no longer require SageMaker dependencies, and gates SageMaker-only operations with a clear error. Extendsembed()to accept/passoutput_dimensionandembedding_typesand returns raw dict embeddings when the API responds by type.Bumps SDK version to
5.20.6and adds unit + gated integration tests covering signing, mode init, and new embed params.Written by Cursor Bugbot for commit 54d22f2. This will update automatically on new commits. Configure here.