headroom-ai 0.2.13__py3-none-any.whl

headroom/compression/__init__.py

@@ -0,0 +1,42 @@
+"""Universal compression with ML-based content detection.
+
+This module provides intelligent, automatic compression that:
+1. Detects content type using ML (Magika)
+2. Preserves structure (keys, signatures, templates)
+3. Compresses content with LLMLingua
+4. Enables retrieval via CCR
+
+Quick Start:
+    # One-liner for simple use
+    from headroom.compression import compress
+    result = compress(content)
+
+    # Or with configuration
+    from headroom.compression import UniversalCompressor, UniversalCompressorConfig
+
+    config = UniversalCompressorConfig(compression_ratio_target=0.5)
+    compressor = UniversalCompressor(config=config)
+    result = compressor.compress(content)
+"""
+
+from headroom.compression.detector import ContentType, MagikaDetector
+from headroom.compression.masks import StructureMask
+from headroom.compression.universal import (
+    CompressionResult,
+    UniversalCompressor,
+    UniversalCompressorConfig,
+    compress,
+)
+
+__all__ = [
+    # Simple API
+    "compress",
+    # Full API
+    "UniversalCompressor",
+    "UniversalCompressorConfig",
+    "CompressionResult",
+    # Advanced
+    "MagikaDetector",
+    "ContentType",
+    "StructureMask",
+]

headroom/compression/detector.py

