Repository navigation
Python: Fix Bedrock Cohere embeddings when embedding_types is set - #14678
alexcastgonz (pcbeingused333) wants to merge 2 commits into
Conversation
When embedding_types is set, Cohere on Bedrock returns the embeddings keyed by type (a dict) instead of a list of vectors, and parse_text_embedding_response rejected it with ServiceInvalidResponseError. Return the float embedding when present, otherwise the first type in the response.
GitHub Actions
Active
semantic-kernel-automation
Bot
added
the
python
There was a problem hiding this comment.
🟡 Changes recommended
Typed embedding responses can accept malformed vectors or select inconsistent encodings across a batch, potentially causing dimension errors.
2 open issues
What changed in this PR
This PR extends the Python Bedrock Cohere embedding parser to handle responses keyed by embedding type, addressing failures when embedding_types is requested.
Changes:
- Accept typed responses, preferring
floatand otherwise selecting the first returned type. - Preserve the existing plain-list response path.
- Add tests for typed embeddings, request forwarding, and empty or malformed typed responses.
| File | Description |
|---|---|
python/tests/unit/connectors/ai/bedrock/services/test_bedrock_text_embedding_generation.py |
Adds successful and invalid typed-response test cases. |
python/semantic_kernel/connectors/ai/bedrock/services/model_provider/bedrock_cohere.py |
Adds extraction of embeddings from type-keyed responses. |
Review effort: Balanced
| if isinstance(embeddings, dict): | ||
| # With `embedding_types` set, the embeddings are keyed by type ({"float": [[...]], "int8": [[...]]}). | ||
| # Return the float embedding when it was requested, otherwise the first type in the response. | ||
| embeddings = embeddings.get("float") or next(iter(embeddings.values()), None) |
There was a problem hiding this comment.
Confirmed on the current revision. The parser returns 0.1 for {"float": [0.1, 0.2]}, [] for {"float": [[]]}, and [1, 2] for {"int8": [[1, 2]], "float": []}. These cases still need validation before this is ready.
| if isinstance(embeddings, dict): | ||
| # With `embedding_types` set, the embeddings are keyed by type ({"float": [[...]], "int8": [[...]]}). | ||
| # Return the float embedding when it was requested, otherwise the first type in the response. | ||
| embeddings = embeddings.get("float") or next(iter(embeddings.values()), None) |
There was a problem hiding this comment.
Confirmed with the current parser: {"int8": [[1, 2]], "binary": [[1]]} returns [1, 2], while reversing those keys returns [1]. Identical requested embedding types can therefore produce different vector dimensions across texts. Selection needs a stable type policy rather than JSON object order.
…priority Select float by key presence, fall back in a fixed priority (int8, uint8, binary, ubinary) independent of JSON key order, and raise ServiceInvalidResponseError when the selected value is not a non-empty numeric vector.
GitHub Actions
Active
|
Thanks, both points are addressed in cad9abe. With This comment was drafted with Claude Code. |
Motivation and Context
BedrockTextEmbeddingaccepts the Cohereembedding_typesoption and forwards it in the request body. With it set, Cohere on Bedrock returns the embeddings keyed by type ({"embeddings": {"float": [[...]], "int8": [[...]]}}) instead of a plain list of vectors, andparse_text_embedding_responseonly accepted a list, so the call failed withServiceInvalidResponseError: The response from Cohere model does not contain embeddings.even though the request succeeded.Description
parse_text_embedding_responseinbedrock_cohere.pynow also accepts the keyed form: it returns thefloatembedding when that type is present, otherwise the first type in the response. The plain list form and the error for missing or malformed embeddings are unchanged.Tests (
tests/unit/connectors/ai/bedrock/services/test_bedrock_text_embedding_generation.py): typed responses (float only, float plus int8, int8 only) and the plain list throughgenerate_embeddingswithembedding_typesin the settings, and empty or malformed typed responses still raisingServiceInvalidResponseError.tests/unit/connectors/ai/bedrock: 123 passed; 3 of the new cases fail without the change. Not run against the real Bedrock API.Contribution Checklist
Disclosure: this change was prepared with an AI coding assistant (Claude) working for the account owner.