"""Universal compressor with ML-based detection and structure preservation. This is the main entry point for compression. It: 1. Detects content type using Magika (ML) 2. Extracts structure using appropriate handler 3. Compresses non-structural content with Kompress 4. Optionally stores original in CCR for retrieval Usage: compressor = UniversalCompressor() result = compressor.compress(content) # Result contains: # - compressed: The compressed content # - compression_ratio: original_tokens / compressed_tokens # - content_type: Detected content type # - preservation_ratio: Fraction of content preserved as structure """ from __future__ import annotations import logging from collections.abc import Callable from dataclasses import dataclass, field from typing import Any from headroom.compression.detector import ( ContentType, DetectionResult, FallbackDetector, get_detector, ) from headroom.compression.handlers.base import ( NoOpHandler, StructureHandler, ) from headroom.compression.handlers.code_handler import CodeStructureHandler from headroom.compression.handlers.json_handler import JSONStructureHandler from headroom.compression.masks import ( StructureMask, compute_entropy_mask_for_content, mask_to_spans, ) logger = logging.getLogger(__name__) @dataclass class UniversalCompressorConfig: """Configuration for UniversalCompressor. Attributes: use_magika: Use ML-based detection (requires magika package). use_kompress: Use Kompress for content compression. use_entropy_preservation: Preserve high-entropy tokens (UUIDs, etc.). entropy_threshold: Threshold for entropy-based preservation. min_content_length: Minimum content length to compress. compression_ratio_target: Target compression ratio (0.0-1.0). ccr_enabled: Store originals in CCR for retrieval. """ use_magika: bool = True use_kompress: bool = True use_entropy_preservation: bool = True entropy_threshold: float = 0.85 min_content_length: int = 100 compression_ratio_target: float = 0.3 # Target 70% reduction ccr_enabled: bool = True @dataclass class CompressionResult: """Result from compression. Attributes: compressed: The compressed content. original: The original content (for reference). compression_ratio: compressed_length / original_length. tokens_before: Estimated token count before compression. tokens_after: Estimated token count after compression. content_type: Detected content type. detection_confidence: Confidence of content type detection. handler_used: Name of structure handler used. preservation_ratio: Fraction of content marked as structural. ccr_key: CCR storage key (if CCR enabled). metadata: Additional metadata. """ compressed: str original: str compression_ratio: float tokens_before: int tokens_after: int content_type: ContentType detection_confidence: float handler_used: str preservation_ratio: float ccr_key: str | None = None metadata: dict = field(default_factory=dict) @property def tokens_saved(self) -> int: """Number of tokens saved.""" return max(0, self.tokens_before - self.tokens_after) @property def savings_percentage(self) -> float: """Percentage of tokens saved.""" if self.tokens_before == 0: return 0.0 return (self.tokens_saved / self.tokens_before) * 100 class UniversalCompressor: """Universal compressor with ML detection and structure preservation. This compressor automatically: 1. Detects content type (JSON, code, logs, text) using ML 2. Extracts structure (keys, signatures, templates) 3. Preserves structure while compressing content 4. Stores original for CCR retrieval Example: >>> compressor = UniversalCompressor() >>> result = compressor.compress('{"users": [{"id": 1, "name": "Alice"}]}') >>> print(result.content_type) # ContentType.JSON >>> print(result.compressed) # Structure preserved, values compressed """ def __init__( self, config: UniversalCompressorConfig | None = None, handlers: dict[ContentType, StructureHandler] | None = None, compress_fn: Callable[[str], str] | None = None, ): """Initialize the compressor. Args: config: Compression configuration. handlers: Custom handlers for content types. compress_fn: Custom compression function. If None, uses Kompress when available, else simple truncation. """ self.config = config or UniversalCompressorConfig() # Initialize detector if self.config.use_magika: self._detector = get_detector(prefer_magika=True) else: self._detector = FallbackDetector() # Initialize handlers self._handlers: dict[ContentType, StructureHandler] = handlers or { ContentType.JSON: JSONStructureHandler(), ContentType.CODE: CodeStructureHandler(), } self._noop_handler = NoOpHandler() # Initialize compression function self._compress_fn = compress_fn or self._get_default_compress_fn() # CCR store (lazy initialized) self._ccr_store: Any | None = None def _get_default_compress_fn(self) -> Callable[[str], str]: """Get default compression function. Returns Kompress wrapper if available, else simple truncation. """ if self.config.use_kompress: try: return self._kompress_compress except ImportError: logger.info("Kompress not available, using simple compression") return self._simple_compress def _kompress_compress(self, text: str) -> str: """Compress using Kompress. Args: text: Text to compress. Returns: Compressed text. """ try: from headroom.transforms.kompress_compressor import KompressCompressor compressor = KompressCompressor() result = compressor.compress(text) return result.compressed except ImportError: return self._simple_compress(text) except Exception as e: logger.warning("Kompress compression failed: %s", e) return self._simple_compress(text) def _simple_compress(self, text: str) -> str: """Simple compression fallback (truncation with indicator). Args: text: Text to compress. Returns: Truncated text with indicator. """ target_len = int(len(text) * self.config.compression_ratio_target) if len(text) <= target_len: return text # Keep first and last portions. The separator must contain no control # characters (RFC 8259 §7): this fallback can run on a span inside a # JSON string value, and a raw newline there produces invalid JSON. keep_start = target_len * 2 // 3 keep_end = target_len // 3 return text[:keep_start] + " ...[compressed]... " + text[-keep_end:] def compress( self, content: str, content_type: ContentType | None = None, **kwargs: Any, ) -> CompressionResult: """Compress content with structure preservation. Args: content: Content to compress. content_type: Override content type detection. **kwargs: Handler-specific options. Returns: CompressionResult with compressed content and metadata. """ # Handle empty/short content if not content or len(content) < self.config.min_content_length: return CompressionResult( compressed=content, original=content, compression_ratio=1.0, tokens_before=self._estimate_tokens(content), tokens_after=self._estimate_tokens(content), content_type=ContentType.UNKNOWN, detection_confidence=0.0, handler_used="none", preservation_ratio=1.0, metadata={"skipped": "content too short"}, ) # Detect content type if content_type is None: detection = self._detector.detect(content) else: detection = DetectionResult( content_type=content_type, confidence=1.0, raw_label="override", ) # Get handler for content type handler = self._handlers.get(detection.content_type, self._noop_handler) # Tokenize content (character-level for masks) tokens = list(content) # Get structure mask from handler handler_result = handler.get_mask(content, tokens, **kwargs) structure_mask = handler_result.mask # Optionally add entropy-based preservation. Score whole words (mapped # back to character positions) rather than the character-level `tokens` # above — single-character tokens never reach the min length, which # would make entropy preservation a silent no-op on plain text. if self.config.use_entropy_preservation: entropy_mask = compute_entropy_mask_for_content( content, threshold=self.config.entropy_threshold, ) # Union: preserve if either mask says preserve structure_mask = structure_mask.union(entropy_mask) # Apply compression to non-structural parts compressed = self._compress_with_mask(content, structure_mask) # Estimate tokens tokens_before = self._estimate_tokens(content) tokens_after = self._estimate_tokens(compressed) # Store in CCR if enabled ccr_key = None if self.config.ccr_enabled: ccr_key = self._store_in_ccr(content, compressed) return CompressionResult( compressed=compressed, original=content, compression_ratio=len(compressed) / len(content) if content else 1.0, tokens_before=tokens_before, tokens_after=tokens_after, content_type=detection.content_type, detection_confidence=detection.confidence, handler_used=handler_result.handler_name, preservation_ratio=structure_mask.preservation_ratio, ccr_key=ccr_key, metadata={ "detection": { "raw_label": detection.raw_label, "language": detection.language, }, "handler": handler_result.metadata, }, ) def _compress_with_mask(self, content: str, mask: StructureMask) -> str: """Apply compression respecting structure mask. Args: content: Original content. mask: Structure mask. Returns: Compressed content with structure preserved. """ spans = mask_to_spans(mask) result_parts: list[str] = [] for span in spans: span_content = content[span.start : span.end] if span.is_structural: # Preserve structural content result_parts.append(span_content) else: # Compress non-structural content if len(span_content) > 50: # Only compress if substantial compressed = self._compress_fn(span_content) result_parts.append(compressed) else: result_parts.append(span_content) return "".join(result_parts) def _estimate_tokens(self, text: str) -> int: """Estimate token count. Uses simple heuristic: ~4 characters per token. Args: text: Text to estimate. Returns: Estimated token count. """ if not text: return 0 # Simple estimation: ~4 chars per token on average return len(text) // 4 def _store_in_ccr(self, original: str, compressed: str) -> str | None: """Store original in CCR for retrieval. Args: original: Original content. compressed: Compressed content. Returns: CCR key if stored, None otherwise. """ try: if self._ccr_store is None: from headroom.cache.compression_store import CompressionStore self._ccr_store = CompressionStore() key = self._ccr_store.store( original, compressed, original_tokens=self._estimate_tokens(original), compressed_tokens=self._estimate_tokens(compressed), ) return str(key) if key else None except ImportError: logger.debug("CCR store not available") return None except Exception as e: logger.warning("Failed to store in CCR: %s", e) return None def compress_batch( self, contents: list[str], **kwargs: Any, ) -> list[CompressionResult]: """Compress multiple contents. More efficient than calling compress() in a loop for ML detection. Args: contents: List of contents to compress. **kwargs: Handler-specific options. Returns: List of CompressionResults. """ if not contents: return [] # Batch detection if hasattr(self._detector, "detect_batch"): detections = self._detector.detect_batch(contents) else: detections = [self._detector.detect(c) for c in contents] # Compress each with detected type results = [] for content, detection in zip(contents, detections): result = self.compress( content, content_type=detection.content_type, **kwargs, ) results.append(result) return results def get_handler(self, content_type: ContentType) -> StructureHandler: """Get handler for content type. Args: content_type: Content type. Returns: Handler for the content type. """ return self._handlers.get(content_type, self._noop_handler) def register_handler( self, content_type: ContentType, handler: StructureHandler, ) -> None: """Register a custom handler for a content type. Args: content_type: Content type to handle. handler: Handler instance. """ self._handlers[content_type] = handler def compress(content: str, **kwargs: Any) -> CompressionResult: """Convenience function for one-off compression. Args: content: Content to compress. **kwargs: Passed to UniversalCompressor.compress(). Returns: CompressionResult. Example: >>> from headroom.compression import compress >>> result = compress('{"users": [{"id": 1}, {"id": 2}]}') >>> print(result.compressed) """ compressor = UniversalCompressor() return compressor.compress(content, **kwargs)