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..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,12 +90,34 @@ 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.""" - 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": [[...]]}). + # 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.") - 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..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 @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. +import json from unittest.mock import ANY, Mock, patch import boto3 @@ -233,3 +234,62 @@ 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]), + # 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]), + ], +) +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}, + {"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}) + + 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"])