diff --git a/python/packages/foundry/agent_framework_foundry/_embedding_client.py b/python/packages/foundry/agent_framework_foundry/_embedding_client.py index d9f5e3fb8f..8050198ab3 100644 --- a/python/packages/foundry/agent_framework_foundry/_embedding_client.py +++ b/python/packages/foundry/agent_framework_foundry/_embedding_client.py @@ -22,6 +22,7 @@ load_settings, ) from agent_framework._telemetry import IS_TELEMETRY_ENABLED, USER_AGENT_KEY, get_user_agent, mark_feature_used +from agent_framework.exceptions import IntegrationInvalidResponseException from agent_framework.observability import EmbeddingTelemetryLayer from azure.ai.inference.aio import EmbeddingsClient, ImageEmbeddingsClient from azure.ai.inference.models import ImageEmbeddingInput @@ -428,7 +429,22 @@ async def get_embeddings( for item in sorted(openai_response.data, key=lambda value: value.index): original_idx = text_items[item.index][0] if encoding == "base64" and isinstance(item.embedding, str): - raw = base64.b64decode(item.embedding) + # base64.b64decode raises binascii.Error for bad padding/alphabet and + # UnicodeEncodeError for non-ASCII input; both are ValueError subclasses. + try: + raw = base64.b64decode(item.embedding) + except ValueError as exc: + raise IntegrationInvalidResponseException( + f"The base64 embedding returned for input index {original_idx} could not be " + f"decoded: {exc}" + ) from exc + if len(raw) % 4 != 0: + # struct.unpack would fail with a bare "unpack requires a buffer of N bytes", + # which says nothing about which input broke or why. + raise IntegrationInvalidResponseException( + f"The base64 embedding returned for input index {original_idx} decoded to " + f"{len(raw)} bytes, which is not a multiple of the 4 bytes a float32 value occupies." + ) vector = list(struct.unpack(f"<{len(raw) // 4}f", raw)) else: vector = [float(value) for value in item.embedding] diff --git a/python/packages/foundry/tests/foundry/test_foundry_embedding_client.py b/python/packages/foundry/tests/foundry/test_foundry_embedding_client.py index e56507488b..91a0e994bf 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_embedding_client.py +++ b/python/packages/foundry/tests/foundry/test_foundry_embedding_client.py @@ -10,6 +10,7 @@ import pytest from agent_framework import Content, SecretString from agent_framework._telemetry import USER_AGENT_KEY, get_user_agent +from agent_framework.exceptions import IntegrationInvalidResponseException from azure.core.credentials import AzureKeyCredential from azure.identity.aio import AzureCliCredential @@ -266,6 +267,60 @@ async def test_project_client_options_and_response_order(self) -> None: ) assert [embedding.vector for embedding in result] == [[0.1], [0.3]] + async def test_project_client_base64_rejects_length_that_is_not_a_multiple_of_four(self) -> None: + """A truncated base64 payload must name the input, not escape as a bare struct.error.""" + import base64 + import struct + + # One float plus a stray byte: 5 bytes, so no whole number of float32 values fits. + b64_str = base64.b64encode(struct.pack("<1f", 0.5) + b"\x00").decode() + + item = MagicMock() + item.embedding = b64_str + item.index = 0 + response = MagicMock() + response.data = [item] + response.model = "test-model" + response.usage = None + + openai_client = MagicMock() + openai_client.base_url = "https://test.services.ai.azure.com/api/projects/test/openai/v1/" + openai_client.embeddings.create = AsyncMock(return_value=response) + openai_client.close = AsyncMock() + project_client = MagicMock() + project_client.get_openai_client.return_value = openai_client + + client = RawFoundryEmbeddingClient(model="text-embedding-3-small", project_client=project_client) + + with pytest.raises(IntegrationInvalidResponseException, match="input index 0") as exc_info: + await client.get_embeddings(["hello"], options={"encoding_format": "base64"}) + + assert "5 bytes" in str(exc_info.value) + + async def test_project_client_base64_rejects_undecodable_payload(self) -> None: + """A payload that cannot be decoded must name the input rather than raise binascii.Error.""" + item = MagicMock() + item.embedding = "!!!not-base64!!!" + item.index = 0 + response = MagicMock() + response.data = [item] + response.model = "test-model" + response.usage = None + + openai_client = MagicMock() + openai_client.base_url = "https://test.services.ai.azure.com/api/projects/test/openai/v1/" + openai_client.embeddings.create = AsyncMock(return_value=response) + openai_client.close = AsyncMock() + project_client = MagicMock() + project_client.get_openai_client.return_value = openai_client + + client = RawFoundryEmbeddingClient(model="text-embedding-3-small", project_client=project_client) + + with pytest.raises(IntegrationInvalidResponseException, match="input index 0") as exc_info: + await client.get_embeddings(["hello"], options={"encoding_format": "base64"}) + + assert exc_info.value.__cause__ is not None + async def test_project_mode_rejects_images_before_sending_text(self) -> None: """Project OpenAI embedding deployments reject image inputs without partial requests.""" openai_client = _make_openai_client() diff --git a/python/packages/openai/agent_framework_openai/_embedding_client.py b/python/packages/openai/agent_framework_openai/_embedding_client.py index 2c06e49615..9560f99cc9 100644 --- a/python/packages/openai/agent_framework_openai/_embedding_client.py +++ b/python/packages/openai/agent_framework_openai/_embedding_client.py @@ -12,6 +12,7 @@ from agent_framework._settings import SecretString from agent_framework._telemetry import USER_AGENT_KEY, mark_feature_used from agent_framework._types import Embedding, EmbeddingGenerationOptions, GeneratedEmbeddings, UsageDetails +from agent_framework.exceptions import IntegrationInvalidResponseException from agent_framework.observability import EmbeddingTelemetryLayer from openai import AsyncAzureOpenAI, AsyncOpenAI @@ -311,8 +312,15 @@ async def get_embeddings( for item in sorted(response.data, key=lambda item: item.index): vector: list[float] if encoding == "base64" and isinstance(item.embedding, str): - # Decode base64-encoded floats (little-endian IEEE 754) - raw = base64.b64decode(item.embedding) + # Decode base64-encoded floats (little-endian IEEE 754). + # base64.b64decode raises binascii.Error for bad padding/alphabet and + # UnicodeEncodeError for non-ASCII input; both are ValueError subclasses. + try: + raw = base64.b64decode(item.embedding) + except ValueError as exc: + raise IntegrationInvalidResponseException( + f"The base64 embedding returned for input index {item.index} could not be decoded: {exc}" + ) from exc 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. 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 4692150390..b9c8913a50 100644 --- a/python/packages/openai/tests/openai/test_openai_embedding_client.py +++ b/python/packages/openai/tests/openai/test_openai_embedding_client.py @@ -12,6 +12,7 @@ from agent_framework import SecretString, SupportsGetEmbeddings from agent_framework._telemetry import get_feature_token from agent_framework.exceptions import SettingNotFoundError +from agent_framework.exceptions import IntegrationInvalidResponseException from openai.types import CreateEmbeddingResponse from openai.types import Embedding as OpenAIEmbedding from openai.types.create_embedding_response import Usage @@ -322,6 +323,38 @@ async def test_openai_base64_misaligned_payload_raises_descriptive_error(openai_ assert "cannot be split into 4-byte little-endian floats" in message +@pytest.mark.parametrize( + "payload", + [ + pytest.param("!!!not-base64!!!", id="invalid-characters"), + pytest.param("\u00e9", id="non-ascii"), + ], +) +async def test_openai_base64_decoding_rejects_undecodable_payload( + payload: str, + openai_unit_test_env: dict[str, str], +) -> None: + """A payload that cannot be decoded must name the input rather than raise binascii.Error.""" + mock_item = MagicMock() + mock_item.embedding = payload + mock_item.index = 1 + + mock_response = MagicMock() + mock_response.data = [mock_item] + mock_response.model = "text-embedding-3-small" + mock_response.usage = MagicMock(prompt_tokens=3, total_tokens=3) + + client = OpenAIEmbeddingClient() + client.client = MagicMock() + client.client.embeddings = MagicMock() + client.client.embeddings.create = AsyncMock(return_value=mock_response) + + with pytest.raises(IntegrationInvalidResponseException, match="input index 1") as exc_info: + await client.get_embeddings(["test"], options={"encoding_format": "base64"}) + + assert exc_info.value.__cause__ is not None + + async def test_openai_error_when_no_model() -> None: client = cast(Any, object.__new__(OpenAIEmbeddingClient)) client.model = None