diff --git a/python/packages/openai/agent_framework_openai/_embedding_client.py b/python/packages/openai/agent_framework_openai/_embedding_client.py index c3ab22bdd13..8c4d7bb1d55 100644 --- a/python/packages/openai/agent_framework_openai/_embedding_client.py +++ b/python/packages/openai/agent_framework_openai/_embedding_client.py @@ -296,6 +296,13 @@ async def get_embeddings( if encoding == "base64" and isinstance(item.embedding, str): # Decode base64-encoded floats (little-endian IEEE 754) raw = base64.b64decode(item.embedding) + if len(raw) % 4: + # A misaligned buffer would otherwise fail inside struct.unpack with a + # bare buffer-size error and no hint about which response item broke. + raise ValueError( + f"Base64 encoded embedding for input index {item.index} decoded to {len(raw)} bytes, " + "which cannot be split into 4-byte little-endian floats. The response payload is malformed." + ) vector = list(struct.unpack(f"<{len(raw) // 4}f", raw)) else: vector = item.embedding diff --git a/python/packages/openai/tests/openai/test_openai_embedding_client.py b/python/packages/openai/tests/openai/test_openai_embedding_client.py index ba57b8ee754..454559988da 100644 --- a/python/packages/openai/tests/openai/test_openai_embedding_client.py +++ b/python/packages/openai/tests/openai/test_openai_embedding_client.py @@ -233,6 +233,47 @@ async def test_openai_base64_decoding(openai_unit_test_env: dict[str, str]) -> N assert abs(expected - actual) < 1e-6 +async def test_openai_base64_misaligned_payload_raises_descriptive_error(openai_unit_test_env: dict[str, str]) -> None: + """A truncated base64 payload must fail with a clear error naming the input index, + not a bare struct buffer-size error.""" + import base64 + import struct + + # Input 0 is well formed; input 1 is one float (4 bytes) plus a stray byte, so it + # decodes to 5 bytes. The response lists them in reverse order to make sure the + # diagnostic reports the item's own index rather than its position in the list. + aligned = base64.b64encode(struct.pack("<1f", 0.25)).decode() + misaligned = base64.b64encode(struct.pack("<1f", 0.5) + b"\x00").decode() + + good_item = MagicMock() + good_item.embedding = aligned + good_item.index = 0 + + bad_item = MagicMock() + bad_item.embedding = misaligned + bad_item.index = 1 + + mock_response = MagicMock() + mock_response.data = [bad_item, good_item] + mock_response.model = "text-embedding-3-small" + mock_response.usage = MagicMock(prompt_tokens=6, total_tokens=6) + + client = OpenAIEmbeddingClient() + client.client = MagicMock() + client.client.embeddings = MagicMock() + client.client.embeddings.create = AsyncMock(return_value=mock_response) + + options: OpenAIEmbeddingOptions = {"encoding_format": "base64"} + + with pytest.raises(ValueError) as exc_info: + await client.get_embeddings(["first", "second"], options=options) + + message = str(exc_info.value) + assert "input index 1" in message + assert "decoded to 5 bytes" in message + assert "cannot be split into 4-byte little-endian floats" in message + + async def test_openai_error_when_no_model() -> None: client = cast(Any, object.__new__(OpenAIEmbeddingClient)) client.model = None