From 7d5139667b03d2a0e034893316c94c2e5aade55b Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Thu, 29 Oct 2026 08:09:24 -0400 Subject: [PATCH 1/2] Python: Fix Bedrock Cohere embeddings when embedding_types is set 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. --- .../services/model_provider/bedrock_cohere.py | 10 ++++- .../test_bedrock_text_embedding_generation.py | 45 +++++++++++++++++++ 2 files changed, 53 insertions(+), 2 deletions(-) diff --git a/python/semantic_kernel/connectors/ai/bedrock/services/model_provider/bedrock_cohere.py b/python/semantic_kernel/connectors/ai/bedrock/services/model_provider/bedrock_cohere.py index 3e498758ae72..c6d0429e4681 100644 --- a/python/semantic_kernel/connectors/ai/bedrock/services/model_provider/bedrock_cohere.py +++ b/python/semantic_kernel/connectors/ai/bedrock/services/model_provider/bedrock_cohere.py @@ -92,10 +92,16 @@ def get_text_embedding_request_body(text: str, settings: BedrockEmbeddingPromptE def parse_text_embedding_response(response: dict[str, Any]) -> list[float]: """Parse the response from text embedding for Cohere Command models.""" - if "embeddings" not in response or not isinstance(response["embeddings"], list) or len(response["embeddings"]) == 0: + embeddings = response.get("embeddings") + 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) + + if not isinstance(embeddings, list) or len(embeddings) == 0: raise ServiceInvalidResponseError("The response from Cohere model does not contain embeddings.") - return response.get("embeddings")[0] # type: ignore + return embeddings[0] # type: ignore # endregion diff --git a/python/tests/unit/connectors/ai/bedrock/services/test_bedrock_text_embedding_generation.py b/python/tests/unit/connectors/ai/bedrock/services/test_bedrock_text_embedding_generation.py index a591f458c204..fd51fda30608 100644 --- a/python/tests/unit/connectors/ai/bedrock/services/test_bedrock_text_embedding_generation.py +++ b/python/tests/unit/connectors/ai/bedrock/services/test_bedrock_text_embedding_generation.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. +import json from unittest.mock import ANY, Mock, patch import boto3 @@ -233,3 +234,47 @@ async def test_bedrock_text_embedding_with_invalid_response( with pytest.raises(ServiceInvalidResponseError): await bedrock_text_embedding.generate_embeddings(["hello", "world"]) + + +@pytest.mark.parametrize( + "embeddings, expected", + [ + ({"float": [[0.1, 0.2, 0.3]]}, [0.1, 0.2, 0.3]), + ({"int8": [[1, 2, 3]], "float": [[0.1, 0.2, 0.3]]}, [0.1, 0.2, 0.3]), + ({"int8": [[1, 2, 3]]}, [1, 2, 3]), + ([[0.1, 0.2, 0.3]], [0.1, 0.2, 0.3]), + ], +) +async def test_bedrock_cohere_text_embedding_with_embedding_types(embeddings, expected) -> None: + """Cohere returns the embeddings keyed by type when `embedding_types` is requested.""" + body = Mock() + body.read.return_value = json.dumps({"embeddings": embeddings}) + + with patch.object(MockBedrockRuntimeClient, "invoke_model", return_value={"body": body}) as mock_model_invoke: + bedrock_text_embedding = BedrockTextEmbedding( + model_id="cohere.command", + runtime_client=MockBedrockRuntimeClient(), + client=MockBedrockClient(), + ) + + settings = BedrockEmbeddingPromptExecutionSettings(extension_data={"embedding_types": ["float", "int8"]}) + response = await bedrock_text_embedding.generate_embeddings(["hello"], settings) + + assert json.loads(mock_model_invoke.call_args.kwargs["body"])["embedding_types"] == ["float", "int8"] + assert response.tolist() == [expected] + + +@pytest.mark.parametrize("embeddings", [{}, {"float": []}, {"float": 0.1}]) +async def test_bedrock_cohere_text_embedding_with_invalid_typed_response(embeddings) -> None: + body = Mock() + body.read.return_value = json.dumps({"embeddings": embeddings}) + + with patch.object(MockBedrockRuntimeClient, "invoke_model", return_value={"body": body}): + bedrock_text_embedding = BedrockTextEmbedding( + model_id="cohere.command", + runtime_client=MockBedrockRuntimeClient(), + client=MockBedrockClient(), + ) + + with pytest.raises(ServiceInvalidResponseError): + await bedrock_text_embedding.generate_embeddings(["hello"]) From cad9abe5c2909c575779db347d5795642af11ad9 Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Mon, 2 Nov 2026 09:06:33 -0500 Subject: [PATCH 2/2] fix(bedrock): validate typed Cohere embeddings and use a stable type 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. --- .../services/model_provider/bedrock_cohere.py | 20 +++++++++++++++++-- .../test_bedrock_text_embedding_generation.py | 17 +++++++++++++++- 2 files changed, 34 insertions(+), 3 deletions(-) diff --git a/python/semantic_kernel/connectors/ai/bedrock/services/model_provider/bedrock_cohere.py b/python/semantic_kernel/connectors/ai/bedrock/services/model_provider/bedrock_cohere.py index c6d0429e4681..e9151b336c97 100644 --- a/python/semantic_kernel/connectors/ai/bedrock/services/model_provider/bedrock_cohere.py +++ b/python/semantic_kernel/connectors/ai/bedrock/services/model_provider/bedrock_cohere.py @@ -90,13 +90,29 @@ def get_text_embedding_request_body(text: str, settings: BedrockEmbeddingPromptE }) +# Preference order when the response is keyed by embedding type. It is fixed (not the JSON key order) +# so the same requested types always yield the same representation for every text. +_EMBEDDING_TYPE_PRIORITY = ("float", "int8", "uint8", "binary", "ubinary") + + +def _is_numeric_vector(value: Any) -> bool: + return ( + isinstance(value, list) + and len(value) > 0 + and all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in value) + ) + + def parse_text_embedding_response(response: dict[str, Any]) -> list[float]: """Parse the response from text embedding for Cohere Command models.""" embeddings = response.get("embeddings") 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) + # Pick the first present type in a fixed priority order, then validate that it holds a non-empty vector. + selected = next((embeddings[t] for t in _EMBEDDING_TYPE_PRIORITY if t in embeddings), None) + if not isinstance(selected, list) or len(selected) == 0 or not _is_numeric_vector(selected[0]): + raise ServiceInvalidResponseError("The response from Cohere model does not contain valid embeddings.") + return selected[0] # type: ignore if not isinstance(embeddings, list) or len(embeddings) == 0: raise ServiceInvalidResponseError("The response from Cohere model does not contain embeddings.") diff --git a/python/tests/unit/connectors/ai/bedrock/services/test_bedrock_text_embedding_generation.py b/python/tests/unit/connectors/ai/bedrock/services/test_bedrock_text_embedding_generation.py index fd51fda30608..b021d1672f6f 100644 --- a/python/tests/unit/connectors/ai/bedrock/services/test_bedrock_text_embedding_generation.py +++ b/python/tests/unit/connectors/ai/bedrock/services/test_bedrock_text_embedding_generation.py @@ -242,6 +242,9 @@ async def test_bedrock_text_embedding_with_invalid_response( ({"float": [[0.1, 0.2, 0.3]]}, [0.1, 0.2, 0.3]), ({"int8": [[1, 2, 3]], "float": [[0.1, 0.2, 0.3]]}, [0.1, 0.2, 0.3]), ({"int8": [[1, 2, 3]]}, [1, 2, 3]), + # Selection follows a fixed priority, not the JSON key order. + ({"binary": [[1]], "int8": [[1, 2]]}, [1, 2]), + ({"int8": [[1, 2]], "binary": [[1]]}, [1, 2]), ([[0.1, 0.2, 0.3]], [0.1, 0.2, 0.3]), ], ) @@ -264,7 +267,19 @@ async def test_bedrock_cohere_text_embedding_with_embedding_types(embeddings, ex assert response.tolist() == [expected] -@pytest.mark.parametrize("embeddings", [{}, {"float": []}, {"float": 0.1}]) +@pytest.mark.parametrize( + "embeddings", + [ + {}, + {"float": []}, + {"float": 0.1}, + {"float": [0.1, 0.2]}, + {"float": [[]]}, + {"int8": [[1, 2]], "float": []}, + {"float": [["a", "b"]]}, + {"unknown": [[1, 2]]}, + ], +) async def test_bedrock_cohere_text_embedding_with_invalid_typed_response(embeddings) -> None: body = Mock() body.read.return_value = json.dumps({"embeddings": embeddings})