From 5ce528e3d39a532792e79ee2fb43cafc56277878 Mon Sep 17 00:00:00 2001 From: Syed Ahmed Mubasiruddin Date: Fri, 28 Aug 2026 14:41:48 +0000 Subject: [PATCH 1/4] fix: batch semantic router embeddings --- .../auto_router/litellm_encoder.py | 27 ++++++++++++++----- .../router_strategy/test_litellm_encoder.py | 22 ++++++++++++++- 2 files changed, 41 insertions(+), 8 deletions(-) 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] From ae12d40eab33976230e466c2383abab2dd5efec3 Mon Sep 17 00:00:00 2001 From: Syed Ahmed Mubasiruddin Date: Fri, 28 Aug 2026 14:58:49 +0000 Subject: [PATCH 2/4] fix: allow mutable embedding accumulator --- litellm/router_strategy/auto_router/litellm_encoder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/router_strategy/auto_router/litellm_encoder.py b/litellm/router_strategy/auto_router/litellm_encoder.py index 81c9dec6ff5..71267f7d955 100644 --- a/litellm/router_strategy/auto_router/litellm_encoder.py +++ b/litellm/router_strategy/auto_router/litellm_encoder.py @@ -147,7 +147,7 @@ class LiteLLMRouterEncoder(CustomDenseEncoder, AsymmetricDenseMixin): if self.litellm_router_instance is None: raise ValueError("litellm_router_instance is not set") - embeddings: list[list[float]] = [] + embeddings: list[list[float]] = [] # mutable-ok: used to accumulate embeddings in batches for i in range(0, len(docs), EMBEDDING_BATCH_SIZE): batch = docs[i : i + EMBEDDING_BATCH_SIZE] From 78b8a5f415597d453c6f37fb20fa4e2fdade8933 Mon Sep 17 00:00:00 2001 From: Syed Ahmed Mubasiruddin Date: Mon, 31 Aug 2026 09:42:26 +0000 Subject: [PATCH 3/4] fix: declare batched embedding accumulator final --- litellm/router_strategy/auto_router/litellm_encoder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/router_strategy/auto_router/litellm_encoder.py b/litellm/router_strategy/auto_router/litellm_encoder.py index 71267f7d955..d68754743f0 100644 --- a/litellm/router_strategy/auto_router/litellm_encoder.py +++ b/litellm/router_strategy/auto_router/litellm_encoder.py @@ -147,7 +147,7 @@ class LiteLLMRouterEncoder(CustomDenseEncoder, AsymmetricDenseMixin): if self.litellm_router_instance is None: raise ValueError("litellm_router_instance is not set") - embeddings: list[list[float]] = [] # mutable-ok: used to accumulate embeddings in batches + embeddings: Final[list[list[float]]] = [] # mutable-ok: used to accumulate embeddings in batches for i in range(0, len(docs), EMBEDDING_BATCH_SIZE): batch = docs[i : i + EMBEDDING_BATCH_SIZE] From ebb34775faf5b7801925973440948b7f36dac9b1 Mon Sep 17 00:00:00 2001 From: Syed Ahmed Mubasiruddin Date: Mon, 31 Aug 2026 09:59:43 +0000 Subject: [PATCH 4/4] fix: allow batch embedding result reassignment --- litellm/router_strategy/auto_router/litellm_encoder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/router_strategy/auto_router/litellm_encoder.py b/litellm/router_strategy/auto_router/litellm_encoder.py index d68754743f0..1cb53ceb9ec 100644 --- a/litellm/router_strategy/auto_router/litellm_encoder.py +++ b/litellm/router_strategy/auto_router/litellm_encoder.py @@ -153,7 +153,7 @@ class LiteLLMRouterEncoder(CustomDenseEncoder, AsymmetricDenseMixin): batch = docs[i : i + EMBEDDING_BATCH_SIZE] try: - embeds: Final = await self.litellm_router_instance.aembedding( + embeds = await self.litellm_router_instance.aembedding( # rebind-ok: one result per batch input=self._clamp(batch), model=self.model_name, **kwargs,