headroom/transforms/tool_crusher.py

"""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)