From 8db3805b5b6e6fc425325bc10b36aaeb6015ba1d Mon Sep 17 00:00:00 2001 From: Syed Ahmed Mubasiruddin Date: Thu, 27 Aug 2026 14:27:11 +0000 Subject: [PATCH] fix: truncate router encoder inputs --- .../auto_router/litellm_encoder.py | 22 +++++++++++++++---- .../router_strategy/test_litellm_encoder.py | 21 ++++++++++++++++++ 2 files changed, 39 insertions(+), 4 deletions(-) diff --git a/litellm/router_strategy/auto_router/litellm_encoder.py b/litellm/router_strategy/auto_router/litellm_encoder.py index 1b34785b6fe..2c757b5216a 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,22 @@ 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) + embeddings: list[list[float]] = [] + + for i in range(0, len(docs), EMBEDDING_BATCH_SIZE): + batch = docs[i : i + EMBEDDING_BATCH_SIZE] + + embeds: Final = await self.litellm_router_instance.aembedding( + input=self._clamp(batch), + model=self.model_name, + **kwargs, + ) + + embeddings.extend(litellm_to_list(embeds)) + + return embeddings + except Exception as e: raise ValueError(f"{self.type.capitalize()} API call failed. Error: {e}") from e diff --git a/tests/test_litellm/router_strategy/test_litellm_encoder.py b/tests/test_litellm/router_strategy/test_litellm_encoder.py index ebd6efe309c..03e2bca30c1 100644 --- a/tests/test_litellm/router_strategy/test_litellm_encoder.py +++ b/tests/test_litellm/router_strategy/test_litellm_encoder.py @@ -108,6 +108,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]