feat: separate out remaining dependencies and improve tests

cli/alora/intrinsic_uploader.py

@@ -12,8 +12,20 @@
 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/m.py

@@ -6,13 +6,19 @@
 evaluation). Run ``m --help`` to see all available sub-commands.
 """
 
-import typer
+try:
+    import typer
+except ImportError as e:
+    raise ImportError(
+        "The 'm' CLI requires extra dependencies. "
+        'Please install them with: pip install "mellea[cli]"'
+    ) 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)
 
@@ -30,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

@@ -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,
@@ -124,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)

docs/alora.md

@@ -9,7 +9,7 @@ Mellea provides a command-line interface for training and uploading [LoRA](https
 From the root of the repository:
 
 ```bash
-pip install mellea
+pip install "mellea[cli]"
 huggingface-cli login  # Optional: only needed for uploads
 ```
 

docs/docs/advanced/lora-and-alora-adapters.md

@@ -11,7 +11,7 @@ schemes not well-represented in general training data. Mellea lets you train a
 [aLoRA](https://github.com/IBM/activated-lora) adapter on your own labeled dataset
 and use it as a requirement validator in any Mellea program.
 
-**Prerequisites:** `pip install mellea`, `m` CLI available. Training requires a GPU or
+**Prerequisites:** `pip install "mellea[cli]"`. Training requires a GPU or
 Apple Silicon Mac with sufficient VRAM for the chosen base model. Uploading requires a
 Hugging Face account.
 

docs/docs/getting-started/installation.md

@@ -25,6 +25,7 @@ pip install "mellea[litellm]"    # LiteLLM multi-provider (Anthropic, Bedrock, e
 pip install "mellea[hf]"         # HuggingFace transformers for local inference
 pip install "mellea[watsonx]"    # IBM WatsonX
 pip install "mellea[tools]"      # Tool and agent dependencies (LangChain, smolagents)
+pip install "mellea[cli]"        # m serve, m alora, m decompose CLI commands
 pip install "mellea[telemetry]"  # OpenTelemetry tracing and metrics
 ```
 
@@ -33,6 +34,7 @@ uv add "mellea[litellm]"        # LiteLLM multi-provider (Anthropic, Bedrock, et
 uv add "mellea[hf]"             # HuggingFace transformers for local inference
 uv add "mellea[watsonx]"        # IBM WatsonX
 uv add "mellea[tools]"          # Tool and agent dependencies (LangChain, smolagents)
+uv add "mellea[cli]"            # m serve, m alora, m decompose CLI commands
 uv add "mellea[telemetry]"      # OpenTelemetry tracing and metrics
 ```
 

docs/docs/guide/m-decompose.md

@@ -11,7 +11,7 @@ description: "Break complex tasks into ordered, executable subtasks with the m d
 3. Generate a prompt template for each subtask
 4. Output a ready-to-run Python script that executes each subtask in order
 
-**Prerequisites:** Mellea installed (`uv add mellea`), Ollama running locally (or an OpenAI-compatible endpoint).
+**Prerequisites:** Mellea installed (`uv add "mellea[cli]"`), Ollama running locally (or an OpenAI-compatible endpoint).
 
 ## Basic usage
 

docs/docs/how-to/refactor-prompts-with-cli.md

@@ -5,7 +5,7 @@ description: "Use m decompose to break a complex prompt into typed, validated ge
 # diataxis: how-to
 ---
 
-**Prerequisites:** `pip install mellea`, Ollama running locally (or an
+**Prerequisites:** `pip install "mellea[cli]"`, Ollama running locally (or an
 OpenAI-compatible endpoint).
 
 When a single prompt grows too long or asks the LLM to do too many things at

docs/docs/integrations/m-serve.md

@@ -8,7 +8,7 @@ description: "Run a Mellea program as an OpenAI-compatible chat endpoint with m
 any LLM client — LangChain, the OpenAI SDK, `curl` — call your Mellea program as if
 it were a model.
 
-**Prerequisites:** `pip install mellea`.
+**Prerequisites:** `pip install "mellea[cli]"`.
 
 ## The serve() function
 

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/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 (

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