feat: better litellm lazy loading / app speed-up (#806)

Author: mhordynski

.github/scripts/check_changelog_update.sh

@@ -1,10 +1,13 @@
 #!/bin/bash
 
-echo "Fetching main branch..."
-git fetch origin main --depth=1
+# Use the provided base branch or default to 'develop'
+BASE_BRANCH=${1:-develop}
 
-echo "Identifying changed files between the current branch and main branch..."
-CHANGED_FILES=$(git diff --name-only origin/main | tr '\n' ' ')
+echo "Fetching $BASE_BRANCH branch..."
+git fetch origin $BASE_BRANCH --depth=1
+
+echo "Identifying changed files between the current branch and $BASE_BRANCH branch..."
+CHANGED_FILES=$(git diff --name-only origin/$BASE_BRANCH | tr '\n' ' ')
 
 if [ -z "$CHANGED_FILES" ]; then
   echo "No files have been changed in this branch."
@@ -29,7 +32,7 @@ fi
 echo "Found changes in the following packages: $CHANGED_PACKAGES"
 
 # Look for "Changelog-ignore: <package-name>" in the commit message (possibly multiple entries in separate lines)
-IGNORED_PACKAGES=$(git log --pretty=format:%B origin/main..HEAD | grep -oP '^Changelog-ignore: \K[^ ]+' | sort -u)
+IGNORED_PACKAGES=$(git log --pretty=format:%B origin/$BASE_BRANCH..HEAD | grep -oP '^Changelog-ignore: \K[^ ]+' | sort -u)
 
 for IGNORED_PACKAGE in $IGNORED_PACKAGES; do
   if echo "$CHANGED_PACKAGES" | grep -q "^$IGNORED_PACKAGE$"; then
@@ -42,7 +45,7 @@ for PACKAGE in $CHANGED_PACKAGES; do
   CHANGELOG="packages/$PACKAGE/CHANGELOG.md"
   echo "Checking changelog for package: $PACKAGE"
 
-  if ! diff -u <(git show origin/main:$CHANGELOG | grep -Pzo '(?s)(## Unreleased.*?)(?=\n## |\Z)' | tr -d '\0') <(grep -Pzo '(?s)(## Unreleased.*?)(?=\n## |\Z)' $CHANGELOG | tr -d '\0') | grep -q '^\+'; then
+  if ! diff -u <(git show origin/$BASE_BRANCH:$CHANGELOG | grep -Pzo '(?s)(## Unreleased.*?)(?=\n## |\Z)' | tr -d '\0') <(grep -Pzo '(?s)(## Unreleased.*?)(?=\n## |\Z)' $CHANGELOG | tr -d '\0') | grep -q '^\+'; then
     echo "No updates detected in changelog for package $PACKAGE. Please add an entry under '## Unreleased'."
     exit 1
   fi

packages/ragbits-cli/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- feat: add current working directory to Python path for better CLI module discovery
+
 ## 1.2.2 (2025-08-08)
 
 ### Changed

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

@@ -1,6 +1,7 @@
 import importlib.util
 import os
 import pkgutil
+import sys
 from pathlib import Path
 from typing import Annotated
 
@@ -60,6 +61,7 @@ def autoregister() -> None:
         for module in pkgutil.iter_modules(ragbits.__path__)
         if module.ispkg and module.name != "cli" and (Path(module.module_finder.path) / module.name / "cli.py").exists()  # type: ignore
     ]
+    sys.path.append(os.getcwd())
 
     for module in cli_enabled_modules:
         register_func = importlib.import_module(f"ragbits.{module.name}.cli").register

packages/ragbits-core/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- feat: improve app startup speed with lazy LiteLLM loading
+
 
 - Feat: added support for IVFFlat indexing and Halfvec datatype
 

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

@@ -1,6 +1,5 @@
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast
 
-import litellm
 from typing_extensions import Self
 
 from ragbits.core.audit.traces import trace
@@ -14,6 +13,10 @@
 )
 from ragbits.core.options import Options
 from ragbits.core.types import NOT_GIVEN, NotGiven
+from ragbits.core.utils.lazy_litellm import LazyLiteLLM
+
+if TYPE_CHECKING:
+    from litellm import Router
 
 
 class LiteLLMEmbedderOptions(Options):
@@ -28,7 +31,7 @@ class LiteLLMEmbedderOptions(Options):
     encoding_format: str | None | NotGiven = NOT_GIVEN
 
 
-class LiteLLMEmbedder(DenseEmbedder[LiteLLMEmbedderOptions]):
+class LiteLLMEmbedder(DenseEmbedder[LiteLLMEmbedderOptions], LazyLiteLLM):
     """
     Client for creating text embeddings using LiteLLM API.
     """
@@ -44,7 +47,7 @@ def __init__(
         base_url: str | None = None,  # Alias for api_base
         api_key: str | None = None,
         api_version: str | None = None,
-        router: litellm.Router | None = None,
+        router: "Router | None" = None,
     ) -> None:
         """
         Constructs the LiteLLMEmbeddingClient.
@@ -119,7 +122,7 @@ async def embed_text(self, data: list[str], options: LiteLLMEmbedderOptions | No
             options=merged_options.dict(),
         ) as outputs:
             try:
-                entrypoint = self.router or litellm
+                entrypoint = self.router or self._litellm
                 response = await entrypoint.aembedding(
                     input=data,
                     model=self.model_name,
@@ -128,11 +131,11 @@ async def embed_text(self, data: list[str], options: LiteLLMEmbedderOptions | No
                     api_version=self.api_version,
                     **merged_options.dict(),
                 )
-            except litellm.openai.APIConnectionError as exc:
+            except self._litellm.openai.APIConnectionError as exc:
                 raise EmbeddingConnectionError() from exc
-            except litellm.openai.APIStatusError as exc:
+            except self._litellm.openai.APIStatusError as exc:
                 raise EmbeddingStatusError(exc.message, exc.status_code) from exc
-            except litellm.openai.APIResponseValidationError as exc:
+            except self._litellm.openai.APIResponseValidationError as exc:
                 raise EmbeddingResponseError() from exc
 
             if not response.data:
@@ -158,7 +161,7 @@ def from_config(cls, config: dict[str, Any]) -> Self:
             LiteLLMEmbedder: An initialized LiteLLMEmbedder instance.
         """
         if "router" in config:
-            router = litellm.router.Router(model_list=config["router"])
+            router = cls._get_litellm_module().router.Router(model_list=config["router"])
             config["router"] = router
 
         # Map base_url to api_base if present

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

@@ -1,50 +1,5 @@
-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
 
-_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"]
+__all__ = ["LLM", "LiteLLM", "LiteLLMOptions", "LocalLLM", "LocalLLMOptions", "ToolCall", "Usage"]