fix: resolve AWS client SigV4 signing, forced SageMaker dep, and missing embed params#728
Open
fern-support wants to merge 6 commits intomainfrom
Open
fix: resolve AWS client SigV4 signing, forced SageMaker dep, and missing embed params#728fern-support wants to merge 6 commits intomainfrom
fern-support wants to merge 6 commits intomainfrom
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>
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 client initialization paths; failures would surface as auth/signature errors or runtime mode mismatches, though changes are small and covered by unit/integration tests.
Overview
Fixes AWS request signing by ensuring SigV4 signs with the rewritten Bedrock/SageMaker
hostheader after the URL is updated.Updates
cohere_aws.Clientto support both SageMaker and Bedrock via a newmodeparameter (defaulting toMode.SAGEMAKER), conditionally initializing the correct boto clients and guarding SageMaker-only operations with a clear runtime error.Extends
embed()to acceptoutput_dimensionandembedding_types, passes them through to Bedrock/SageMaker requests, and returns raw dict embeddings when the API responds with embeddings-by-type. Adds unit tests for all fixes plus new gated AWS integration coverage for Bedrock clients.Written by Cursor Bugbot for commit 5f77e85. This will update automatically on new commits. Configure here.