move check

wandb · Feb 5, 2025 · 7c1ed3a · 7c1ed3a
1 parent 14f6e17
commit 7c1ed3a
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 6 deletions.
diff --git a/weave/scorers/__init__.py b/weave/scorers/__init__.py
@@ -1,9 +1,7 @@
-try:
-    from litellm import acompletion  # noqa: F401
-except ImportError:
-    raise ImportError(
-        "litellm is required to use the LLM-powered scorers, please install it with `pip install litellm`"
-    )
+from weave.scorers.initialization import check_litellm_installation
+
+# Make sure litellm is available
+check_litellm_installation()
 
 from weave.scorers.classification_scorer import (
     MultiTaskBinaryClassificationF1,

diff --git a/weave/scorers/initialization.py b/weave/scorers/initialization.py
@@ -0,0 +1,13 @@
+"""Module to handle initialization of LLM-related dependencies."""
+
+
+def check_litellm_installation():
+    """Check if litellm is installed and raise an informative error if not."""
+    try:
+        from litellm import acompletion  # noqa: F401
+
+        return acompletion
+    except ImportError:
+        raise ImportError(
+            "litellm is required to use the LLM-powered scorers, please install it with `pip install litellm`"
+        )