feat: add helpful import guards for all optional dependencies

Author: jakelorocco

cli/alora/intrinsic_uploader.py

@@ -12,8 +12,14 @@
 import tempfile
 from typing import Literal
 
-import git
-from huggingface_hub import HfFolder, RepoUrl, create_repo, upload_file, upload_folder
+try:
+    import git
+    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. "
+        'Please install them with: pip install "mellea[hf]"'
+    ) from e
 
 
 def upload_intrinsic(

cli/alora/train.py

@@ -12,19 +12,25 @@
 import sys
 import warnings
 
-import torch
-import typer
-from datasets import Dataset
-from peft import LoraConfig, get_peft_model
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    TrainerCallback,
-    TrainerControl,
-    TrainerState,
-    TrainingArguments,
-)
-from trl import DataCollatorForCompletionOnlyLM, SFTConfig, SFTTrainer
+try:
+    import torch
+    import typer
+    from datasets import Dataset
+    from peft import LoraConfig, get_peft_model
+    from transformers import (
+        AutoModelForCausalLM,
+        AutoTokenizer,
+        TrainerCallback,
+        TrainerControl,
+        TrainerState,
+        TrainingArguments,
+    )
+    from trl import DataCollatorForCompletionOnlyLM, SFTConfig, SFTTrainer
+except ImportError as e:
+    raise ImportError(
+        "The 'm alora' command requires extra dependencies. "
+        'Please install them with: pip install "mellea[hf]"'
+    ) from e
 
 # Handle MPS with old PyTorch versions on macOS only
 # Accelerate's GradScaler requires PyTorch >= 2.8.0 for MPS

cli/alora/upload.py

@@ -8,7 +8,13 @@
 
 import os
 
-from huggingface_hub import HfApi, HfFolder, create_repo, upload_folder
+try:
+    from huggingface_hub import HfApi, HfFolder, create_repo, upload_folder
+except ImportError as e:
+    raise ImportError(
+        "The 'm alora upload' command requires extra dependencies. "
+        'Please install them with: pip install "mellea[hf]"'
+    ) from e
 
 
 def upload_model(weight_path: str, model_name: str, private: bool = True):

cli/serve/app.py

@@ -6,10 +6,16 @@
 import time
 import uuid
 
-import typer
-import uvicorn
-from fastapi import FastAPI
-from fastapi.responses import JSONResponse
+try:
+    import typer
+    import uvicorn
+    from fastapi import FastAPI
+    from fastapi.responses import JSONResponse
+except ImportError as e:
+    raise ImportError(
+        "The 'm serve' command requires extra dependencies. "
+        'Please install them with: pip install "mellea[server]"'
+    ) from e
 
 from .models import (
     ChatCompletion,

mellea/backends/huggingface.py

@@ -14,18 +14,24 @@
 from collections.abc import Callable, Coroutine, Sequence
 from typing import Any, overload
 
-import llguidance
-import llguidance.hf
-import llguidance.torch
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from transformers.cache_utils import DynamicCache
-from transformers.generation.logits_process import LogitsProcessorList
-from transformers.generation.streamers import AsyncTextIteratorStreamer
-from transformers.generation.utils import GenerateDecoderOnlyOutput
-from transformers.modeling_utils import PreTrainedModel
-from transformers.tokenization_utils import PreTrainedTokenizer
-from transformers.trainer_utils import set_seed
+try:
+    import llguidance
+    import llguidance.hf
+    import llguidance.torch
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from transformers.cache_utils import DynamicCache
+    from transformers.generation.logits_process import LogitsProcessorList
+    from transformers.generation.streamers import AsyncTextIteratorStreamer
+    from transformers.generation.utils import GenerateDecoderOnlyOutput
+    from transformers.modeling_utils import PreTrainedModel
+    from transformers.tokenization_utils import PreTrainedTokenizer
+    from transformers.trainer_utils import set_seed
+except ImportError as e:
+    raise ImportError(
+        "The HuggingFace backend requires extra dependencies. "
+        'Please install them with: pip install "mellea[hf]"'
+    ) from e
 
 from ..backends import kv_block_helpers
 from ..core import (

mellea/backends/kv_block_helpers.py

@@ -11,10 +11,16 @@
 from functools import reduce
 from typing import Any
 
-import torch
-from transformers import PreTrainedModel
-from transformers.cache_utils import DynamicCache
-from transformers.tokenization_utils_base import BatchEncoding
+try:
+    import torch
+    from transformers import PreTrainedModel
+    from transformers.cache_utils import DynamicCache
+    from transformers.tokenization_utils_base import BatchEncoding
+except ImportError as e:
+    raise ImportError(
+        "The HuggingFace backend requires extra dependencies. "
+        'Please install them with: pip install "mellea[hf]"'
+    ) from e
 
 TokenizedCacheIterleaving = Iterable[BatchEncoding | DynamicCache]
 LegacyCache = Any

mellea/formatters/granite/retrievers/util.py

@@ -10,8 +10,14 @@
 import zipfile
 
 # Third Party
-import pyarrow as pa  # type: ignore[import-not-found]
-import pyarrow.json as pj  # type: ignore[import-not-found]
+try:
+    import pyarrow as pa  # type: ignore[import-not-found]
+    import pyarrow.json as pj  # type: ignore[import-not-found]
+except ImportError as e:
+    raise ImportError(
+        "The granite retrievers module requires extra dependencies. "
+        'Please install them with: pip install "mellea[granite_retriever]"'
+    ) from e
 
 
 def download_mtrag_corpus(target_dir: str, corpus_name: str) -> pathlib.Path:

mellea/stdlib/components/docs/richdocument.py

@@ -13,11 +13,17 @@
 import io
 from pathlib import Path
 
-from docling.datamodel.base_models import InputFormat
-from docling.datamodel.pipeline_options import PdfPipelineOptions
-from docling.document_converter import DocumentConverter, PdfFormatOption
-from docling_core.types.doc.document import DoclingDocument, TableItem
-from docling_core.types.io import DocumentStream
+try:
+    from docling.datamodel.base_models import InputFormat
+    from docling.datamodel.pipeline_options import PdfPipelineOptions
+    from docling.document_converter import DocumentConverter, PdfFormatOption
+    from docling_core.types.doc.document import DoclingDocument, TableItem
+    from docling_core.types.io import DocumentStream
+except ImportError as e:
+    raise ImportError(
+        "RichDocument requires extra dependencies. "
+        'Please install them with: pip install "mellea[docling]"'
+    ) from e
 
 from ....backends.tools import MelleaTool
 from ....core import CBlock, Component, ModelOutputThunk, TemplateRepresentation

pyproject.toml

@@ -97,6 +97,7 @@ cli = [
 server = [
     "uvicorn",
     "fastapi",
+    "mellea[cli]",
 ]
 
 sandbox = [