fix: missing nice import errors and missing test scenarios

Author: jakelorocco

cli/alora/intrinsic_uploader.py

@@ -14,7 +14,13 @@
 
 try:
     import git
-    from huggingface_hub import HfFolder, RepoUrl, create_repo, upload_file, upload_folder
+    from huggingface_hub import (
+        HfFolder,
+        RepoUrl,
+        create_repo,
+        upload_file,
+        upload_folder,
+    )
 except ImportError as e:
     raise ImportError(
         "The 'm alora upload' command requires extra dependencies. "

cli/m.py

@@ -8,17 +8,17 @@
 
 try:
     import typer
-except ImportError:
-    raise SystemExit(
+except ImportError as e:
+    raise ImportError(
         "The 'm' CLI requires extra dependencies. "
         'Please install them with: pip install "mellea[cli]"'
-    ) from None
+    ) from e
 
 from cli.alora.commands import alora_app
 from cli.decompose import app as decompose_app
 from cli.eval.commands import eval_app
 from cli.fix import fix_app
-from cli.serve.app import serve
+from cli.serve.commands import serve
 
 cli = typer.Typer(name="m", no_args_is_help=True)
 
@@ -36,10 +36,10 @@ def callback() -> None:
 
 # Typer assumes that all commands are in the same file/module.
 # Use this workaround to separate out functionality. Can still be called
-# as if added with @cli.command() (ie `m serve` here).
+# as if added with @cli.command() (ie `m serve` here). If we don't use this
+# approach, we would have to use `m server <subcommand>` instead.
 cli.command(name="serve")(serve)
 
-
 # Add new subcommand groups by importing and adding with `cli.add_typer()`
 # as documented: https://typer.tiangolo.com/tutorial/subcommands/add-typer/#put-them-together.
 cli.add_typer(alora_app)

cli/serve/app.py

@@ -130,14 +130,7 @@ async def endpoint(request: ChatCompletionRequest):
     return endpoint
 
 
-def serve(
-    script_path: str = typer.Argument(
-        default="docs/examples/m_serve/example.py",
-        help="Path to the Python script to import and serve",
-    ),
-    host: str = typer.Option("0.0.0.0", help="Host to bind to"),
-    port: int = typer.Option(8080, help="Port to bind to"),
-):
+def run_server(script_path: str, host: str = "0.0.0.0", port: int = 8080):
     """Serve a FastAPI endpoint for a given script."""
     module = load_module_from_path(script_path)
     route_path = "/v1/chat/completions"

cli/serve/commands.py

@@ -0,0 +1,22 @@
+"""Typer command definition for ``m serve``.
+
+Separates the CLI interface (typer annotations) from the server implementation
+(FastAPI, uvicorn) so that ``m --help`` works without the ``server`` extra installed.
+The heavy server dependencies are only imported when ``m serve`` is actually invoked.
+"""
+
+import typer
+
+
+def serve(
+    script_path: str = typer.Argument(
+        default="docs/examples/m_serve/example.py",
+        help="Path to the Python script to import and serve",
+    ),
+    host: str = typer.Option("0.0.0.0", help="Host to bind to"),
+    port: int = typer.Option(8080, help="Port to bind to"),
+):
+    """Serve a FastAPI endpoint for a given script."""
+    from cli.serve.app import run_server
+
+    run_server(script_path=script_path, host=host, port=port)

mellea/backends/litellm.py

@@ -7,9 +7,15 @@
 from collections.abc import Callable, Coroutine, Sequence
 from typing import Any, overload
 
-import litellm
-import litellm.litellm_core_utils
-import litellm.litellm_core_utils.get_supported_openai_params
+try:
+    import litellm
+    import litellm.litellm_core_utils
+    import litellm.litellm_core_utils.get_supported_openai_params
+except ImportError as e:
+    raise ImportError(
+        "The LiteLLM backend requires extra dependencies. "
+        'Please install them with: pip install "mellea[litellm]"'
+    ) from e
 
 from ..backends import model_ids
 from ..core import (

mellea/backends/watsonx.py

@@ -10,9 +10,15 @@
 from dataclasses import fields
 from typing import Any, overload
 
-from ibm_watsonx_ai import APIClient, Credentials
-from ibm_watsonx_ai.foundation_models import ModelInference
-from ibm_watsonx_ai.foundation_models.schema import TextChatParameters
+try:
+    from ibm_watsonx_ai import APIClient, Credentials
+    from ibm_watsonx_ai.foundation_models import ModelInference
+    from ibm_watsonx_ai.foundation_models.schema import TextChatParameters
+except ImportError as e:
+    raise ImportError(
+        "The Watsonx backend requires extra dependencies. "
+        'Please install them with: pip install "mellea[watsonx]"'
+    ) from e
 
 from ..backends import ModelIdentifier, model_ids
 from ..core import (

test/package/test_dependency_isolation.py

@@ -60,11 +60,11 @@
     "docling": ["from mellea.stdlib.components.docs.richdocument import RichDocument"],
     "granite_retriever": [
         "import sentence_transformers",
-        "import pyarrow",
+        "from mellea.formatters.granite.retrievers.util import download_mtrag_corpus",
         "import elasticsearch",
     ],
-    "cli": ["import typer"],
-    "server": ["import uvicorn", "import fastapi"],
+    "cli": ["from cli.m import cli"],
+    "server": ["from cli.serve.app import run_server"],
     "sandbox": ["import llm_sandbox"],
     "hooks": [
         "import cpex"  # We directly import a non-Mellea class/package here since the hooks classes are defined regardless of if cpex is installed.
@@ -89,6 +89,11 @@
 # Meta-groups that compose other extras (no dedicated isolation test needed).
 META_GROUPS = {"backends", "all"}
 
+# Extras whose IMPORTS entries all go through guarded mellea/cli modules
+# (i.e. ImportError will contain a mellea[<extra>] hint). Allows us to
+# check for a nice ImportError message.
+GUARDED_EXTRAS = {*BACKEND_EXTRAS, "server", "docling", "cli"}
+
 # Extras that have a corresponding test_<name> function in this module.
 # Used to determine if we are missing tests for a given optional-dependency group.
 TESTED_EXTRAS = {
@@ -125,32 +130,38 @@ def _parse_optional_dependency_groups() -> set[str]:
     return set(data.get("project", {}).get("optional-dependencies", {}).keys())
 
 
-def _backend_fail_imports(exclude: set[str] | str = "") -> list[str]:
-    """Return backend import statements that should fail, excluding the given extras.
+def _normalize_exclude(exclude: set[str] | str = "") -> set[str]:
+    """Normalize an exclude argument to a set of extra names."""
+    if isinstance(exclude, str):
+        return {exclude} if exclude else set()
+    return exclude
+
+
+def _backend_fail_imports(exclude: set[str] | str = "") -> list[tuple[str, str]]:
+    """Return ``(import_stmt, extra_name)`` pairs for backends expected to fail.
 
     Collects import statements from all ``BACKEND_EXTRAS`` except those in
     ``exclude``, and also excludes ``vllm`` on macOS since it is unsupported there.
+    Each tuple pairs the import statement with the extra name so the checker
+    script can verify the ``ImportError`` contains a ``mellea[<extra>]`` hint.
 
     Args:
         exclude (set[str] | str): Backend extras to omit (their imports are
             expected to succeed). Pass a set of extra names, a single name, or
             an empty string to include all backends.
 
     Returns:
-        list[str]: Import statements that should raise ``ImportError`` in the
-            isolated environment.
+        list[tuple[str, str]]: Pairs of ``(import_statement, extra_name)``.
     """
-    if isinstance(exclude, str):
-        exclude = {exclude} if exclude else set()
-    extras = BACKEND_EXTRAS - exclude
+    extras = BACKEND_EXTRAS - _normalize_exclude(exclude)
     if IS_MACOS:
         extras -= {"vllm"}
-    return [stmt for name in sorted(extras) for stmt in IMPORTS[name]]
+    return [(stmt, name) for name in sorted(extras) for stmt in IMPORTS[name]]
 
 
 def _build_check_script(
     should_succeed: list[str] = [],
-    should_fail: list[str] = [],
+    should_fail: list[tuple[str, str]] = [],
     flag_checks: list[tuple[str, str, bool]] = [],
 ) -> str:
     """Build a self-contained Python script that tests imports and prints failures.
@@ -161,7 +172,9 @@ def _build_check_script(
 
     Args:
         should_succeed (list[str]): Import statements that must execute without error.
-        should_fail (list[str]): Import statements that must raise ``ImportError``.
+        should_fail (list[tuple[str, str]]): Tuples of
+            ``(import_statement, extra_name)`` — the import must raise
+            ``ImportError`` and the message must contain ``mellea[<extra_name>]``.
         flag_checks (list[tuple[str, str, bool]]): Tuples of
             ``(module_path, attribute_name, expected_value)`` for boolean flag
             assertions (e.g. ``("mellea.telemetry.tracing", "_OTEL_AVAILABLE", True)``).
@@ -178,12 +191,16 @@ def _build_check_script(
         lines.append(f"    failures.append(f'SHOULD SUCCEED: {stmt}: {{e}}')")
         lines.append("")
 
-    for stmt in should_fail:
+    for stmt, extra in should_fail:
+        hint = f"mellea[{extra}]"
         lines.append("try:")
         lines.append(f"    {stmt}")
         lines.append(f"    failures.append('SHOULD FAIL but succeeded: {stmt}')")
-        lines.append("except ImportError:")
-        lines.append("    pass")
+        lines.append("except ImportError as e:")
+        lines.append(f"    if '{hint}' not in str(e):")
+        lines.append(
+            f"        failures.append(f'MISSING HINT {hint} in: {stmt}: {{e}}')"
+        )
         lines.append("except Exception:")
         lines.append("    pass  # non-ImportError counts as unavailable")
         lines.append("")
@@ -209,7 +226,7 @@ def _build_check_script(
 def _run_check(
     extra: str | None = None,
     should_succeed: list[str] = [],
-    should_fail: list[str] = [],
+    should_fail: list[tuple[str, str]] = [],
     flag_checks: list[tuple[str, str, bool]] = [],
 ) -> None:
     """Build and run an import check script in an isolated uv environment.
@@ -223,7 +240,9 @@ def _run_check(
         extra (str | None): The optional-dependency extra to install
             (e.g. ``"hf"``). Pass ``None`` to install core mellea only.
         should_succeed (list[str]): Import statements that must execute without error.
-        should_fail (list[str]): Import statements that must raise ``ImportError``.
+        should_fail (list[tuple[str, str]]): Tuples of
+            ``(import_statement, extra_name)`` — must raise ``ImportError``
+            with ``mellea[<extra_name>]`` in the message.
         flag_checks (list[tuple[str, str, bool]]): Tuples of
             ``(module_path, attribute_name, expected_value)`` for boolean flag
             assertions.
@@ -232,7 +251,8 @@ def _run_check(
         ValueError: If any import statement appears in both ``should_succeed``
             and ``should_fail``.
     """
-    overlap = set(should_succeed) & set(should_fail)
+    fail_stmts = {stmt for stmt, _ in should_fail}
+    overlap = set(should_succeed) & fail_stmts
     if overlap:
         raise ValueError(
             f"Same imports in both should_succeed and should_fail: {overlap}"
@@ -286,10 +306,17 @@ def _inverted_flag_checks(extra: str) -> list[tuple[str, str, bool]]:
 
 
 def test_core_only() -> None:
-    """Core mellea with no extras: basic imports work, optional backends fail."""
+    """Core mellea with no extras: basic imports work, optional extras fail with hints."""
     _run_check(
         should_succeed=IMPORTS["core"],
-        should_fail=_backend_fail_imports(exclude=""),
+        should_fail=[
+            *_backend_fail_imports(exclude=""),
+            *[
+                (stmt, name)
+                for name in sorted(GUARDED_EXTRAS - BACKEND_EXTRAS)
+                for stmt in IMPORTS[name]
+            ],
+        ],
         flag_checks=[
             *_inverted_flag_checks("telemetry"),
             *_inverted_flag_checks("hooks"),
@@ -298,7 +325,7 @@ def test_core_only() -> None:
 
 
 def test_hf() -> None:
-    """mellea[hf]: HuggingFace backend imports succeed, others fail."""
+    """mellea[hf]: HuggingFace backend imports succeed, others fail with hints."""
     _run_check(
         extra="hf",
         should_succeed=[*IMPORTS["core"], *IMPORTS["hf"]],
@@ -308,7 +335,7 @@ def test_hf() -> None:
 
 @pytest.mark.skipif(IS_MACOS, reason="vllm is not supported on macOS")
 def test_vllm() -> None:
-    """mellea[vllm]: vLLM backend imports succeed, others fail."""
+    """mellea[vllm]: vLLM backend imports succeed, others fail with hints."""
     _run_check(
         extra="vllm",
         should_succeed=[*IMPORTS["core"], *IMPORTS["vllm"]],
@@ -317,7 +344,7 @@ def test_vllm() -> None:
 
 
 def test_litellm() -> None:
-    """mellea[litellm]: LiteLLM backend imports succeed, others fail."""
+    """mellea[litellm]: LiteLLM backend imports succeed, others fail with hints."""
     _run_check(
         extra="litellm",
         should_succeed=[*IMPORTS["core"], *IMPORTS["litellm"]],
@@ -326,7 +353,7 @@ def test_litellm() -> None:
 
 
 def test_watsonx() -> None:
-    """mellea[watsonx]: Watsonx backend imports succeed, others fail."""
+    """mellea[watsonx]: Watsonx backend imports succeed, others fail with hints."""
     _run_check(
         extra="watsonx",
         should_succeed=[*IMPORTS["core"], *IMPORTS["watsonx"]],
@@ -462,7 +489,7 @@ def test_checker_detects_should_succeed_failure() -> None:
 
 def test_checker_detects_should_fail_that_succeeds() -> None:
     """Checker script exits non-zero when a should_fail import actually works."""
-    script = _build_check_script(should_fail=["import json"])
+    script = _build_check_script(should_fail=[("import json", "fake")])
     result = _run_script_raw(script)
     assert result.returncode != 0
     assert "SHOULD FAIL but succeeded" in result.stderr
@@ -478,11 +505,21 @@ def test_checker_detects_wrong_flag_value() -> None:
     assert "FLAG sys.maxsize" in result.stderr
 
 
+def test_checker_detects_missing_hint() -> None:
+    """Checker script exits non-zero when ImportError lacks the expected install hint."""
+    # Raise an ImportError without the expected hint substring
+    script = _build_check_script(
+        should_fail=[("import no_such_module_xyz", "fake_extra")]
+    )
+    result = _run_script_raw(script)
+    assert result.returncode != 0
+    assert "MISSING HINT mellea[fake_extra]" in result.stderr
+
+
 def test_checker_passes_when_all_correct() -> None:
     """Checker script exits zero when all checks pass."""
     script = _build_check_script(
         should_succeed=["import json", "from os.path import join"],
-        should_fail=["import no_such_module_xyz"],
         flag_checks=[("sys", "maxsize", sys.maxsize)],  # type: ignore[list-item]
     )
     result = _run_script_raw(script)