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

@@ -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),
     )
 

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
 

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