diff --git a/litellm/router_strategy/auto_router/litellm_encoder.py b/litellm/router_strategy/auto_router/litellm_encoder.py index 1b34785b6fe..81c9dec6ff5 100644 --- a/litellm/router_strategy/auto_router/litellm_encoder.py +++ b/litellm/router_strategy/auto_router/litellm_encoder.py @@ -7,6 +7,8 @@ from semantic_router.encoders.base import AsymmetricDenseMixin import litellm from litellm._logging import verbose_router_logger +EMBEDDING_BATCH_SIZE = 512 + if TYPE_CHECKING: from litellm.router import Router else: @@ -144,10 +146,21 @@ class LiteLLMRouterEncoder(CustomDenseEncoder, AsymmetricDenseMixin): async def aencode_documents(self, docs: list[str], **kwargs) -> list[list[float]]: if self.litellm_router_instance is None: raise ValueError("litellm_router_instance is not set") - try: - embeds: Final = await self.litellm_router_instance.aembedding( - input=self._clamp(docs), model=self.model_name, **kwargs - ) - return litellm_to_list(embeds) - except Exception as e: - raise ValueError(f"{self.type.capitalize()} API call failed. Error: {e}") from e + + embeddings: list[list[float]] = [] + + for i in range(0, len(docs), EMBEDDING_BATCH_SIZE): + batch = docs[i : i + EMBEDDING_BATCH_SIZE] + + try: + embeds: Final = await self.litellm_router_instance.aembedding( + input=self._clamp(batch), + model=self.model_name, + **kwargs, + ) + except Exception as e: + raise ValueError(f"{self.type.capitalize()} API call failed. Error: {e}") from e + + embeddings.extend(litellm_to_list(embeds)) + + return embeddings diff --git a/tests/test_litellm/router_strategy/test_litellm_encoder.py b/tests/test_litellm/router_strategy/test_litellm_encoder.py index ebd6efe309c..2ae6105fb6f 100644 --- a/tests/test_litellm/router_strategy/test_litellm_encoder.py +++ b/tests/test_litellm/router_strategy/test_litellm_encoder.py @@ -4,7 +4,6 @@ from typing import Any, Final import pytest - import litellm from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder @@ -108,6 +107,27 @@ class TestEmbeddingInputCap: assert router.embedded_inputs == [["y" * 100]] + @pytest.mark.asyncio + async def test_should_batch_large_document_embeddings(self): + router: Final = RecordingRouter() + docs: Final = [f"tool-{i}" for i in range(513)] + + embeddings = await _encoder(router).aencode_documents(docs) + + assert [len(batch) for batch in router.embedded_inputs] == [512, 1] + assert [doc for batch in router.embedded_inputs for doc in batch] == docs + assert len(embeddings) == len(docs) + + @pytest.mark.asyncio + async def test_should_not_create_extra_embedding_batch_at_limit(self): + router: Final = RecordingRouter() + docs: Final = [f"tool-{i}" for i in range(512)] + + embeddings = await _encoder(router).aencode_documents(docs) + + assert [len(batch) for batch in router.embedded_inputs] == [512] + assert len(embeddings) == len(docs) + def test_should_leave_docs_within_the_cap_untouched(self): router: Final = RecordingRouter() docs: Final = ["a short prompt", "b" * 100]