Skip to content

Conversation

@acastro2
Copy link
Contributor

Summary

Fixes support for AWS Bedrock Cohere v4 embedding models (cohere.embed-v4:0) by handling the new response format introduced in v4. The integration now automatically detects and handles both v3 and v4 response formats while maintaining full backward compatibility.

Problem

When using AWS Bedrock Cohere v4 embedding models through the BedrockEmbedding class, the embedding process fails with a validation error due to a format change in the response.

Root Cause: Cohere v4 models return embeddings in a different format than v3:

  • v3 format: {"embeddings": [[vec1], [vec2]]}
  • v4 format (multiple types): {"embeddings": {"float": [[vec1], [vec2]]}}

Solution

Added a new _parse_cohere_response() helper function that:

  • Detects and handles all Cohere response formats (v3 and v4)
  • Supports nested format: {"embeddings": {"float": [...]}}
  • Supports direct format: {"float": [...]}
  • Provides clear error messages for unexpected formats
  • Maintains 100% backward compatibility with v3 models

Changes

Code (base.py)

  • Added Models.COHERE_EMBED_V4 = "cohere.embed-v4:0" to supported models enum
  • Added _parse_cohere_response() function with comprehensive format handling
  • Updated PROVIDER_SPECIFIC_IDENTIFIERS to use the new parser

Tests (test_bedrock.py)

  • Added 4 comprehensive test cases for v4 format handling
  • Updated test_list_supported_models() to include v4 model

Documentation (README.md)

  • Added comprehensive usage examples
  • Documented all supported models including cohere.embed-v4:0
  • Added Cohere v4 support section with format explanation

Testing

All 13 tests passing (11 existing + 4 new)
Test coverage: 68% (exceeds 50% minimum requirement)
All linting checks pass (ruff, black, mypy, prettier, codespell)
Backward compatibility maintained - All existing v3 tests still pass

Verification

Verified against official AWS Bedrock documentation:

Related

Similar to the Cohere v4 support added to llama-index-embeddings-cohere in PR #18464

@acastro2 acastro2 force-pushed the bedrock-coherev4-support branch from 5966bfb to e26ae0e Compare October 15, 2025 20:27
@acastro2 acastro2 marked this pull request as ready for review October 15, 2025 20:27
@acastro2
Copy link
Contributor Author

Manually tested against my code and its working!

@dosubot dosubot bot added the size:L This PR changes 100-499 lines, ignoring generated files. label Oct 15, 2025
@acastro2 acastro2 force-pushed the bedrock-coherev4-support branch from 911200d to 4cf7998 Compare October 16, 2025 13:46
@acastro2 acastro2 force-pushed the bedrock-coherev4-support branch from 4cf7998 to 8d7fbed Compare October 16, 2025 13:49
Copy link
Member

@AstraBert AstraBert left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm!

And thanks for adding the readme :)

@AstraBert AstraBert merged commit 3899161 into run-llama:main Oct 16, 2025
11 checks passed
@dosubot dosubot bot added the lgtm This PR has been approved by a maintainer label Oct 16, 2025
@acastro2 acastro2 deleted the bedrock-coherev4-support branch October 16, 2025 16:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

lgtm This PR has been approved by a maintainer size:L This PR changes 100-499 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants