Fix proxy crash when torch not installed (kompress lazy imports)

The proxy startup crashed with `ModuleNotFoundError: No module named
'torch'` when installed with just `[proxy]` extras because
kompress_compressor.py had unconditional top-level torch imports.
Moved torch/transformers imports to be lazy so the module is safely
importable without the [ml] extra. Added tests for import safety.
Bumped version to 0.4.5

headroom/__init__.py

@@ -153,7 +153,7 @@ from .transforms import (
     TransformPipeline,
 )
 
-__version__ = "0.4.4"
+__version__ = "0.4.5"
 
 __all__ = [
     # Main client

headroom/transforms/kompress_compressor.py

@@ -3,8 +3,7 @@
 Drop-in replacement for LLMLingua-2. Auto-downloads the model from
 HuggingFace (chopratejas/kompress-base) on first use.
 
-No extra pip install needed — uses transformers + safetensors
-which are already Headroom dependencies.
+Requires the [ml] extra: pip install headroom-ai[ml]
 
 Usage:
     >>> from headroom.transforms.kompress_compressor import KompressCompressor
@@ -20,10 +19,6 @@ import threading
 from dataclasses import dataclass
 from typing import Any
 
-import torch
-import torch.nn as nn
-from transformers import AutoModel, AutoTokenizer
-
 from ..config import TransformResult
 from ..tokenizer import Tokenizer
 from .base import Transform
@@ -39,64 +34,95 @@ _kompress_tokenizer = None
 _kompress_lock = threading.Lock()
 
 
-# ── Model Architecture (must match training) ──────────────────────────
-
+def is_kompress_available() -> bool:
+    """Check if Kompress dependencies are available (requires [ml] extra)."""
+    try:
+        import huggingface_hub  # noqa: F401
+        import safetensors  # noqa: F401
+        import torch  # noqa: F401
+        import transformers  # noqa: F401
 
-class HeadroomCompressorModel(nn.Module):
-    """Dual-head ModernBERT: token classification + span importance CNN."""
+        return True
+    except ImportError:
+        return False
 
-    def __init__(self, model_name: str = "answerdotai/ModernBERT-base"):
-        super().__init__()
-        self.encoder = AutoModel.from_pretrained(model_name, attn_implementation="eager")
-        hidden_size = self.encoder.config.hidden_size  # 768
 
-        # Head 1: Token keep/discard
-        self.token_dropout = nn.Dropout(0.1)
-        self.token_head = nn.Linear(hidden_size, 2)
+# ── Model Architecture (must match training) ──────────────────────────
+# torch/transformers are imported lazily — only when actually needed.
+# This allows `from kompress_compressor import is_kompress_available`
+# to work without torch installed.
+
+
+def _get_model_class() -> type:
+    """Return the HeadroomCompressorModel class, importing torch on demand."""
+    import torch
+    import torch.nn as nn
+    from transformers import AutoModel
+
+    class HeadroomCompressorModel(nn.Module):
+        """Dual-head ModernBERT: token classification + span importance CNN."""
+
+        def __init__(self, model_name: str = "answerdotai/ModernBERT-base"):
+            super().__init__()
+            self.encoder = AutoModel.from_pretrained(model_name, attn_implementation="eager")
+            hidden_size = self.encoder.config.hidden_size  # 768
+
+            # Head 1: Token keep/discard
+            self.token_dropout = nn.Dropout(0.1)
+            self.token_head = nn.Linear(hidden_size, 2)
+
+            # Head 2: Span importance (1D CNN)
+            self.span_conv = nn.Sequential(
+                nn.Conv1d(hidden_size, 256, kernel_size=5, padding=2),
+                nn.GELU(),
+                nn.Conv1d(256, 1, kernel_size=3, padding=1),
+                nn.Sigmoid(),
+            )
 
-        # Head 2: Span importance (1D CNN)
-        self.span_conv = nn.Sequential(
-            nn.Conv1d(hidden_size, 256, kernel_size=5, padding=2),
-            nn.GELU(),
-            nn.Conv1d(256, 1, kernel_size=3, padding=1),
-            nn.Sigmoid(),
-        )
+        def get_keep_mask(
+            self, input_ids: torch.Tensor, attention_mask: torch.Tensor
+        ) -> torch.Tensor:
+            """Get per-token keep/discard decision. True = keep."""
+            with torch.no_grad():
+                hidden = self.encoder(input_ids, attention_mask=attention_mask).last_hidden_state
 
-    def get_keep_mask(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
-        """Get per-token keep/discard decision. True = keep."""
-        with torch.no_grad():
-            hidden = self.encoder(input_ids, attention_mask=attention_mask).last_hidden_state
+                # Token head: binary classifier — argmax decides keep/discard
+                token_logits = self.token_head(hidden)  # [B, L, 2]
+                token_keep = (
+                    token_logits[:, :, 1] > token_logits[:, :, 0]
+                )  # True if class 1 > class 0
 
-            # Token head: binary classifier — argmax decides keep/discard
-            token_logits = self.token_head(hidden)  # [B, L, 2]
-            token_keep = token_logits[:, :, 1] > token_logits[:, :, 0]  # True if class 1 > class 0
+                # Span head: boost tokens in important spans
+                # If a token is borderline but its span is important, keep it
+                span_scores = self.span_conv(hidden.transpose(1, 2)).squeeze(1)
+                span_boost = span_scores > 0.5  # span says this region matters
 
-            # Span head: boost tokens in important spans
-            # If a token is borderline but its span is important, keep it
-            span_scores = self.span_conv(hidden.transpose(1, 2)).squeeze(1)
-            span_boost = span_scores > 0.5  # span says this region matters
+                # Keep if: token head says keep, OR token is borderline and span says keep
+                token_probs = torch.softmax(token_logits, dim=-1)[:, :, 1]
+                borderline = (token_probs > 0.3) & (token_probs <= 0.5)
+                keep = token_keep | (borderline & span_boost)
 
-            # Keep if: token head says keep, OR token is borderline and span says keep
-            token_probs = torch.softmax(token_logits, dim=-1)[:, :, 1]
-            borderline = (token_probs > 0.3) & (token_probs <= 0.5)
-            keep = token_keep | (borderline & span_boost)
+                return keep  # type: ignore[no-any-return]
 
-            return keep  # type: ignore[no-any-return]
+        def get_scores(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+            """Get per-token importance scores (for ranking when target_ratio is set)."""
+            with torch.no_grad():
+                hidden = self.encoder(input_ids, attention_mask=attention_mask).last_hidden_state
+                token_probs = torch.softmax(self.token_head(hidden), dim=-1)[:, :, 1]
+                span_scores = self.span_conv(hidden.transpose(1, 2)).squeeze(1)
+                return token_probs * (0.5 + 0.5 * span_scores)  # type: ignore[no-any-return]
 
-    def get_scores(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
-        """Get per-token importance scores (for ranking when target_ratio is set)."""
-        with torch.no_grad():
-            hidden = self.encoder(input_ids, attention_mask=attention_mask).last_hidden_state
-            token_probs = torch.softmax(self.token_head(hidden), dim=-1)[:, :, 1]
-            span_scores = self.span_conv(hidden.transpose(1, 2)).squeeze(1)
-            return token_probs * (0.5 + 0.5 * span_scores)  # type: ignore[no-any-return]
+    return HeadroomCompressorModel
 
 
 # ── Model Loading ─────────────────────────────────────────────────────
 
 
-def _load_kompress(device: str = "auto") -> tuple[HeadroomCompressorModel, Any]:
+def _load_kompress(device: str = "auto") -> tuple[Any, Any]:
     """Download from HuggingFace and load the Kompress model."""
+    import torch
+    from transformers import AutoTokenizer
+
     global _kompress_model, _kompress_tokenizer
 
     with _kompress_lock:
@@ -111,6 +137,7 @@ def _load_kompress(device: str = "auto") -> tuple[HeadroomCompressorModel, Any]:
         weights_path = hf_hub_download(HF_MODEL_ID, "model.safetensors")
 
         # Load architecture
+        HeadroomCompressorModel = _get_model_class()
         model = HeadroomCompressorModel()
 
         # Load trained weights
@@ -139,19 +166,6 @@ def _load_kompress(device: str = "auto") -> tuple[HeadroomCompressorModel, Any]:
         return model, tokenizer
 
 
-def is_kompress_available() -> bool:
-    """Check if Kompress dependencies are available (requires [ml] extra)."""
-    try:
-        import huggingface_hub  # noqa: F401
-        import safetensors  # noqa: F401
-        import torch  # noqa: F401
-        import transformers  # noqa: F401
-
-        return True
-    except ImportError:
-        return False
-
-
 def unload_kompress_model() -> bool:
     """Unload the Kompress model to free memory."""
     global _kompress_model, _kompress_tokenizer
@@ -159,8 +173,13 @@ def unload_kompress_model() -> bool:
         if _kompress_model is not None:
             _kompress_model = None
             _kompress_tokenizer = None
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
+            try:
+                import torch
+
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+            except ImportError:
+                pass
             return True
     return False
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "headroom-ai"
-version = "0.4.4"
+version = "0.4.5"
 description = "The Context Optimization Layer for LLM Applications - Cut costs by 50-90%"
 readme = "README.md"
 license = "Apache-2.0"

tests/test_transforms/test_kompress_compressor.py

@@ -0,0 +1,204 @@
+"""Tests for Kompress compressor.
+
+Covers:
+- Lazy imports: module importable without torch installed
+- is_kompress_available(): correct detection of [ml] extra
+- KompressConfig / KompressResult: dataclass defaults
+- KompressCompressor: passthrough for short content, fallback on error
+- Transform interface: apply() method
+"""
+
+from unittest.mock import MagicMock, patch
+
+# ── Import safety (the whole point of the fix) ─────────────────────────
+
+
+class TestLazyImports:
+    """The module must be importable without torch/transformers."""
+
+    def test_is_kompress_available_importable(self) -> None:
+        """is_kompress_available can be imported even without torch."""
+        from headroom.transforms.kompress_compressor import is_kompress_available
+
+        # Should return bool (True or False depending on environment)
+        result = is_kompress_available()
+        assert isinstance(result, bool)
+
+    def test_module_import_without_torch(self) -> None:
+        """Importing the module with torch blocked should not raise."""
+        import sys
+
+        # Block torch imports
+        with patch.dict(sys.modules, {"torch": None, "torch.nn": None}):
+            # Force re-evaluation of is_kompress_available
+            from headroom.transforms.kompress_compressor import is_kompress_available
+
+            # Should gracefully return False, not crash
+            assert is_kompress_available() is False
+
+    def test_dataclasses_importable_without_torch(self) -> None:
+        """KompressConfig, KompressResult, KompressCompressor are importable without torch."""
+        from headroom.transforms.kompress_compressor import (
+            KompressCompressor,  # noqa: F401
+            KompressConfig,
+            KompressResult,
+        )
+
+        # These don't need torch to instantiate
+        config = KompressConfig()
+        assert config.device == "auto"
+        assert config.enable_ccr is True
+
+        result = KompressResult(
+            compressed="hello",
+            original="hello world",
+            original_tokens=2,
+            compressed_tokens=1,
+            compression_ratio=0.5,
+        )
+        assert result.tokens_saved == 1
+        assert result.savings_percentage == 50.0
+
+
+# ── KompressResult ──────────────────────────────────────────────────────
+
+
+class TestKompressResult:
+    def test_tokens_saved(self) -> None:
+        from headroom.transforms.kompress_compressor import KompressResult
+
+        r = KompressResult(
+            compressed="a b",
+            original="a b c d",
+            original_tokens=4,
+            compressed_tokens=2,
+            compression_ratio=0.5,
+        )
+        assert r.tokens_saved == 2
+
+    def test_tokens_saved_no_negative(self) -> None:
+        from headroom.transforms.kompress_compressor import KompressResult
+
+        r = KompressResult(
+            compressed="a b c d e",
+            original="a b c",
+            original_tokens=3,
+            compressed_tokens=5,
+            compression_ratio=1.67,
+        )
+        assert r.tokens_saved == 0
+
+    def test_savings_percentage_zero_tokens(self) -> None:
+        from headroom.transforms.kompress_compressor import KompressResult
+
+        r = KompressResult(
+            compressed="",
+            original="",
+            original_tokens=0,
+            compressed_tokens=0,
+            compression_ratio=1.0,
+        )
+        assert r.savings_percentage == 0.0
+
+    def test_default_model(self) -> None:
+        from headroom.transforms.kompress_compressor import HF_MODEL_ID, KompressResult
+
+        r = KompressResult(
+            compressed="x",
+            original="x y",
+            original_tokens=2,
+            compressed_tokens=1,
+            compression_ratio=0.5,
+        )
+        assert r.model_used == HF_MODEL_ID
+
+
+# ── KompressCompressor (without model) ──────────────────────────────────
+
+
+class TestKompressCompressorPassthrough:
+    """Test compressor behavior that doesn't require the actual model."""
+
+    def test_short_content_passthrough(self) -> None:
+        """Content under 10 words should pass through unchanged."""
+        from headroom.transforms.kompress_compressor import KompressCompressor
+
+        compressor = KompressCompressor()
+        result = compressor.compress("hello world")
+        assert result.compressed == "hello world"
+        assert result.compression_ratio == 1.0
+        assert result.original_tokens == 2
+        assert result.compressed_tokens == 2
+
+    def test_empty_content_passthrough(self) -> None:
+        from headroom.transforms.kompress_compressor import KompressCompressor
+
+        compressor = KompressCompressor()
+        result = compressor.compress("")
+        assert result.compressed == ""
+        assert result.compression_ratio == 1.0
+
+    def test_fallback_on_model_error(self) -> None:
+        """If _load_kompress fails, compress should return passthrough."""
+        from headroom.transforms.kompress_compressor import KompressCompressor
+
+        compressor = KompressCompressor()
+        long_text = " ".join(f"word{i}" for i in range(20))
+
+        with patch(
+            "headroom.transforms.kompress_compressor._load_kompress",
+            side_effect=RuntimeError("no model"),
+        ):
+            result = compressor.compress(long_text)
+            assert result.compressed == long_text
+            assert result.compression_ratio == 1.0
+
+
+# ── Transform interface ─────────────────────────────────────────────────
+
+
+class TestKompressTransformInterface:
+    def test_apply_short_messages_unchanged(self) -> None:
+        """Messages with <10 words should pass through apply() unchanged."""
+        from headroom.transforms.kompress_compressor import KompressCompressor
+
+        compressor = KompressCompressor()
+        messages = [
+            {"role": "user", "content": "hello"},
+            {"role": "tool", "content": "short"},
+        ]
+        tokenizer = MagicMock()
+        tokenizer.count_text = MagicMock(return_value=5)
+
+        result = compressor.apply(messages, tokenizer)
+        assert len(result.messages) == 2
+        assert result.messages[0]["content"] == "hello"
+        assert result.messages[1]["content"] == "short"
+
+    def test_apply_preserves_user_messages(self) -> None:
+        """User messages should never be compressed."""
+        from headroom.transforms.kompress_compressor import KompressCompressor
+
+        compressor = KompressCompressor()
+        long_text = " ".join(f"word{i}" for i in range(50))
+        messages = [{"role": "user", "content": long_text}]
+        tokenizer = MagicMock()
+        tokenizer.count_text = MagicMock(return_value=50)
+
+        with patch(
+            "headroom.transforms.kompress_compressor._load_kompress",
+            side_effect=RuntimeError("should not be called"),
+        ):
+            result = compressor.apply(messages, tokenizer)
+            assert result.messages[0]["content"] == long_text
+
+
+# ── unload_kompress_model ───────────────────────────────────────────────
+
+
+class TestUnloadKompressModel:
+    def test_unload_when_no_model(self) -> None:
+        from headroom.transforms.kompress_compressor import unload_kompress_model
+
+        # Should return False when no model is loaded
+        assert unload_kompress_model() is False