@@ -0,0 +1,424 @@
+"""ML-based content type detection using Google's Magika.
+
+Magika is a deep learning model for content type detection that:
+- Runs locally (~5ms latency)
+- Supports 100+ content types
+- Has 99%+ accuracy on supported types
+- Requires no configuration
+
+This replaces rule-based detection with learned detection.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from magika import Magika
+    from magika.types import MagikaResult
+
+logger = logging.getLogger(__name__)
+
+# Lazy-loaded Magika instance (singleton)
+_magika_instance: Magika | None = None
+
+
+class ContentType(Enum):
+    """High-level content categories for compression routing."""
+
+    JSON = "json"
+    CODE = "code"
+    LOG = "log"
+    MARKDOWN = "markdown"
+    TEXT = "text"
+    UNKNOWN = "unknown"
+
+
+@dataclass
+class DetectionResult:
+    """Result of ML-based content detection."""
+
+    content_type: ContentType
+    confidence: float  # 0.0 to 1.0
+    raw_label: str  # Original Magika label
+    language: str | None = None  # For code: python, javascript, etc.
+    metadata: dict = field(default_factory=dict)
+
+
+# Map Magika labels to our content types
+# This is the ONLY place where we map labels - no hardcoding elsewhere
+_CODE_LABELS = frozenset(
+    {
+        "python",
+        "javascript",
+        "typescript",
+        "go",
+        "rust",
+        "java",
+        "c",
+        "cpp",
+        "csharp",
+        "ruby",
+        "php",
+        "swift",
+        "kotlin",
+        "scala",
+        "shell",
+        "bash",
+        "powershell",
+        "sql",
+        "r",
+        "perl",
+        "lua",
+        "haskell",
+        "elixir",
+        "erlang",
+        "clojure",
+        "ocaml",
+        "fsharp",
+        "dart",
+        "julia",
+        "zig",
+        "nim",
+        "crystal",
+        "v",
+        "solidity",
+        "move",
+        "cairo",
+        "vyper",
+    }
+)
+
+_STRUCTURED_LABELS = frozenset(
+    {
+        "json",
+        "jsonl",
+        "yaml",
+        "toml",
+        "xml",
+        "html",
+        "csv",
+        "tsv",
+        "ini",
+        "properties",
+    }
+)
+
+_LOG_LABELS = frozenset(
+    {
+        "log",
+        "syslog",
+    }
+)
+
+_MARKDOWN_LABELS = frozenset(
+    {
+        "markdown",
+        "rst",
+        "asciidoc",
+        "org",
+    }
+)
+
+
+def _get_magika() -> Magika:
+    """Get or create the singleton Magika instance.
+
+    Lazy-loads on first use to avoid import cost if not needed.
+    """
+    global _magika_instance
+    if _magika_instance is None:
+        try:
+            from magika import Magika
+
+            _magika_instance = Magika()
+            logger.debug("Magika model loaded successfully")
+        except ImportError as e:
+            raise ImportError(
+                "Magika is required for ML-based content detection. "
+                "Install with: pip install magika"
+            ) from e
+    return _magika_instance
+
+
+def _magika_available() -> bool:
+    """Check if Magika is available without loading it."""
+    try:
+        import magika  # noqa: F401
+
+        return True
+    except ImportError:
+        return False
+
+
+class MagikaDetector:
+    """ML-based content type detector using Google's Magika.
+
+    This detector uses a deep learning model to identify content types
+    without relying on file extensions or brittle regex patterns.
+
+    Example:
+        detector = MagikaDetector()
+        result = detector.detect('def hello(): print("hi")')
+        # result.content_type == ContentType.CODE
+        # result.language == "python"
+    """
+
+    def __init__(self, min_confidence: float = 0.5):
+        """Initialize the detector.
+
+        Args:
+            min_confidence: Minimum confidence threshold. Below this,
+                returns ContentType.UNKNOWN.
+        """
+        self.min_confidence = min_confidence
+        self._magika: Magika | None = None
+
+    def _ensure_magika(self) -> Magika:
+        """Ensure Magika is loaded."""
+        if self._magika is None:
+            self._magika = _get_magika()
+        return self._magika
+
+    def detect(self, content: str) -> DetectionResult:
+        """Detect content type using ML.
+
+        Args:
+            content: The content to analyze.
+
+        Returns:
+            DetectionResult with type, confidence, and metadata.
+
+        Example:
+            >>> detector = MagikaDetector()
+            >>> result = detector.detect('{"users": [{"id": 1}]}')
+            >>> result.content_type
+            ContentType.JSON
+        """
+        if not content or not content.strip():
+            return DetectionResult(
+                content_type=ContentType.UNKNOWN,
+                confidence=0.0,
+                raw_label="empty",
+            )
+
+        # Get Magika prediction
+        magika = self._ensure_magika()
+        result: MagikaResult = magika.identify_bytes(content.encode("utf-8"))
+
+        raw_label = result.output.ct_label
+        confidence = result.output.score
+
+        # Map to our content type
+        content_type, language = self._map_label(raw_label)
+
+        # Apply confidence threshold
+        if confidence < self.min_confidence:
+            content_type = ContentType.UNKNOWN
+
+        return DetectionResult(
+            content_type=content_type,
+            confidence=confidence,
+            raw_label=raw_label,
+            language=language,
+            metadata={
+                "magika_group": result.output.group,
+                "magika_mime": result.output.mime_type,
+            },
+        )
+
+    def detect_batch(self, contents: list[str]) -> list[DetectionResult]:
+        """Detect content types for multiple contents.
+
+        More efficient than calling detect() in a loop.
+
+        Args:
+            contents: List of content strings to analyze.
+
+        Returns:
+            List of DetectionResults in same order as input.
+        """
+        if not contents:
+            return []
+
+        magika = self._ensure_magika()
+        results = []
+
+        # Convert to bytes for Magika
+        byte_contents = [c.encode("utf-8") for c in contents]
+
+        # Batch detection
+        magika_results = magika.identify_bytes_batch(byte_contents)
+
+        for content, magika_result in zip(contents, magika_results):
+            if not content or not content.strip():
+                results.append(
+                    DetectionResult(
+                        content_type=ContentType.UNKNOWN,
+                        confidence=0.0,
+                        raw_label="empty",
+                    )
+                )
+                continue
+
+            raw_label = magika_result.output.ct_label
+            confidence = magika_result.output.score
+            content_type, language = self._map_label(raw_label)
+
+            if confidence < self.min_confidence:
+                content_type = ContentType.UNKNOWN
+
+            results.append(
+                DetectionResult(
+                    content_type=content_type,
+                    confidence=confidence,
+                    raw_label=raw_label,
+                    language=language,
+                    metadata={
+                        "magika_group": magika_result.output.group,
+                        "magika_mime": magika_result.output.mime_type,
+                    },
+                )
+            )
+
+        return results
+
+    def _map_label(self, label: str) -> tuple[ContentType, str | None]:
+        """Map Magika label to our ContentType.
+
+        Args:
+            label: Raw Magika label (e.g., "python", "json").
+
+        Returns:
+            Tuple of (ContentType, optional language).
+        """
+        label_lower = label.lower()
+
+        # Check code languages
+        if label_lower in _CODE_LABELS:
+            return ContentType.CODE, label_lower
+
+        # Check structured data
+        if label_lower in _STRUCTURED_LABELS:
+            # JSON gets its own type for specialized handling
+            if label_lower in ("json", "jsonl"):
+                return ContentType.JSON, None
+            # Other structured data treated as JSON-like
+            return ContentType.JSON, None
+
+        # Check logs
+        if label_lower in _LOG_LABELS:
+            return ContentType.LOG, None
+
+        # Check markdown/docs
+        if label_lower in _MARKDOWN_LABELS:
+            return ContentType.MARKDOWN, None
+
+        # Text types
+        if label_lower in ("txt", "text", "ascii", "utf8", "empty"):
+            return ContentType.TEXT, None
+
+        # Default: treat as text
+        return ContentType.TEXT, None
+
+    @staticmethod
+    def is_available() -> bool:
+        """Check if Magika is available."""
+        return _magika_available()
+
+
+class FallbackDetector:
+    """Simple fallback detector when Magika is not available.
+
+    Uses basic heuristics - not as accurate but requires no dependencies.
+    """
+
+    def __init__(self, min_confidence: float = 0.5):
+        """Initialize the fallback detector."""
+        self.min_confidence = min_confidence
+
+    def detect(self, content: str) -> DetectionResult:
+        """Detect content type using simple heuristics.
+
+        Args:
+            content: The content to analyze.
+
+        Returns:
+            DetectionResult with type and confidence.
+        """
+        if not content or not content.strip():
+            return DetectionResult(
+                content_type=ContentType.UNKNOWN,
+                confidence=0.0,
+                raw_label="empty",
+            )
+
+        stripped = content.strip()
+
+        # JSON detection (simple but effective)
+        if stripped.startswith(("{", "[")):
+            try:
+                import json
+
+                json.loads(stripped)
+                return DetectionResult(
+                    content_type=ContentType.JSON,
+                    confidence=1.0,
+                    raw_label="json",
+                )
+            except (json.JSONDecodeError, ValueError):
+                pass
+
+        # Code detection (look for common patterns)
+        code_indicators = [
+            "def ",
+            "class ",
+            "function ",
+            "import ",
+            "const ",
+            "let ",
+            "var ",
+            "func ",
+            "fn ",
+            "pub ",
+            "package ",
+        ]
+        if any(indicator in content for indicator in code_indicators):
+            return DetectionResult(
+                content_type=ContentType.CODE,
+                confidence=0.7,
+                raw_label="code",
+            )
+
+        # Log detection
+        log_indicators = ["ERROR", "WARN", "INFO", "DEBUG", "FATAL"]
+        if any(indicator in content for indicator in log_indicators):
+            return DetectionResult(
+                content_type=ContentType.LOG,
+                confidence=0.6,
+                raw_label="log",
+            )
+
+        # Default to text
+        return DetectionResult(
+            content_type=ContentType.TEXT,
+            confidence=0.5,
+            raw_label="text",
+        )
+
+
+def get_detector(prefer_magika: bool = True) -> MagikaDetector | FallbackDetector:
+    """Get the best available detector.
+
+    Args:
+        prefer_magika: If True, use Magika if available.
+
+    Returns:
+        MagikaDetector if available and preferred, else FallbackDetector.
+    """
+    if prefer_magika and MagikaDetector.is_available():
+        return MagikaDetector()
+    return FallbackDetector()

headroom/compression/handlers/__init__.py

@@ -0,0 +1,22 @@
+"""Structure handlers for different content types.
+
+Each handler knows how to extract structural information from a specific
+content type and create a StructureMask marking what should be preserved.
+
+Handlers don't compress - they only identify structure. The actual
+compression is done by LLMLingua on the non-structural parts.
+"""
+
+from headroom.compression.handlers.base import (
+    HandlerResult,
+    StructureHandler,
+)
+from headroom.compression.handlers.code_handler import CodeStructureHandler
+from headroom.compression.handlers.json_handler import JSONStructureHandler
+
+__all__ = [
+    "StructureHandler",
+    "HandlerResult",
+    "JSONStructureHandler",
+    "CodeStructureHandler",
+]

headroom/compression/handlers/base.py

@@ -0,0 +1,219 @@
+"""Base class and protocol for structure handlers.
+
+Structure handlers extract structural information from content and create
+masks identifying what should be preserved during compression.
+
+The handler protocol is simple:
+1. get_mask(content) -> StructureMask
+2. can_handle(content) -> bool (optional)
+
+Handlers are content-type specific but domain-agnostic. A JSONStructureHandler
+preserves JSON keys whether it's user data, search results, or config files.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Protocol, runtime_checkable
+
+from headroom.compression.masks import StructureMask
+
+
+@dataclass
+class HandlerResult:
+    """Result from a structure handler.
+
+    Contains the mask plus metadata about what was detected.
+    """
+
+    mask: StructureMask
+    handler_name: str
+    confidence: float = 1.0  # How confident the handler is in its detection
+    metadata: dict = field(default_factory=dict)
+
+    @property
+    def preservation_ratio(self) -> float:
+        """Fraction of content marked for preservation."""
+        return self.mask.preservation_ratio
+
+
+@runtime_checkable
+class StructureHandler(Protocol):
+    """Protocol for structure handlers.
+
+    Any class implementing get_mask() can be used as a handler.
+    """
+
+    @property
+    def name(self) -> str:
+        """Handler name for logging and metadata."""
+        ...
+
+    def get_mask(
+        self,
+        content: str,
+        tokens: list[str] | None = None,
+        **kwargs: Any,
+    ) -> HandlerResult:
+        """Extract structure mask from content.
+
+        Args:
+            content: The content to analyze.
+            tokens: Pre-tokenized content (optional). If not provided,
+                handler should tokenize internally.
+            **kwargs: Handler-specific options.
+
+        Returns:
+            HandlerResult with mask and metadata.
+        """
+        ...
+
+    def can_handle(self, content: str) -> bool:
+        """Check if this handler can process the content.
+
+        Default implementation returns True. Override for handlers
+        that need to verify content format before processing.
+
+        Args:
+            content: The content to check.
+
+        Returns:
+            True if handler can process this content.
+        """
+        ...
+
+
+class BaseStructureHandler(ABC):
+    """Base implementation for structure handlers.
+
+    Provides common functionality and enforces the handler interface.
+    Subclasses must implement _extract_mask().
+    """
+
+    def __init__(self, name: str | None = None):
+        """Initialize the handler.
+
+        Args:
+            name: Optional handler name. Defaults to class name.
+        """
+        self._name = name or self.__class__.__name__
+
+    @property
+    def name(self) -> str:
+        """Handler name."""
+        return self._name
+
+    def get_mask(
+        self,
+        content: str,
+        tokens: list[str] | None = None,
+        **kwargs: Any,
+    ) -> HandlerResult:
+        """Extract structure mask from content.
+
+        This is the main entry point. It handles common logic like
+        empty content and delegates to _extract_mask() for the
+        content-specific logic.
+
+        Args:
+            content: The content to analyze.
+            tokens: Pre-tokenized content (optional).
+            **kwargs: Handler-specific options.
+
+        Returns:
+            HandlerResult with mask and metadata.
+        """
+        # Handle empty content
+        if not content or not content.strip():
+            tokens = tokens or []
+            return HandlerResult(
+                mask=StructureMask.empty(tokens),
+                handler_name=self.name,
+                confidence=0.0,
+                metadata={"empty": True},
+            )
+
+        # Tokenize if not provided
+        if tokens is None:
+            tokens = self._tokenize(content)
+
+        # Delegate to subclass
+        return self._extract_mask(content, tokens, **kwargs)
+
+    def can_handle(self, content: str) -> bool:
+        """Check if this handler can process the content.
+
+        Default implementation returns True. Override for handlers
+        that need to verify content format.
+
+        Args:
+            content: The content to check.
+
+        Returns:
+            True if handler can process this content.
+        """
+        return True
+
+    @abstractmethod
+    def _extract_mask(
+        self,
+        content: str,
+        tokens: list[str],
+        **kwargs: Any,
+    ) -> HandlerResult:
+        """Extract structure mask from content.
+
+        Subclasses implement this to provide content-specific logic.
+
+        Args:
+            content: The content to analyze (non-empty, stripped).
+            tokens: Tokenized content.
+            **kwargs: Handler-specific options.
+
+        Returns:
+            HandlerResult with mask and metadata.
+        """
+        ...
+
+    def _tokenize(self, content: str) -> list[str]:
+        """Default tokenization - character-level.
+
+        Subclasses may override for more sophisticated tokenization.
+        For mask purposes, character-level is often sufficient and
+        aligns well with LLMLingua's token-level compression.
+
+        Args:
+            content: Content to tokenize.
+
+        Returns:
+            List of tokens (characters by default).
+        """
+        # Simple character-level tokenization
+        # This aligns well with structure detection (we mark ranges)
+        return list(content)
+
+
+class NoOpHandler(BaseStructureHandler):
+    """Handler that marks everything as compressible.
+
+    Used as a fallback when no structure is detected.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the no-op handler."""
+        super().__init__(name="noop")
+
+    def _extract_mask(
+        self,
+        content: str,
+        tokens: list[str],
+        **kwargs: Any,
+    ) -> HandlerResult:
+        """Return mask with everything compressible."""
+        return HandlerResult(
+            mask=StructureMask.empty(tokens),
+            handler_name=self.name,
+            confidence=1.0,
+            metadata={"reason": "no structure detected"},
+        )