From 3d71e846b913ddbedd974c530c92609f698d1cf0 Mon Sep 17 00:00:00 2001 From: Aasif-Simpplr Date: Mon, 5 Oct 2026 15:29:05 +0530 Subject: [PATCH 1/2] fix(embedding): return one mock embedding per input mock_response on embedding calls always returned a single vector at index 0, however many inputs were sent, so batch callers got back fewer rows than they asked for. The mock now returns the vector once per input, with matching indexes, for string, string list and token inputs --- litellm/litellm_core_utils/mock_functions.py | 21 +++++++++++++++-- litellm/main.py | 2 +- .../litellm_core_utils/test_mock_functions.py | 23 +++++++++++++++++++ 3 files changed, 43 insertions(+), 3 deletions(-) create mode 100644 tests/unit/litellm_core_utils/test_mock_functions.py diff --git a/litellm/litellm_core_utils/mock_functions.py b/litellm/litellm_core_utils/mock_functions.py index ffbe5f72357..94d0ef8e423 100644 --- a/litellm/litellm_core_utils/mock_functions.py +++ b/litellm/litellm_core_utils/mock_functions.py @@ -1,3 +1,5 @@ +from collections.abc import Sequence + from ..types.utils import ( Embedding, EmbeddingResponse, @@ -7,14 +9,29 @@ from ..types.utils import ( ) -def mock_embedding(model: str, mock_response: list[float] | None): +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) + + +def mock_embedding( + model: str, + embedding_input: str | Sequence[str] | Sequence[int] | Sequence[Sequence[int]], + mock_response: list[float] | None, +): if mock_response is None: mock_response = [0.0] * 1536 elif mock_response == "error": raise Exception("Mock error") return EmbeddingResponse( model=model, - data=[Embedding(embedding=mock_response, index=0, object="embedding")], + data=[ + Embedding(embedding=mock_response, index=index, object="embedding") + for index in range(_embedding_input_count(embedding_input)) + ], usage=Usage(prompt_tokens=10, completion_tokens=0), ) diff --git a/litellm/main.py b/litellm/main.py index 122de9a02c0..05a6bd9ca02 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6452,7 +6452,7 @@ def embedding( ) if mock_response is not None: - return mock_embedding(model=model, mock_response=mock_response) + return mock_embedding(model=model, embedding_input=input, mock_response=mock_response) try: response: EmbeddingResponse | Coroutine[object, object, EmbeddingResponse] | None = None diff --git a/tests/unit/litellm_core_utils/test_mock_functions.py b/tests/unit/litellm_core_utils/test_mock_functions.py new file mode 100644 index 00000000000..841336f9917 --- /dev/null +++ b/tests/unit/litellm_core_utils/test_mock_functions.py @@ -0,0 +1,23 @@ +from typing import Final + +import pytest + +import litellm + + +@pytest.mark.parametrize( + ("embedding_input", "expected_count"), + [ + ("hello", 1), + (["first", "second", "third"], 3), + ([101, 102, 103], 1), + ([[101, 102], [103, 104]], 2), + ], +) +def test_mock_embedding_returns_one_embedding_per_input(embedding_input: object, expected_count: int) -> None: + vector: Final = [0.1, 0.2] + + response: Final = litellm.embedding(model="openai/test-model", input=embedding_input, mock_response=vector) + + assert [row["index"] for row in response.data] == list(range(expected_count)) + assert [row["embedding"] for row in response.data] == [vector] * expected_count From 44a2e74123d33df774c3f9a58c2a3e73c0c7a0ce Mon Sep 17 00:00:00 2001 From: Aasif-Simpplr Date: Mon, 5 Oct 2026 15:50:47 +0530 Subject: [PATCH 2/2] fix(embedding): reject mock input counts a provider would reject An empty input list or more than 2048 inputs now raises BadRequestError instead of returning zero rows or one row per item without limit, matching OpenAI's documented bounds --- litellm/litellm_core_utils/mock_functions.py | 27 ++++++++++++------- .../litellm_core_utils/test_mock_functions.py | 7 +++++ 2 files changed, 25 insertions(+), 9 deletions(-) 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])