headroom-ai 0.2.13__py3-none-any.whl

headroom/transforms/text_compressor.py

@@ -0,0 +1,259 @@
+"""Generic text compressor for plain text content.
+
+This module provides a fallback compressor for plain text that doesn't match
+any specialized format (search results, logs, code, diffs). Uses line-based
+sampling with anchor preservation.
+
+Compression Strategy:
+1. Identify anchor lines (contain context keywords)
+2. Keep first N and last M lines
+3. Sample from middle based on line importance
+4. Add summary of omitted content
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+
+
+@dataclass
+class TextCompressorConfig:
+    """Configuration for text compression."""
+
+    # Line limits
+    keep_first_lines: int = 10
+    keep_last_lines: int = 10
+    max_total_lines: int = 50
+
+    # Sampling
+    sample_every_n_lines: int = 10
+
+    # Anchor detection
+    anchor_keywords: list[str] = field(default_factory=list)
+    boost_pattern_lines: bool = True
+
+    # CCR integration
+    enable_ccr: bool = True
+    min_lines_for_ccr: int = 100
+
+
+class TextCompressor:
+    """Compresses generic plain text.
+
+    Example:
+        >>> compressor = TextCompressor()
+        >>> result = compressor.compress(large_text, context="find errors")
+        >>> print(result.compressed)
+    """
+
+    # Patterns that indicate important lines
+    _IMPORTANT_PATTERNS = [
+        re.compile(r"\b(error|exception|fail|warning)\b", re.IGNORECASE),
+        re.compile(r"\b(important|note|todo|fixme)\b", re.IGNORECASE),
+        re.compile(r"^#+\s"),  # Markdown headers
+        re.compile(r"^\*\*"),  # Bold text
+        re.compile(r"^>\s"),  # Quotes
+    ]
+
+    def __init__(self, config: TextCompressorConfig | None = None):
+        """Initialize text compressor.
+
+        Args:
+            config: Compression configuration.
+        """
+        self.config = config or TextCompressorConfig()
+
+    def compress(self, content: str, context: str = "") -> TextCompressionResult:
+        """Compress text content.
+
+        Args:
+            content: Raw text content.
+            context: User query context for anchor detection.
+
+        Returns:
+            TextCompressionResult with compressed output.
+        """
+        lines = content.split("\n")
+
+        if len(lines) <= self.config.max_total_lines:
+            return TextCompressionResult(
+                compressed=content,
+                original=content,
+                original_line_count=len(lines),
+                compressed_line_count=len(lines),
+                compression_ratio=1.0,
+            )
+
+        # Score lines by importance
+        scored_lines = self._score_lines(lines, context)
+
+        # Select lines
+        selected = self._select_lines(scored_lines, lines)
+
+        # Format output
+        compressed = self._format_output(selected, len(lines))
+
+        ratio = len(compressed) / max(len(content), 1)
+
+        # Store in CCR if significant compression
+        cache_key = None
+        if self.config.enable_ccr and len(lines) >= self.config.min_lines_for_ccr and ratio < 0.7:
+            cache_key = self._store_in_ccr(content, compressed, len(lines))
+            if cache_key:
+                compressed += f"\n[{len(lines)} lines compressed. hash={cache_key}]"
+
+        return TextCompressionResult(
+            compressed=compressed,
+            original=content,
+            original_line_count=len(lines),
+            compressed_line_count=len(selected),
+            compression_ratio=ratio,
+            cache_key=cache_key,
+        )
+
+    def _score_lines(self, lines: list[str], context: str) -> list[tuple[int, str, float]]:
+        """Score lines by importance."""
+        context_lower = context.lower()
+        context_words = set(context_lower.split()) if context else set()
+        anchor_keywords = {k.lower() for k in self.config.anchor_keywords}
+
+        scored: list[tuple[int, str, float]] = []
+
+        for i, line in enumerate(lines):
+            score = 0.0
+            line_lower = line.lower()
+
+            # Boost if contains context words
+            for word in context_words:
+                if len(word) > 2 and word in line_lower:
+                    score += 0.3
+
+            # Boost if contains anchor keywords
+            for keyword in anchor_keywords:
+                if keyword in line_lower:
+                    score += 0.4
+
+            # Boost if matches important patterns
+            if self.config.boost_pattern_lines:
+                for pattern in self._IMPORTANT_PATTERNS:
+                    if pattern.search(line):
+                        score += 0.2
+                        break
+
+            # Small boost for non-empty lines
+            if line.strip():
+                score += 0.1
+
+            scored.append((i, line, min(1.0, score)))
+
+        return scored
+
+    def _select_lines(
+        self, scored_lines: list[tuple[int, str, float]], original_lines: list[str]
+    ) -> list[tuple[int, str]]:
+        """Select lines to keep."""
+        total = len(scored_lines)
+        selected_indices: set[int] = set()
+
+        # Always keep first N lines
+        for i in range(min(self.config.keep_first_lines, total)):
+            selected_indices.add(i)
+
+        # Always keep last M lines
+        for i in range(max(0, total - self.config.keep_last_lines), total):
+            selected_indices.add(i)
+
+        # Add high-scoring lines
+        high_score_lines = [
+            (idx, line, score)
+            for idx, line, score in scored_lines
+            if score >= 0.3 and idx not in selected_indices
+        ]
+        high_score_lines.sort(key=lambda x: x[2], reverse=True)
+
+        remaining_slots = self.config.max_total_lines - len(selected_indices)
+        for idx, _line, _score in high_score_lines[:remaining_slots]:
+            selected_indices.add(idx)
+            remaining_slots -= 1
+            if remaining_slots <= 0:
+                break
+
+        # Sample from remaining middle lines
+        if remaining_slots > 0:
+            middle_start = self.config.keep_first_lines
+            middle_end = total - self.config.keep_last_lines
+
+            for i in range(middle_start, middle_end, self.config.sample_every_n_lines):
+                if i not in selected_indices:
+                    selected_indices.add(i)
+                    remaining_slots -= 1
+                    if remaining_slots <= 0:
+                        break
+
+        # Sort by line number and return
+        selected = sorted(selected_indices)
+        return [(i, original_lines[i]) for i in selected]
+
+    def _format_output(self, selected: list[tuple[int, str]], total_lines: int) -> str:
+        """Format selected lines with ellipsis markers."""
+        if not selected:
+            return f"[{total_lines} lines omitted]"
+
+        output_lines: list[str] = []
+        prev_idx = -1
+
+        for idx, line in selected:
+            # Add ellipsis if there's a gap
+            if prev_idx >= 0 and idx - prev_idx > 1:
+                gap = idx - prev_idx - 1
+                output_lines.append(f"[... {gap} lines omitted ...]")
+
+            output_lines.append(line)
+            prev_idx = idx
+
+        # Add trailing ellipsis if needed
+        if selected and selected[-1][0] < total_lines - 1:
+            gap = total_lines - selected[-1][0] - 1
+            output_lines.append(f"[... {gap} lines omitted ...]")
+
+        return "\n".join(output_lines)
+
+    def _store_in_ccr(self, original: str, compressed: str, original_count: int) -> str | None:
+        """Store original in CCR for later retrieval."""
+        try:
+            from ..cache.compression_store import get_compression_store
+
+            store = get_compression_store()
+            return store.store(
+                original,
+                compressed,
+                original_item_count=original_count,
+            )
+        except ImportError:
+            return None
+        except Exception:
+            return None
+
+
+@dataclass
+class TextCompressionResult:
+    """Result of text compression."""
+
+    compressed: str
+    original: str
+    original_line_count: int
+    compressed_line_count: int
+    compression_ratio: float
+    cache_key: str | None = None
+
+    @property
+    def tokens_saved_estimate(self) -> int:
+        """Estimate tokens saved."""
+        chars_saved = len(self.original) - len(self.compressed)
+        return max(0, chars_saved // 4)
+
+    @property
+    def lines_omitted(self) -> int:
+        """Number of lines omitted."""
+        return self.original_line_count - self.compressed_line_count

headroom/transforms/tool_crusher.py

@@ -0,0 +1,338 @@
+"""Tool output compression transform for Headroom SDK."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from ..config import ToolCrusherConfig, TransformResult
+from ..tokenizer import Tokenizer
+from ..utils import (
+    compute_short_hash,
+    create_tool_digest_marker,
+    deep_copy_messages,
+    safe_json_dumps,
+    safe_json_loads,
+)
+from .base import Transform
+
+logger = logging.getLogger(__name__)
+
+
+class ToolCrusher(Transform):
+    """
+    Compress tool output to reduce token usage.
+
+    This transform applies conservative compression:
+    - Only compresses tool role messages > min_tokens
+    - Preserves JSON structure (never removes keys)
+    - Truncates arrays to max_items
+    - Truncates long strings
+    - Limits nesting depth
+
+    Safety: If JSON parsing fails, content is returned unchanged.
+    """
+
+    name = "tool_crusher"
+
+    def __init__(self, config: ToolCrusherConfig | None = None):
+        """
+        Initialize tool crusher.
+
+        Args:
+            config: Configuration for compression behavior.
+        """
+        self.config = config or ToolCrusherConfig()
+
+    def should_apply(
+        self,
+        messages: list[dict[str, Any]],
+        tokenizer: Tokenizer,
+        **kwargs: Any,
+    ) -> bool:
+        """Check if any tool messages exceed threshold."""
+        if not self.config.enabled:
+            return False
+
+        for msg in messages:
+            # OpenAI style: role="tool"
+            if msg.get("role") == "tool":
+                content = msg.get("content", "")
+                if isinstance(content, str):
+                    tokens = tokenizer.count_text(content)
+                    if tokens > self.config.min_tokens_to_crush:
+                        return True
+
+            # Anthropic style: role="user" with tool_result content blocks
+            content = msg.get("content")
+            if isinstance(content, list):
+                for block in content:
+                    if isinstance(block, dict) and block.get("type") == "tool_result":
+                        tool_content = block.get("content", "")
+                        if isinstance(tool_content, str):
+                            tokens = tokenizer.count_text(tool_content)
+                            if tokens > self.config.min_tokens_to_crush:
+                                return True
+
+        return False
+
+    def apply(
+        self,
+        messages: list[dict[str, Any]],
+        tokenizer: Tokenizer,
+        **kwargs: Any,
+    ) -> TransformResult:
+        """
+        Apply tool crushing to messages.
+
+        Args:
+            messages: List of messages.
+            tokenizer: Tokenizer for counting.
+            **kwargs: May include 'tool_profiles' for per-tool config.
+
+        Returns:
+            TransformResult with crushed messages.
+        """
+        tool_profiles = kwargs.get("tool_profiles", self.config.tool_profiles)
+
+        tokens_before = tokenizer.count_messages(messages)
+        result_messages = deep_copy_messages(messages)
+        transforms_applied: list[str] = []
+        markers_inserted: list[str] = []
+        warnings: list[str] = []
+
+        crushed_count = 0
+
+        for msg in result_messages:
+            # OpenAI style: role="tool"
+            if msg.get("role") == "tool":
+                content = msg.get("content", "")
+                if not isinstance(content, str):
+                    continue
+
+                # Check token threshold
+                tokens = tokenizer.count_text(content)
+                if tokens <= self.config.min_tokens_to_crush:
+                    continue
+
+                # Get tool-specific profile if available
+                tool_call_id = msg.get("tool_call_id", "")
+                profile = self._get_profile(tool_call_id, tool_profiles)
+
+                # Try to crush
+                crushed, was_modified = self._crush_content(content, profile)
+
+                if was_modified:
+                    # Compute hash of original for marker
+                    original_hash = compute_short_hash(content)
+                    marker = create_tool_digest_marker(original_hash)
+
+                    msg["content"] = crushed + "\n" + marker
+                    crushed_count += 1
+                    markers_inserted.append(marker)
+
+            # Anthropic style: role="user" with tool_result content blocks
+            content = msg.get("content")
+            if isinstance(content, list):
+                for i, block in enumerate(content):
+                    if not isinstance(block, dict):
+                        continue
+                    if block.get("type") != "tool_result":
+                        continue
+
+                    tool_content = block.get("content", "")
+                    if not isinstance(tool_content, str):
+                        continue
+
+                    # Check token threshold
+                    tokens = tokenizer.count_text(tool_content)
+                    if tokens <= self.config.min_tokens_to_crush:
+                        continue
+
+                    # Get tool-specific profile if available
+                    tool_use_id = block.get("tool_use_id", "")
+                    profile = self._get_profile(tool_use_id, tool_profiles)
+
+                    # Try to crush
+                    crushed, was_modified = self._crush_content(tool_content, profile)
+
+                    if was_modified:
+                        # Compute hash of original for marker
+                        original_hash = compute_short_hash(tool_content)
+                        marker = create_tool_digest_marker(original_hash)
+
+                        # Update the content block
+                        content[i]["content"] = crushed + "\n" + marker
+                        crushed_count += 1
+                        markers_inserted.append(marker)
+
+        if crushed_count > 0:
+            transforms_applied.append(f"tool_crush:{crushed_count}")
+            logger.info(
+                "ToolCrusher: compressed %d tool outputs, %d -> %d tokens",
+                crushed_count,
+                tokens_before,
+                tokenizer.count_messages(result_messages),
+            )
+
+        tokens_after = tokenizer.count_messages(result_messages)
+
+        return TransformResult(
+            messages=result_messages,
+            tokens_before=tokens_before,
+            tokens_after=tokens_after,
+            transforms_applied=transforms_applied,
+            markers_inserted=markers_inserted,
+            warnings=warnings,
+        )
+
+    def _get_profile(
+        self,
+        tool_call_id: str,
+        tool_profiles: dict[str, dict[str, Any]],
+    ) -> dict[str, Any]:
+        """Get compression profile for a tool."""
+        # Tool profiles are keyed by tool name, not call ID
+        # For now, use default config
+        # In a real implementation, you'd map call_id -> tool_name
+        return {
+            "max_array_items": self.config.max_array_items,
+            "max_string_length": self.config.max_string_length,
+            "max_depth": self.config.max_depth,
+            "preserve_keys": self.config.preserve_keys,
+        }
+
+    def _crush_content(
+        self,
+        content: str,
+        profile: dict[str, Any],
+    ) -> tuple[str, bool]:
+        """
+        Crush content according to profile.
+
+        Returns:
+            Tuple of (crushed_content, was_modified).
+            If parsing fails, returns (original_content, False).
+        """
+        # Try JSON parse
+        parsed, success = safe_json_loads(content)
+        if not success:
+            # Safety: don't modify unparseable content
+            return content, False
+
+        # Apply crushing
+        crushed = self._crush_value(
+            parsed,
+            depth=0,
+            max_depth=profile.get("max_depth", 5),
+            max_array_items=profile.get("max_array_items", 10),
+            max_string_length=profile.get("max_string_length", 1000),
+        )
+
+        # Serialize back
+        result = safe_json_dumps(crushed, indent=None)
+
+        # Check if actually modified
+        was_modified = result != content.strip()
+
+        return result, was_modified
+
+    def _crush_value(
+        self,
+        value: Any,
+        depth: int,
+        max_depth: int,
+        max_array_items: int,
+        max_string_length: int,
+    ) -> Any:
+        """Recursively crush a value."""
+        if depth >= max_depth:
+            # At max depth, summarize
+            if isinstance(value, dict):
+                return {"__headroom_depth_exceeded": len(value)}
+            elif isinstance(value, list):
+                return {"__headroom_depth_exceeded": len(value)}
+            elif isinstance(value, str) and len(value) > max_string_length:
+                return (
+                    value[:max_string_length]
+                    + f"...[truncated {len(value) - max_string_length} chars]"
+                )
+            return value
+
+        if isinstance(value, dict):
+            return {
+                k: self._crush_value(
+                    v,
+                    depth + 1,
+                    max_depth,
+                    max_array_items,
+                    max_string_length,
+                )
+                for k, v in value.items()
+            }
+
+        elif isinstance(value, list):
+            if len(value) <= max_array_items:
+                return [
+                    self._crush_value(
+                        item,
+                        depth + 1,
+                        max_depth,
+                        max_array_items,
+                        max_string_length,
+                    )
+                    for item in value
+                ]
+            else:
+                # Truncate array
+                truncated = [
+                    self._crush_value(
+                        item,
+                        depth + 1,
+                        max_depth,
+                        max_array_items,
+                        max_string_length,
+                    )
+                    for item in value[:max_array_items]
+                ]
+                truncated.append({"__headroom_truncated": len(value) - max_array_items})
+                return truncated
+
+        elif isinstance(value, str):
+            if len(value) > max_string_length:
+                return (
+                    value[:max_string_length]
+                    + f"...[truncated {len(value) - max_string_length} chars]"
+                )
+            return value
+
+        else:
+            # Numbers, bools, None - pass through
+            return value
+
+
+def crush_tool_output(
+    content: str,
+    config: ToolCrusherConfig | None = None,
+) -> tuple[str, bool]:
+    """
+    Convenience function to crush a single tool output.
+
+    Args:
+        content: The tool output content.
+        config: Optional configuration.
+
+    Returns:
+        Tuple of (crushed_content, was_modified).
+    """
+    cfg = config or ToolCrusherConfig()
+    crusher = ToolCrusher(cfg)
+
+    profile = {
+        "max_array_items": cfg.max_array_items,
+        "max_string_length": cfg.max_string_length,
+        "max_depth": cfg.max_depth,
+        "preserve_keys": cfg.preserve_keys,
+    }
+
+    return crusher._crush_content(content, profile)