diff --git a/litellm/litellm_core_utils/mock_functions.py b/litellm/litellm_core_utils/mock_functions.py index 94d0ef8e423..4d33008d692 100644 --- a/litellm/litellm_core_utils/mock_functions.py +++ b/litellm/litellm_core_utils/mock_functions.py @@ -1,5 +1,7 @@ from collections.abc import Sequence +from typing import Final +from ..exceptions import BadRequestError from ..types.utils import ( Embedding, EmbeddingResponse, @@ -8,13 +10,15 @@ from ..types.utils import ( Usage, ) +# OpenAI caps an embedding input array at 2048 items, as of 2026-10-05: +# https://platform.openai.com/docs/api-reference/embeddings/create +_MAX_EMBEDDING_INPUTS: Final = 2048 + def _embedding_input_count(embedding_input: str | Sequence[str] | Sequence[int] | Sequence[Sequence[int]]) -> int: - match embedding_input: - case str() | [int(), *_]: - return 1 - case _: - return len(embedding_input) + if isinstance(embedding_input, str) or (len(embedding_input) > 0 and isinstance(embedding_input[0], int)): + return 1 + return len(embedding_input) def mock_embedding( @@ -26,12 +30,17 @@ def mock_embedding( mock_response = [0.0] * 1536 elif mock_response == "error": raise Exception("Mock error") + input_count: Final = _embedding_input_count(embedding_input) + if not 1 <= input_count <= _MAX_EMBEDDING_INPUTS: + raise BadRequestError( + message=f"input must contain between 1 and {_MAX_EMBEDDING_INPUTS} items, got {input_count}", + model=model, + llm_provider=None, + body={"param": "input"}, + ) return EmbeddingResponse( model=model, - data=[ - Embedding(embedding=mock_response, index=index, object="embedding") - for index in range(_embedding_input_count(embedding_input)) - ], + data=[Embedding(embedding=mock_response, index=index, object="embedding") for index in range(input_count)], usage=Usage(prompt_tokens=10, completion_tokens=0), ) diff --git a/tests/unit/litellm_core_utils/test_mock_functions.py b/tests/unit/litellm_core_utils/test_mock_functions.py index 841336f9917..e362ae98055 100644 --- a/tests/unit/litellm_core_utils/test_mock_functions.py +++ b/tests/unit/litellm_core_utils/test_mock_functions.py @@ -3,6 +3,7 @@ from typing import Final import pytest import litellm +from litellm.exceptions import BadRequestError @pytest.mark.parametrize( @@ -21,3 +22,9 @@ def test_mock_embedding_returns_one_embedding_per_input(embedding_input: object, assert [row["index"] for row in response.data] == list(range(expected_count)) assert [row["embedding"] for row in response.data] == [vector] * expected_count + + +@pytest.mark.parametrize("embedding_input", [[], ["text"] * 2049]) +def test_mock_embedding_rejects_input_count_a_provider_would_reject(embedding_input: list[str]) -> None: + with pytest.raises(BadRequestError, match="between 1 and 2048 items"): + litellm.embedding(model="openai/test-model", input=embedding_input, mock_response=[0.1, 0.2])