diff --git a/headroom/__init__.py b/headroom/__init__.py index 3d3d2c8f8..90d878572 100644 --- a/headroom/__init__.py +++ b/headroom/__init__.py @@ -153,7 +153,7 @@ from .transforms import ( TransformPipeline, ) -__version__ = "0.5.4" +__version__ = "0.5.5" __all__ = [ # Main client diff --git a/headroom/compression/detector.py b/headroom/compression/detector.py index 59e76898b..3a726a737 100644 --- a/headroom/compression/detector.py +++ b/headroom/compression/detector.py @@ -209,8 +209,8 @@ class MagikaDetector: magika = self._ensure_magika() result: MagikaResult = magika.identify_bytes(content.encode("utf-8")) - raw_label = result.output.ct_label - confidence = result.output.score + raw_label = result.output.label + confidence = result.score # Map to our content type content_type, language = self._map_label(raw_label) @@ -233,8 +233,6 @@ class MagikaDetector: 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. @@ -244,16 +242,9 @@ class MagikaDetector: 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): + for content in contents: if not content or not content.strip(): results.append( DetectionResult( @@ -264,8 +255,9 @@ class MagikaDetector: ) continue - raw_label = magika_result.output.ct_label - confidence = magika_result.output.score + magika_result = self._ensure_magika().identify_bytes(content.encode("utf-8")) + raw_label = magika_result.output.label + confidence = magika_result.score content_type, language = self._map_label(raw_label) if confidence < self.min_confidence: diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index d383002c6..cbc597aa7 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -1995,36 +1995,32 @@ class HeadroomProxy: else: logger.info("Smart Routing: DISABLED (legacy sequential mode)") - # Eagerly load ML compressors at startup (avoids download on first request) - # Kompress requires [ml] extra (torch + transformers). If not installed, skip. + # Eagerly load ALL compressors, parsers, and detectors at startup + # This eliminates cold-start latency spikes on first requests self._kompress_status = "not installed" - from headroom.transforms.kompress_compressor import is_kompress_available + eager_status: dict[str, str] = {} - if is_kompress_available() and self.config.optimize: - logger.info("Kompress: Downloading model (first-time only)...") + if self.config.optimize: + logger.info("Pre-loading compressors and parsers...") for transform in self.anthropic_pipeline.transforms: if hasattr(transform, "eager_load_compressors"): - transform.eager_load_compressors() - self._kompress_status = "enabled" - break - if self._kompress_status == "enabled": - logger.info("Kompress: ENABLED (ModernBERT token compressor)") - else: - if self.config.optimize: - logger.info( - "Kompress: not installed (pip install headroom-ai[ml] for ML compression)" - ) - - # LLMLingua fallback (only loads if Kompress is not available) - if self._kompress_status != "enabled" and self.config.llmlingua_enabled: - for transform in self.anthropic_pipeline.transforms: - if hasattr(transform, "_get_llmlingua"): - llmlingua = transform._get_llmlingua() - if llmlingua: - self._llmlingua_status = "enabled" + eager_status = transform.eager_load_compressors() break - # LLMLingua status + # Update internal status from eager loading results + if eager_status.get("kompress") == "enabled": + self._kompress_status = "enabled" + if eager_status.get("llmlingua") == "enabled": + self._llmlingua_status = "enabled" + if eager_status.get("code_aware") == "enabled": + self._code_aware_status = "enabled" + + # Log component status + if self._kompress_status == "enabled": + logger.info("Kompress: ENABLED (ModernBERT token compressor)") + elif self.config.optimize: + logger.info("Kompress: not installed (pip install headroom-ai[ml] for ML compression)") + if self._llmlingua_status == "enabled": logger.info( f"LLMLingua: ENABLED (device={self.config.llmlingua_device}, " @@ -2037,9 +2033,10 @@ class HeadroomProxy: elif self._llmlingua_status == "disabled": logger.info("LLMLingua: DISABLED") - # Code-aware status if self._code_aware_status == "enabled": logger.info("Code-Aware: ENABLED (AST-based compression)") + if "tree_sitter" in eager_status: + logger.info(f"Tree-Sitter: {eager_status['tree_sitter']}") elif self._code_aware_status == "lazy": logger.info("Code-Aware: LAZY (will load when code content detected)") elif self._code_aware_status == "available": @@ -2049,6 +2046,9 @@ class HeadroomProxy: elif self._code_aware_status == "disabled": logger.info("Code-Aware: DISABLED") + if eager_status.get("magika") == "enabled": + logger.info("Magika: ENABLED (ML content detection)") + # CCR status ccr_features = [] if self.config.ccr_inject_tool: diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index 908fe2dae..297edd00a 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -428,6 +428,7 @@ class ContentRouterConfig: _CODE_FENCE_PATTERN = re.compile(r"^```(\w*)\s*$", re.MULTILINE) _JSON_BLOCK_START = re.compile(r"^\s*[\[{]", re.MULTILINE) _SEARCH_RESULT_PATTERN = re.compile(r"^\S+:\d+:", re.MULTILINE) +_PROSE_PATTERN = re.compile(r"[A-Z][a-z]+\s+\w+\s+\w+") def is_mixed_content(content: str) -> bool: @@ -442,7 +443,7 @@ def is_mixed_content(content: str) -> bool: indicators = { "has_code_fences": bool(_CODE_FENCE_PATTERN.search(content)), "has_json_blocks": bool(_JSON_BLOCK_START.search(content)), - "has_prose": len(re.findall(r"[A-Z][a-z]+\s+\w+\s+\w+", content)) > 5, + "has_prose": len(_PROSE_PATTERN.findall(content)) > 5, "has_search_results": bool(_SEARCH_RESULT_PATTERN.search(content)), } @@ -1168,30 +1169,97 @@ class ContentRouter(Transform): logger.debug("HTMLExtractor not available (install trafilatura)") return self._html_extractor - def eager_load_compressors(self) -> None: + def eager_load_compressors(self) -> dict[str, str]: """Pre-load compressors at startup to avoid first-request latency. - Call this during proxy startup to load models (~5s) - before any requests arrive. + Call this during proxy startup to load models and parsers + before any requests arrive. Eliminates cold-start latency spikes. + + Returns: + Dict of component name -> status string for logging. """ - # Prefer Kompress (faster, smaller, better on structured data) + status: dict[str, str] = {} + + # 1. ML text compressor: Kompress or LLMLingua fallback if self.config.enable_kompress: compressor = self._get_kompress() if compressor: logger.info("Kompress model pre-loaded at startup") - return # No need to also load LLMLingua + status["kompress"] = "enabled" + else: + status["kompress"] = "unavailable" + if "kompress" not in status or status["kompress"] != "enabled": + if self.config.enable_llmlingua: + compressor = self._get_llmlingua() + if compressor: + try: + from .llmlingua_compressor import _get_llmlingua_compressor - if self.config.enable_llmlingua: - compressor = self._get_llmlingua() - if compressor: + device = compressor._resolve_device() + _get_llmlingua_compressor(compressor.config.model_name, device) + logger.info("LLMLingua model pre-loaded at startup") + status["llmlingua"] = "enabled" + except Exception as e: + logger.warning("Failed to pre-load LLMLingua model: %s", e) + status["llmlingua"] = f"failed: {e}" + + # 2. Magika content detector (avoids 100-200ms on first content detection) + try: + from ..compression.detector import _get_magika, _magika_available + + if _magika_available(): + _get_magika() # Initializes the singleton + logger.info("Magika content detector pre-loaded at startup") + status["magika"] = "enabled" + else: + status["magika"] = "not installed" + except Exception as e: + logger.debug("Magika pre-load skipped: %s", e) + status["magika"] = "skipped" + + # 3. CodeAware compressor + common tree-sitter parsers + if self.config.enable_code_aware: + code_compressor = self._get_code_compressor() + if code_compressor: + status["code_aware"] = "enabled" + # Pre-load tree-sitter parsers for common languages + # Each parser is ~50ms to load; doing it here avoids 500ms+ on first code hit try: - from .llmlingua_compressor import _get_llmlingua_compressor + from .code_compressor import _check_tree_sitter_available, _get_parser - device = compressor._resolve_device() - _get_llmlingua_compressor(compressor.config.model_name, device) - logger.info("LLMLingua model pre-loaded at startup") + if _check_tree_sitter_available(): + common_languages = [ + "python", + "javascript", + "typescript", + "go", + "rust", + "java", + "c", + "cpp", + ] + loaded = [] + for lang in common_languages: + try: + _get_parser(lang) + loaded.append(lang) + except (ValueError, ImportError): + pass # Language not available, skip + if loaded: + logger.info("Tree-sitter parsers pre-loaded: %s", ", ".join(loaded)) + status["tree_sitter"] = f"loaded ({len(loaded)} languages)" except Exception as e: - logger.warning("Failed to pre-load LLMLingua model: %s", e) + logger.debug("Tree-sitter pre-load skipped: %s", e) + status["tree_sitter"] = "skipped" + else: + status["code_aware"] = "not installed" + + # 4. SmartCrusher (lightweight init, but ensures import + TOIN ready) + smart_crusher = self._get_smart_crusher() + if smart_crusher: + status["smart_crusher"] = "ready" + + return status def _get_kompress(self) -> Any: """Get KompressCompressor (lazy load). Downloads from HuggingFace on first use.""" diff --git a/headroom/transforms/smart_crusher.py b/headroom/transforms/smart_crusher.py index d6bfcaccc..012e55fd9 100644 --- a/headroom/transforms/smart_crusher.py +++ b/headroom/transforms/smart_crusher.py @@ -92,6 +92,10 @@ _HOSTNAME_PATTERN = re.compile( _QUOTED_STRING_PATTERN = re.compile(r"['\"]([^'\"]{1,50})['\"]") # Short quoted strings _EMAIL_PATTERN = re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b") +# Temporal detection patterns (compiled once, used in SmartAnalyzer._detect_temporal_field) +_ISO_DATETIME_PATTERN = re.compile(r"^\d{4}-\d{2}-\d{2}[T ]\d{2}:\d{2}:\d{2}") +_ISO_DATE_PATTERN = re.compile(r"^\d{4}-\d{2}-\d{2}$") + def extract_query_anchors(text: str) -> set[str]: """Extract query anchors from user text (legacy regex-based method). @@ -628,7 +632,10 @@ def _detect_rare_status_values(items: list[dict], common_fields: set[str]) -> li _ERROR_KEYWORDS_FOR_PRESERVATION = ERROR_KEYWORDS -def _detect_error_items_for_preservation(items: list[dict]) -> list[int]: +def _detect_error_items_for_preservation( + items: list[dict], + item_strings: list[str] | None = None, +) -> list[int]: """Detect items containing error keywords for PRESERVATION guarantee. This is NOT for crushability analysis - it's for ensuring ALL error items @@ -636,6 +643,10 @@ def _detect_error_items_for_preservation(items: list[dict]) -> list[int]: are NEVER dropped, even if errors are common in the dataset. Uses keywords because error semantics are well-defined across domains. + + Args: + items: List of items to check. + item_strings: Pre-computed JSON serializations to avoid redundant json.dumps. """ error_indices: list[int] = [] @@ -643,9 +654,12 @@ def _detect_error_items_for_preservation(items: list[dict]) -> list[int]: if not isinstance(item, dict): continue - # Serialize item to check all content + # Reuse cached serialization if available, otherwise serialize try: - item_str = json.dumps(item).lower() + if item_strings is not None and i < len(item_strings): + item_str = item_strings[i].lower() + else: + item_str = json.dumps(item).lower() except Exception: continue @@ -694,13 +708,18 @@ def _detect_items_by_learned_semantics( if not confident_semantics: return [] + # Pre-compute field name hashes to avoid redundant SHA256 per item + _field_hash_cache: dict[str, str] = {} + for i, item in enumerate(items): if not isinstance(item, dict): continue for field_name, value in item.items(): - # Hash the field name to match TOIN's format - field_hash = hashlib.sha256(field_name.encode()).hexdigest()[:8] + # Hash the field name to match TOIN's format (cached per unique field name) + if field_name not in _field_hash_cache: + _field_hash_cache[field_name] = _hash_field_name(field_name) + field_hash = _field_hash_cache[field_name] if field_hash not in confident_semantics: continue @@ -1080,9 +1099,9 @@ class SmartAnalyzer: Uses STRUCTURAL detection based on value format, not field names. """ - # Check string fields for ISO 8601 patterns - iso_datetime_pattern = re.compile(r"^\d{4}-\d{2}-\d{2}[T ]\d{2}:\d{2}:\d{2}") - iso_date_pattern = re.compile(r"^\d{4}-\d{2}-\d{2}$") + # Check string fields for ISO 8601 patterns (module-level compiled) + iso_datetime_pattern = _ISO_DATETIME_PATTERN + iso_date_pattern = _ISO_DATE_PATTERN for name, stats in field_stats.items(): if stats.field_type == "string": @@ -2464,12 +2483,14 @@ class SmartCrusher(Transform): # Create compression plan with relevance scoring # Pass TOIN preserve_fields so items with those fields get priority # Pass effective_max_items for thread-safe compression + # Pass item_strings to avoid redundant json.dumps across plan methods plan = self._create_plan( analysis, items, query_context, preserve_fields=toin_preserve_fields or None, effective_max_items=effective_max_items, + item_strings=item_strings, ) # Execute compression @@ -2483,7 +2504,8 @@ class SmartCrusher(Transform): and len(result) < len(items) # Only cache if compression actually happened ): store = self._get_compression_store() - original_json = json.dumps(items, default=str) + # Reuse cached item_strings to avoid re-serializing + original_json = "[" + ", ".join(item_strings) + "]" compressed_json = json.dumps(result, default=str) ccr_hash = store.store( @@ -2521,8 +2543,8 @@ class SmartCrusher(Transform): # TOIN: Record compression event for cross-user learning try: - # Calculate token counts (approximate) - original_tokens = len(json.dumps(items, default=str)) // 4 + # Calculate token counts (approximate) - reuse cached item_strings + original_tokens = sum(len(s) for s in item_strings) // 4 compressed_tokens = len(json.dumps(result, default=str)) // 4 toin.record_compression( @@ -2573,18 +2595,29 @@ class SmartCrusher(Transform): # Universal JSON type handlers (string, number, mixed arrays) # ================================================================= - def _compute_k_split(self, items: list, bias: float = 1.0) -> tuple[int, int, int, int]: + def _compute_k_split( + self, + items: list, + bias: float = 1.0, + item_strings: list[str] | None = None, + ) -> tuple[int, int, int, int]: """Compute adaptive K split into first/last/importance slots. Uses the existing Kneedle-based adaptive_sizer for K_total, then splits according to configurable first_fraction / last_fraction. + Args: + items: List of items (used as fallback for serialization). + bias: Compression bias multiplier. + item_strings: Pre-computed JSON serializations to avoid redundant json.dumps. + Returns: (k_total, k_first, k_last, k_importance) """ from .adaptive_sizer import compute_optimal_k - item_strings = [json.dumps(item, default=str) for item in items] + if item_strings is None: + item_strings = [json.dumps(item, default=str) for item in items] k_total = compute_optimal_k( item_strings, bias=bias, @@ -2997,6 +3030,7 @@ class SmartCrusher(Transform): query_context: str = "", preserve_fields: list[str] | None = None, effective_max_items: int | None = None, + item_strings: list[str] | None = None, ) -> CompressionPlan: """Create a detailed compression plan using relevance scoring. @@ -3006,6 +3040,7 @@ class SmartCrusher(Transform): query_context: Context string from user messages for relevance scoring. preserve_fields: TOIN-learned fields that users commonly retrieve. Items with values in these fields get higher priority. + item_strings: Pre-computed JSON serializations to avoid redundant json.dumps. effective_max_items: Thread-safe max items limit (defaults to config value). """ # Use provided effective_max_items or fall back to config @@ -3027,22 +3062,46 @@ class SmartCrusher(Transform): if analysis.recommended_strategy == CompressionStrategy.TIME_SERIES: plan = self._plan_time_series( - analysis, items, plan, query_context, preserve_fields, max_items + analysis, + items, + plan, + query_context, + preserve_fields, + max_items, + item_strings=item_strings, ) elif analysis.recommended_strategy == CompressionStrategy.CLUSTER_SAMPLE: plan = self._plan_cluster_sample( - analysis, items, plan, query_context, preserve_fields, max_items + analysis, + items, + plan, + query_context, + preserve_fields, + max_items, + item_strings=item_strings, ) elif analysis.recommended_strategy == CompressionStrategy.TOP_N: plan = self._plan_top_n( - analysis, items, plan, query_context, preserve_fields, max_items + analysis, + items, + plan, + query_context, + preserve_fields, + max_items, + item_strings=item_strings, ) else: # SMART_SAMPLE or NONE plan = self._plan_smart_sample( - analysis, items, plan, query_context, preserve_fields, max_items + analysis, + items, + plan, + query_context, + preserve_fields, + max_items, + item_strings=item_strings, ) return plan @@ -3055,6 +3114,7 @@ class SmartCrusher(Transform): query_context: str = "", preserve_fields: list[str] | None = None, max_items: int | None = None, + item_strings: list[str] | None = None, ) -> CompressionPlan: """Plan compression for time series data. @@ -3111,7 +3171,12 @@ class SmartCrusher(Transform): # 5. Items with high relevance to query context (PROBABILISTIC semantic match) if query_context: - item_strs = [json.dumps(item, default=str) for item in items] + # Reuse pre-computed item_strings if available + item_strs = ( + item_strings + if item_strings is not None + else [json.dumps(item, default=str) for item in items] + ) scores = self._scorer.score_batch(item_strs, query_context) for i, score in enumerate(scores): if score.score >= self._relevance_threshold: @@ -3138,6 +3203,7 @@ class SmartCrusher(Transform): query_context: str = "", preserve_fields: list[str] | None = None, max_items: int | None = None, + item_strings: list[str] | None = None, ) -> CompressionPlan: """Plan compression for clusterable data (like logs). @@ -3211,7 +3277,12 @@ class SmartCrusher(Transform): # 5. Items with high relevance to query context (PROBABILISTIC semantic match) if query_context: - item_strs = [json.dumps(item, default=str) for item in items] + # Reuse pre-computed item_strings if available + item_strs = ( + item_strings + if item_strings is not None + else [json.dumps(item, default=str) for item in items] + ) scores = self._scorer.score_batch(item_strs, query_context) for i, score in enumerate(scores): if score.score >= self._relevance_threshold: @@ -3238,6 +3309,7 @@ class SmartCrusher(Transform): query_context: str = "", preserve_fields: list[str] | None = None, max_items: int | None = None, + item_strings: list[str] | None = None, ) -> CompressionPlan: """Plan compression for scored/ranked data. @@ -3269,7 +3341,13 @@ class SmartCrusher(Transform): if not score_field: return self._plan_smart_sample( - analysis, items, plan, query_context, preserve_fields, effective_max + analysis, + items, + plan, + query_context, + preserve_fields, + effective_max, + item_strings=item_strings, ) plan.sort_field = score_field @@ -3307,7 +3385,12 @@ class SmartCrusher(Transform): # Only add items that are NOT already in top N but match the query strongly # Use a higher threshold (0.5) since the score field already captures relevance if query_context: - item_strs = [json.dumps(item, default=str) for item in items] + # Reuse pre-computed item_strings if available + item_strs = ( + item_strings + if item_strings is not None + else [json.dumps(item, default=str) for item in items] + ) scores = self._scorer.score_batch(item_strs, query_context) # Higher threshold and limit count to avoid adding everything high_threshold = max(0.5, self._relevance_threshold * 2) @@ -3340,6 +3423,7 @@ class SmartCrusher(Transform): query_context: str = "", preserve_fields: list[str] | None = None, max_items: int | None = None, + item_strings: list[str] | None = None, ) -> CompressionPlan: """Plan smart statistical sampling using STATISTICAL detection. @@ -3415,7 +3499,12 @@ class SmartCrusher(Transform): # 6. Items with high relevance to query context (PROBABILISTIC semantic match) if query_context: - item_strs = [json.dumps(item, default=str) for item in items] + # Reuse pre-computed item_strings if available + item_strs = ( + item_strings + if item_strings is not None + else [json.dumps(item, default=str) for item in items] + ) scores = self._scorer.score_batch(item_strs, query_context) for i, score in enumerate(scores): if score.score >= self._relevance_threshold: diff --git a/pyproject.toml b/pyproject.toml index a41e647b8..bd161ac97 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "headroom-ai" -version = "0.5.4" +version = "0.5.5" description = "The Context Optimization Layer for LLM Applications - Cut costs by 50-90%" readme = "README.md" license = "Apache-2.0" @@ -60,6 +60,7 @@ proxy = [ "httpx[http2]>=0.24.0", "openai>=2.14.0", # OpenAI API format support "mcp>=1.0.0", # MCP server (headroom_compress, retrieve, stats) + "magika>=0.6.0", # ML content detection for ContentRouter ] # AST-based code compression (tree-sitter) code = [