feat: add router option to LiteLLMEmbedder (#457)

Author: mhordynski

packages/ragbits-core/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Add router option to LiteLLMEmbedder (#440)
+
 ## 0.12.0 (2025-03-25)
 - Allow Prompt class to accept the asynchronous response_parser. Change the signature of parse_response method.
 - Fix from_config for LiteLLM class (#441)

packages/ragbits-core/src/ragbits/core/embeddings/litellm.py

@@ -1,4 +1,7 @@
+from typing import Any
+
 import litellm
+from typing_extensions import Self
 
 from ragbits.core.audit import trace
 from ragbits.core.embeddings import Embedder
@@ -38,6 +41,7 @@ def __init__(
         api_base: str | None = None,
         api_key: str | None = None,
         api_version: str | None = None,
+        router: litellm.Router | None = None,
     ) -> None:
         """
         Constructs the LiteLLMEmbeddingClient.
@@ -51,12 +55,14 @@ def __init__(
                 for more information, follow the instructions for your specific vendor in the\
                 [LiteLLM documentation](https://docs.litellm.ai/docs/embedding/supported_embedding).
             api_version: The API version for the call.
+            router: Router to be used to [route requests](https://docs.litellm.ai/docs/routing) to different models.
         """
         super().__init__(default_options=default_options)
         self.model = model
         self.api_base = api_base
         self.api_key = api_key
         self.api_version = api_version
+        self.router = router
 
     async def embed_text(self, data: list[str], options: LiteLLMEmbedderOptions | None = None) -> list[list[float]]:
         """
@@ -85,7 +91,8 @@ async def embed_text(self, data: list[str], options: LiteLLMEmbedderOptions | No
             options=merged_options.dict(),
         ) as outputs:
             try:
-                response = await litellm.aembedding(
+                entrypoint = self.router or litellm
+                response = await entrypoint.aembedding(
                     input=data,
                     model=self.model,
                     api_base=self.api_base,
@@ -110,3 +117,19 @@ async def embed_text(self, data: list[str], options: LiteLLMEmbedderOptions | No
                 outputs.total_tokens = response.usage.total_tokens
 
         return outputs.embeddings
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> Self:
+        """
+        Creates and returns a LiteLLMEmbedder instance.
+
+        Args:
+            config: A configuration object containing the configuration for initializing the LiteLLMEmbedder instance.
+
+        Returns:
+            LiteLLMEmbedder: An initialized LiteLLMEmbedder instance.
+        """
+        if "router" in config:
+            router = litellm.router.Router(model_list=config["router"])
+            config["router"] = router
+        return super().from_config(config)

packages/ragbits-core/tests/unit/embeddings/test_from_config.py

@@ -1,3 +1,5 @@
+import litellm
+
 from ragbits.core.embeddings import Embedder, NoopEmbedder
 from ragbits.core.embeddings.litellm import LiteLLMEmbedder, LiteLLMEmbedderOptions
 from ragbits.core.embeddings.sparse import BagOfTokens, BagOfTokensOptions, SparseEmbedder
@@ -58,3 +60,42 @@ def test_subclass_from_config_bag_of_tokens():
         option1="value1",
         option2="value2",
     )  # type: ignore
+
+
+def test_from_config_with_router():
+    config = ObjectConstructionConfig(
+        type="ragbits.core.embeddings.litellm:LiteLLMEmbedder",
+        config={
+            "model": "text-embedding-3-small",
+            "api_key": "test_api_key",
+            "router": [
+                {
+                    "model_name": "small",
+                    "litellm_params": {
+                        "model": "text-embedding-3-small",
+                        "dimensions": 3000,
+                        "api_key": "test_api_key",
+                    },
+                },
+                {
+                    "model_name": "large",
+                    "litellm_params": {
+                        "model": "text-embedding-3-large",
+                        "api_key": "test_api_key",
+                    },
+                },
+            ],
+        },
+    )
+
+    embedder: Embedder = Embedder.subclass_from_config(config)
+    assert isinstance(embedder, LiteLLMEmbedder)
+    assert embedder.api_base is None
+    assert embedder.model == "text-embedding-3-small"
+    assert embedder.api_key == "test_api_key"
+    assert isinstance(embedder.router, litellm.router.Router)
+    assert len(embedder.router.model_list) == 2
+    assert embedder.router.model_list[0]["model_name"] == "small"
+    assert embedder.router.model_list[0]["litellm_params"]["dimensions"] == 3000
+    assert embedder.router.model_list[1]["model_name"] == "large"
+    assert embedder.router.model_list[1]["litellm_params"]["model"] == "text-embedding-3-large"