fix: resolve AWS client SigV4 signing, forced SageMaker dep, and missing embed params #728
Conversation
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
…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 <[email protected]>
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
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 <[email protected]>
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 <[email protected]>
Fixes mypy attr-defined error in CI. Co-Authored-By: Claude Opus 4.6 <[email protected]>
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 <[email protected]>
| "input_type": input_type | ||
| "input_type": input_type, | ||
| "output_dimension": output_dimension, | ||
| "embedding_types": embedding_types, |
There was a problem hiding this comment.
Embed response parsing breaks with embedding_types parameter
Medium Severity
Adding embedding_types to the request parameters changes the Cohere API response format, but response parsing isn't updated. When embedding_types is specified, the API returns embeddings as a dict (e.g., {"float": [[...]], "int8": [[...]]}) with response_type: "embeddings_by_type", instead of a flat List[List[float]]. Both _bedrock_embed and _sagemaker_embed unconditionally pass response['embeddings'] to the Embeddings constructor, which expects a list. This causes len() and iteration on the returned Embeddings object to silently produce wrong results.
Additional Locations (1)
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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 request signing and AWS client initialization paths, where regressions could break Bedrock/SageMaker calls at runtime. Test coverage was added, but changes affect core connectivity/auth behavior.
Overview
Fixes AWS SigV4 signing for
AwsClientV2by ensuring the host header used during signing is updated after rewriting requests to Bedrock/SageMaker endpoints.Updates
cohere_aws.Clientto support a newmode(defaultMode.SAGEMAKER) that conditionally initializes Bedrock vs SageMaker boto3 clients and gates SageMaker-only operations with a clear error when used in Bedrock mode.Extends
cohere_aws.Client.embed()to accept and forwardoutput_dimensionandembedding_types, and adds unit + gated integration tests covering the signing fix, mode initialization, and new embed params.Written by Cursor Bugbot for commit bae43fa. This will update automatically on new commits. Configure here.