diff --git a/litellm/litellm_core_utils/mock_functions.py b/litellm/litellm_core_utils/mock_functions.py index ffbe5f72357..4d33008d692 100644 --- a/litellm/litellm_core_utils/mock_functions.py +++ b/litellm/litellm_core_utils/mock_functions.py @@ -1,3 +1,7 @@ +from collections.abc import Sequence +from typing import Final + +from ..exceptions import BadRequestError from ..types.utils import ( Embedding, EmbeddingResponse, @@ -6,15 +10,37 @@ 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 mock_embedding(model: str, mock_response: list[float] | None): + +def _embedding_input_count(embedding_input: str | Sequence[str] | Sequence[int] | Sequence[Sequence[int]]) -> int: + 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( + 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") + 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=0, object="embedding")], + 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/litellm/main.py b/litellm/main.py index 794b3cef729..2c67dc5d560 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6971,7 +6971,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..e362ae98055 --- /dev/null +++ b/tests/unit/litellm_core_utils/test_mock_functions.py @@ -0,0 +1,30 @@ +from typing import Final + +import pytest + +import litellm +from litellm.exceptions import BadRequestError + + +@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 + + +@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])