fix: from config for LiteLLM (#454)

Author: kzamlynska

packages/ragbits-conversations/src/ragbits/conversations/history/stores/sql.py

@@ -10,7 +10,7 @@
 from ragbits.conversations.history.stores.base import HistoryStore
 from ragbits.core.options import Options
 from ragbits.core.prompt import ChatFormat
-from ragbits.core.utils.config_handling import ObjectContructionConfig
+from ragbits.core.utils.config_handling import ObjectConstructionConfig
 
 
 class _Base(DeclarativeBase):
@@ -190,6 +190,6 @@ def from_config(cls, config: dict) -> Self:
         Returns:
             An instance of the class initialized with the provided configuration.
         """
-        engine_options = ObjectContructionConfig.model_validate(config["sqlalchemy_engine"])
+        engine_options = ObjectConstructionConfig.model_validate(config["sqlalchemy_engine"])
         config["sqlalchemy_engine"] = create_async_engine(engine_options.config["url"])
         return cls(**config)

packages/ragbits-core/CHANGELOG.md

@@ -4,7 +4,7 @@
 
 ## 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)
 - Fix Qdrant vector store serialization (#419)
 
 ## 0.11.0 (2025-03-25)

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

@@ -1,8 +1,10 @@
 from collections.abc import AsyncGenerator
+from typing import Any
 
 import litellm
 from litellm.utils import CustomStreamWrapper, ModelResponse
 from pydantic import BaseModel
+from typing_extensions import Self
 
 from ragbits.core.audit import trace
 from ragbits.core.llms.base import LLM
@@ -239,3 +241,19 @@ def _get_response_format(
             elif json_mode:
                 response_format = {"type": "json_object"}
         return response_format
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> Self:
+        """
+        Creates and returns a LiteLLM instance.
+
+        Args:
+            config: A configuration object containing the configuration for initializing the LiteLLM instance.
+
+        Returns:
+            LiteLLM: An initialized LiteLLM instance.
+        """
+        if "router" in config:
+            router = litellm.router.Router(model_list=config["router"])
+            config["router"] = router
+        return super().from_config(config)

packages/ragbits-core/src/ragbits/core/utils/config_handling.py

@@ -31,7 +31,7 @@ class NoPreferredConfigError(InvalidConfigError):
 def import_by_path(path: str, default_module: ModuleType | None = None) -> Any:  # noqa: ANN401
     """
     Retrieves and returns an object based on the string in the format of "module.submodule:object_name".
-    If the first part is ommited, the default module is used.
+    If the first part is omitted, the default module is used.
 
     Args:
         path: A string representing the path to the object. This can either be a
@@ -63,7 +63,7 @@ def import_by_path(path: str, default_module: ModuleType | None = None) -> Any:
         raise InvalidConfigError(f"{path} not found in module {default_module}") from err
 
 
-class ObjectContructionConfig(BaseModel):
+class ObjectConstructionConfig(BaseModel):
     """
     A model for object construction configuration.
     """
@@ -87,7 +87,7 @@ class WithConstructionConfig(abc.ABC):
     configuration_key: ClassVar[str]
 
     @classmethod
-    def subclass_from_config(cls, config: ObjectContructionConfig) -> Self:
+    def subclass_from_config(cls, config: ObjectConstructionConfig) -> Self:
         """
         Initializes the class with the provided configuration. May return a subclass of the class,
         if requested by the configuration.
@@ -151,7 +151,7 @@ def preferred_subclass(
         if yaml_path_override:
             preferences = get_config_from_yaml(yaml_path_override)
             if type_config := preferences.get(cls.configuration_key):
-                return cls.subclass_from_config(ObjectContructionConfig.model_validate(type_config))
+                return cls.subclass_from_config(ObjectConstructionConfig.model_validate(type_config))
 
         if factory_path_override:
             return cls.subclass_from_factory(factory_path_override)
@@ -160,7 +160,7 @@ def preferred_subclass(
             return cls.subclass_from_factory(preferred_factory)
 
         if preferred_config := config.preferred_instances_config.get(cls.configuration_key):
-            return cls.subclass_from_config(ObjectContructionConfig.model_validate(preferred_config))
+            return cls.subclass_from_config(ObjectConstructionConfig.model_validate(preferred_config))
 
         raise NoPreferredConfigError(f"Could not find preferred factory or configuration for {cls.configuration_key}")
 
@@ -195,7 +195,7 @@ def __init__(self, default_options: OptionsT | None = None) -> None:
         self.default_options: OptionsT = default_options or self.options_cls()
 
     @classmethod
-    def from_config(cls, config: dict[str, Any]) -> ConfigurableComponent:
+    def from_config(cls, config: dict[str, Any]) -> Self:
         """
         Initializes the class with the provided configuration.
 

packages/ragbits-core/src/ragbits/core/vector_stores/base.py

@@ -10,7 +10,7 @@
 from ragbits.core import vector_stores
 from ragbits.core.embeddings.base import Embedder
 from ragbits.core.options import Options
-from ragbits.core.utils.config_handling import ConfigurableComponent, ObjectContructionConfig
+from ragbits.core.utils.config_handling import ConfigurableComponent, ObjectConstructionConfig
 from ragbits.core.utils.pydantic import SerializableBytes
 
 WhereQuery = dict[str, str | int | float | bool]
@@ -196,6 +196,6 @@ def from_config(cls, config: dict) -> Self:
         options = cls.options_cls(**default_options) if default_options else None
 
         embedder_config = config.pop("embedder")
-        embedder: Embedder = Embedder.subclass_from_config(ObjectContructionConfig.model_validate(embedder_config))
+        embedder: Embedder = Embedder.subclass_from_config(ObjectConstructionConfig.model_validate(embedder_config))
 
         return cls(**config, default_options=options, embedder=embedder)

packages/ragbits-core/src/ragbits/core/vector_stores/chroma.py

@@ -8,7 +8,7 @@
 
 from ragbits.core.audit import trace
 from ragbits.core.embeddings.base import Embedder
-from ragbits.core.utils.config_handling import ObjectContructionConfig, import_by_path
+from ragbits.core.utils.config_handling import ObjectConstructionConfig, import_by_path
 from ragbits.core.utils.dict_transformations import flatten_dict, unflatten_dict
 from ragbits.core.vector_stores.base import (
     EmbeddingType,
@@ -78,7 +78,7 @@ def from_config(cls, config: dict) -> Self:
         Returns:
             An instance of the class initialized with the provided configuration.
         """
-        client_options = ObjectContructionConfig.model_validate(config["client"])
+        client_options = ObjectConstructionConfig.model_validate(config["client"])
         client_cls = import_by_path(client_options.type, chromadb)
         config["client"] = client_cls(**client_options.config)
         return super().from_config(config)

packages/ragbits-core/src/ragbits/core/vector_stores/qdrant.py

@@ -11,7 +11,7 @@
 
 from ragbits.core.audit import trace
 from ragbits.core.embeddings.base import Embedder
-from ragbits.core.utils.config_handling import ObjectContructionConfig, import_by_path
+from ragbits.core.utils.config_handling import ObjectConstructionConfig, import_by_path
 from ragbits.core.utils.dict_transformations import flatten_dict
 from ragbits.core.vector_stores.base import (
     EmbeddingType,
@@ -105,7 +105,7 @@ def from_config(cls, config: dict) -> Self:
         Returns:
             An instance of the class initialized with the provided configuration.
         """
-        client_options = ObjectContructionConfig.model_validate(config["client"])
+        client_options = ObjectConstructionConfig.model_validate(config["client"])
         client_cls = import_by_path(client_options.type, qdrant_client)
         if "limits" in client_options.config:
             limits = httpx.Limits(**client_options.config["limits"])

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

@@ -2,11 +2,11 @@
 from ragbits.core.embeddings.litellm import LiteLLMEmbedder, LiteLLMEmbedderOptions
 from ragbits.core.embeddings.sparse import BagOfTokens, BagOfTokensOptions, SparseEmbedder
 from ragbits.core.types import NOT_GIVEN
-from ragbits.core.utils.config_handling import ObjectContructionConfig
+from ragbits.core.utils.config_handling import ObjectConstructionConfig
 
 
 def test_subclass_from_config_litellm():
-    config = ObjectContructionConfig.model_validate(
+    config = ObjectConstructionConfig.model_validate(
         {
             "type": "ragbits.core.embeddings.litellm:LiteLLMEmbedder",
             "config": {
@@ -32,13 +32,13 @@ def test_subclass_from_config_litellm():
 
 
 def test_subclass_from_config_default_path_litellm():
-    config = ObjectContructionConfig.model_validate({"type": "NoopEmbedder"})
+    config = ObjectConstructionConfig.model_validate({"type": "NoopEmbedder"})
     embedder: Embedder = Embedder.subclass_from_config(config)
     assert isinstance(embedder, NoopEmbedder)
 
 
 def test_subclass_from_config_bag_of_tokens():
-    config = ObjectContructionConfig.model_validate(
+    config = ObjectConstructionConfig.model_validate(
         {
             "type": "ragbits.core.embeddings.sparse:BagOfTokens",
             "config": {

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

@@ -1,10 +1,12 @@
+import litellm
+
 from ragbits.core.llms import LLM
 from ragbits.core.llms.litellm import LiteLLM, LiteLLMOptions
-from ragbits.core.utils.config_handling import ObjectContructionConfig
+from ragbits.core.utils.config_handling import ObjectConstructionConfig
 
 
 def test_subclass_from_config():
-    config = ObjectContructionConfig.model_validate(
+    config = ObjectConstructionConfig.model_validate(
         {
             "type": "ragbits.core.llms.litellm:LiteLLM",
             "config": {
@@ -27,6 +29,45 @@ def test_subclass_from_config():
 
 
 def test_subclass_from_config_default_path():
-    config = ObjectContructionConfig.model_validate({"type": "LiteLLM"})
+    config = ObjectConstructionConfig.model_validate({"type": "LiteLLM"})
+    llm: LLM = LLM.subclass_from_config(config)
+    assert isinstance(llm, LiteLLM)
+
+
+def test_from_config_with_router():
+    config = ObjectConstructionConfig(
+        type="ragbits.core.llms.litellm:LiteLLM",
+        config={
+            "model_name": "gpt-4-turbo",
+            "api_key": "test_api_key",
+            "router": [
+                {
+                    "model_name": "gpt-4o",
+                    "litellm_params": {
+                        "model": "azure/gpt-4o-eval-1",
+                        "api_key": "test_api_key",
+                        "api_version": "2024-07-19-test",
+                        "api_base": "https://test-api.openai.azure.com",
+                    },
+                },
+                {
+                    "model_name": "gpt-4o",
+                    "litellm_params": {
+                        "model": "azure/gpt-4o-eval-2",
+                        "api_key": "test_api_key",
+                        "api_version": "2024-07-19-test",
+                        "api_base": "https://test-api.openai.azure.com",
+                    },
+                },
+            ],
+        },
+    )
+
     llm: LLM = LLM.subclass_from_config(config)
     assert isinstance(llm, LiteLLM)
+    assert llm.base_url is None
+    assert llm.model_name == "gpt-4-turbo"
+    assert llm.api_key == "test_api_key"
+    assert isinstance(llm.router, litellm.router.Router)
+    assert len(llm.router.model_list) == 2
+    assert llm.router.model_list[0]["model_name"] == "gpt-4o"

packages/ragbits-core/tests/unit/utils/test_config_handling.py

@@ -5,7 +5,7 @@
 
 from ragbits.core.config import CoreConfig, core_config
 from ragbits.core.utils._pyproject import get_config_instance
-from ragbits.core.utils.config_handling import InvalidConfigError, ObjectContructionConfig, WithConstructionConfig
+from ragbits.core.utils.config_handling import InvalidConfigError, ObjectConstructionConfig, WithConstructionConfig
 
 projects_dir = Path(__file__).parent / "testprojects"
 
@@ -40,7 +40,7 @@ def test_default_from_config():
 
 
 def test_subclass_from_config():
-    config = ObjectContructionConfig.model_validate(
+    config = ObjectConstructionConfig.model_validate(
         {
             "type": "ExampleSubclass",
             "config": {"foo": "foo", "bar": 1},
@@ -53,7 +53,7 @@ def test_subclass_from_config():
 
 
 def test_incorrect_subclass_from_config():
-    config = ObjectContructionConfig.model_validate(
+    config = ObjectConstructionConfig.model_validate(
         {
             "type": "ExampleWithNoDefaultModule",  # Not a subclass of ExampleClassWithConfigMixin
             "config": {"foo": "foo", "bar": 1},
@@ -64,7 +64,7 @@ def test_incorrect_subclass_from_config():
 
 
 def test_no_default_module():
-    config = ObjectContructionConfig.model_validate(
+    config = ObjectConstructionConfig.model_validate(
         {
             "type": "ExampleWithNoDefaultModule",
             "config": {"foo": "foo", "bar": 1},