Merge 8db3805b5b into bd7e9c1997

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

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]