feat(lazy-loading): Decreasing the time needed to start the app (#753)

Co-authored-by: GlockPL <[email protected]>
Co-authored-by: Mateusz Hordyński <[email protected]>

Author: GlockPL

packages/ragbits-core/CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Added Lazy loading of dependencies in local.py and during importing of LiteLLM
 - Add tool_choice parameter to LLM interface (#738)
 - Fix Prompt consumes same iterator twice leading to no data added to chat (#768)
 

packages/ragbits-core/pyproject.toml

@@ -36,7 +36,7 @@ dependencies = [
     "pydantic>=2.9.1,<3.0.0",
     "typer>=0.12.5,<1.0.0",
     "tomli>=2.0.2,<3.0.0",
-    "litellm>=1.55.0,<2.0.0",
+    "litellm>=1.74.0,<2.0.0",
     "aiohttp>=3.10.8,<4.0.0",
     "filetype>=1.2.0,<2.0.0",
     "griffe>=1.7.3,<2.0.0"

packages/ragbits-core/src/ragbits/core/__init__.py

@@ -1,12 +1,28 @@
 import os
+from concurrent.futures import ThreadPoolExecutor
 
 import typer
 
 from ragbits.core.audit.traces import set_trace_handlers
-from ragbits.core.config import import_modules_from_config
+
+_config_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="config-import")
+_config_future = None
+
+
+def _import_and_run_config() -> None:
+    from ragbits.core.config import import_modules_from_config
+
+    import_modules_from_config()
+
+
+def ensure_config_loaded() -> None:
+    """Wait for config import to complete if it hasn't already."""
+    if _config_future:
+        _config_future.result()
+
 
 if os.getenv("RAGBITS_VERBOSE", "0") == "1":
     typer.echo('Verbose mode is enabled with environment variable "RAGBITS_VERBOSE".')
     set_trace_handlers("cli")
 
-import_modules_from_config()
+_config_future = _config_executor.submit(_import_and_run_config)

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

@@ -1,5 +1,50 @@
+import threading
+from concurrent.futures import Future, ThreadPoolExecutor
+from functools import cache
+
 from .base import LLM, ToolCall, Usage
-from .litellm import LiteLLM, LiteLLMOptions
 from .local import LocalLLM, LocalLLMOptions
 
-__all__ = ["LLM", "LiteLLM", "LiteLLMOptions", "LocalLLM", "LocalLLMOptions", "ToolCall", "Usage"]
+_import_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="litellm-import")
+_litellm_future: Future[tuple[type, type]] | None = None
+_import_lock = threading.Lock()
+
+
+@cache
+def _import_litellm() -> tuple[type, type]:
+    from .litellm import LiteLLM, LiteLLMOptions
+
+    return LiteLLM, LiteLLMOptions
+
+
+def _start_litellm_import() -> None:
+    global _litellm_future  # noqa: PLW0603
+    with _import_lock:
+        if _litellm_future is None:
+            _litellm_future = _import_executor.submit(_import_litellm)
+
+
+def __getattr__(name: str) -> type:
+    if name in ("LiteLLM", "LiteLLMOptions"):
+        _start_litellm_import()
+        if _litellm_future is not None:
+            LiteLLM, LiteLLMOptions = _litellm_future.result()
+        else:
+            # Fallback to synchronous import if future is None
+            LiteLLM, LiteLLMOptions = _import_litellm()
+
+        if name == "LiteLLM":
+            return LiteLLM
+        elif name == "LiteLLMOptions":
+            return LiteLLMOptions
+
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+
+# Dynamic __all__ to handle lazy-loaded LiteLLM imports
+__all__ = ["LLM", "LocalLLM", "LocalLLMOptions", "ToolCall", "Usage"]
+
+
+def __dir__() -> list[str]:
+    """Return available module attributes including lazy-loaded ones."""
+    return __all__ + ["LiteLLM", "LiteLLMOptions"]

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

@@ -1,4 +1,5 @@
 import asyncio
+import threading
 import time
 from collections.abc import AsyncGenerator, Callable, Iterable
 from typing import Any, Literal
@@ -102,6 +103,15 @@ def __init__(
         self.custom_model_cost_config = custom_model_cost_config
         if custom_model_cost_config:
             litellm.register_model(custom_model_cost_config)
+        else:
+
+            def download_and_register_model_cost() -> None:
+                litellm.register_model(
+                    model_cost="https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
+                )
+
+            thread = threading.Thread(target=download_and_register_model_cost, daemon=True)
+            thread.start()
 
     def get_model_id(self) -> str:
         """

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

@@ -2,22 +2,17 @@
 import threading
 import time
 from collections.abc import AsyncGenerator, Iterable
-
-try:
-    import accelerate  # noqa: F401
-    import torch  # noqa: F401
-    from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer  # noqa: F401
-
-    HAS_LOCAL_LLM = True
-except ImportError:
-    HAS_LOCAL_LLM = False
+from typing import TYPE_CHECKING, Any
 
 from ragbits.core.audit.metrics import record_metric
 from ragbits.core.audit.metrics.base import LLMMetric, MetricType
 from ragbits.core.llms.base import LLM, LLMOptions, ToolChoice
 from ragbits.core.prompt.base import BasePrompt
 from ragbits.core.types import NOT_GIVEN, NotGiven
 
+if TYPE_CHECKING:
+    from transformers import TextIteratorStreamer
+
 
 class LocalLLMOptions(LLMOptions):
     """
@@ -69,8 +64,10 @@ def __init__(
             ImportError: If the 'local' extra requirements are not installed.
             ValueError: If the model was not trained as a chat model.
         """
-        if not HAS_LOCAL_LLM:
+        deps = self._lazy_import_local_deps()
+        if deps is None:
             raise ImportError("You need to install the 'local' extra requirements to use local LLM models")
+        torch, AutoModelForCausalLM, AutoTokenizer, self.TextIteratorStreamer = deps
 
         super().__init__(model_name, default_options)
         self.model = AutoModelForCausalLM.from_pretrained(
@@ -87,6 +84,16 @@ def __init__(
         self._price_per_prompt_token = price_per_prompt_token
         self._price_per_completion_token = price_per_completion_token
 
+    @staticmethod
+    def _lazy_import_local_deps() -> tuple[Any, Any, Any, Any] | None:
+        try:
+            import torch
+            from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+
+            return torch, AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+        except ImportError:
+            return None
+
     def get_model_id(self) -> str:
         """
         Returns the model id.
@@ -212,7 +219,7 @@ async def _call_streaming(
         input_ids = self.tokenizer.apply_chat_template(prompt.chat, add_generation_prompt=True, return_tensors="pt").to(
             self.model.device
         )
-        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True)
+        streamer = self.TextIteratorStreamer(self.tokenizer, skip_prompt=True)
         generation_kwargs = dict(streamer=streamer, **options.dict())
         generation_thread = threading.Thread(target=self.model.generate, args=(input_ids,), kwargs=generation_kwargs)
 
@@ -221,7 +228,7 @@ async def streamer_to_async_generator(
         ) -> AsyncGenerator[dict, None]:
             output_tokens = 0
             generation_thread.start()
-            for text in streamer:
+            for text in streamer:  # type: ignore[attr-defined]
                 if text:
                     output_tokens += 1
                     if output_tokens == 1:
@@ -270,3 +277,20 @@ async def streamer_to_async_generator(
             )
 
         return streamer_to_async_generator(streamer=streamer, generation_thread=generation_thread)
+
+
+def __getattr__(name: str) -> type:
+    """Allow access to transformers classes for testing purposes."""
+    if name in ("AutoModelForCausalLM", "AutoTokenizer", "TextIteratorStreamer"):
+        try:
+            from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+
+            transformers_classes = {
+                "AutoModelForCausalLM": AutoModelForCausalLM,
+                "AutoTokenizer": AutoTokenizer,
+                "TextIteratorStreamer": TextIteratorStreamer,
+            }
+            return transformers_classes[name]
+        except ImportError:
+            pass
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")

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

@@ -438,13 +438,19 @@ async def test_init_registers_model_with_custom_cost_config():
         mock_register.assert_called_once_with(custom_config)
 
 
-async def test_init_does_not_register_model_if_no_cost_config_is_provided():
-    """Test that the model is not registered if no cost config is provided."""
+async def test_init_registers_default_model_cost_when_no_custom_config_provided():
+    """Test that the default model cost config is registered when no custom config is provided."""
+    import time
+
     with patch("litellm.register_model") as mock_register:
         LiteLLM(
             model_name="some_model",
         )
-        mock_register.assert_not_called()
+        # Give the thread a moment to complete
+        time.sleep(0.1)
+        mock_register.assert_called_once_with(
+            model_cost="https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
+        )
 
 
 async def test_pickling_registers_model_with_custom_cost_config():