From 9d455fb8571b09a0cf24bb270c3e45ed90956f55 Mon Sep 17 00:00:00 2001 From: Divyam Talwar Date: Thu, 17 Sep 2026 09:00:41 +0530 Subject: [PATCH] fix: reject malformed embedding indices --- .../embedding/providers/openai_compatible.py | 26 +++++- .../test_embedding_response_validation.py | 88 +++++++++++++++++++ 2 files changed, 110 insertions(+), 4 deletions(-) create mode 100644 datamind/tests/test_embedding_response_validation.py diff --git a/datamind/capabilities/embedding/providers/openai_compatible.py b/datamind/capabilities/embedding/providers/openai_compatible.py index cb87271..42a91a0 100644 --- a/datamind/capabilities/embedding/providers/openai_compatible.py +++ b/datamind/capabilities/embedding/providers/openai_compatible.py @@ -147,6 +147,10 @@ async def _call(self, inputs: list[str]) -> list[list[float]]: ) resp.raise_for_status() body = resp.json() + if not isinstance(body, dict): + raise ExternalServiceError( + "embedding", "response body must be an object", + ) data = body.get("data") or [] if not data: raise ExternalServiceError( @@ -159,9 +163,21 @@ async def _call(self, inputs: list[str]) -> list[list[float]]: f"response count mismatch: expected {len(inputs)}, got {len(data)}", ) try: - ordered = sorted(data, key=lambda row: int(row["index"])) - indices = [int(row["index"]) for row in ordered] - vecs = [list(row["embedding"]) for row in ordered] + indexed_vectors = [] + for row in data: + if not isinstance(row, dict): + raise TypeError("each response row must be an object") + index = row["index"] + # JSON booleans are distinct from integer indices even + # though bool is an int subclass in Python. Do not + # coerce floats or strings: truncation can turn a + # malformed response into a seemingly complete batch. + if isinstance(index, bool) or not isinstance(index, int): + raise ValueError("response index must be a JSON integer") + indexed_vectors.append((index, list(row["embedding"]))) + ordered = sorted(indexed_vectors, key=lambda item: item[0]) + indices = [index for index, _vector in ordered] + vecs = [vector for _index, vector in ordered] except (KeyError, TypeError, ValueError) as exc: raise ExternalServiceError( "embedding", "response rows require integer index and embedding", @@ -183,7 +199,9 @@ async def _call(self, inputs: list[str]) -> list[list[float]]: f"dimension mismatch for {self._model}: expected {self.dimension}, got {actual_dimension}", ) if any( - not isinstance(value, (int, float)) or not math.isfinite(float(value)) + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(float(value)) for vec in vecs for value in vec ): raise ExternalServiceError("embedding", "response contains non-finite vector values") diff --git a/datamind/tests/test_embedding_response_validation.py b/datamind/tests/test_embedding_response_validation.py new file mode 100644 index 0000000..66e640b --- /dev/null +++ b/datamind/tests/test_embedding_response_validation.py @@ -0,0 +1,88 @@ +"""Strict validation for OpenAI-compatible embedding response rows.""" +from __future__ import annotations + +import json + +import httpx +import pytest + +from datamind.capabilities.embedding.providers.openai_compatible import ( + OpenAICompatibleEmbedding, +) +from datamind.core.errors import ExternalServiceError + + +def _client(body, calls): + def handler(request: httpx.Request) -> httpx.Response: + calls.append(request) + return httpx.Response( + 200, + content=json.dumps(body).encode("utf-8"), + headers={"content-type": "application/json"}, + ) + + return httpx.AsyncClient(transport=httpx.MockTransport(handler)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "indices", + [ + [0.9, 1.9], # fractional values must not be truncated + [True, 1], # bool is an int subclass in Python, not a JSON index + ["0", 1], # numeric strings are not the documented JSON shape + [-1, 0], + [0, 0], + [0, 2], + ], +) +async def test_rejects_malformed_response_indices_without_retry(indices): + calls = [] + body = { + "data": [ + {"index": indices[0], "embedding": [1.0, 2.0]}, + {"index": indices[1], "embedding": [3.0, 4.0]}, + ], + } + client = _client(body, calls) + embedding = OpenAICompatibleEmbedding( + api_key="test", model="custom", dimension=2, + client=client, max_retries=5, + ) + with pytest.raises(ExternalServiceError): + await embedding.embed_texts(["a", "b"]) + assert len(calls) == 1 + await embedding.aclose() + + +@pytest.mark.asyncio +async def test_rejects_boolean_vector_coordinates(): + calls = [] + client = _client( + {"data": [{"index": 0, "embedding": [True, 0.0]}]}, calls, + ) + embedding = OpenAICompatibleEmbedding( + api_key="test", model="custom", dimension=2, + client=client, max_retries=0, + ) + with pytest.raises(ExternalServiceError): + await embedding.embed_texts(["a"]) + await embedding.aclose() + + +@pytest.mark.asyncio +async def test_reorders_valid_integer_indices(): + calls = [] + client = _client( + {"data": [ + {"index": 1, "embedding": [3.0, 4.0]}, + {"index": 0, "embedding": [1.0, 2.0]}, + ]}, + calls, + ) + embedding = OpenAICompatibleEmbedding( + api_key="test", model="custom", dimension=2, + client=client, + ) + assert await embedding.embed_texts(["a", "b"]) == [[1.0, 2.0], [3.0, 4.0]] + await embedding.aclose()