headroom-ai 0.2.13__py3-none-any.whl

headroom/transforms/pipeline.py

@@ -0,0 +1,297 @@
+"""Transform pipeline orchestration for Headroom SDK."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from ..config import (
+    CacheAlignerConfig,
+    DiffArtifact,
+    HeadroomConfig,
+    RollingWindowConfig,
+    ToolCrusherConfig,
+    TransformDiff,
+    TransformResult,
+)
+from ..tokenizer import Tokenizer
+from ..utils import deep_copy_messages
+from .base import Transform
+from .cache_aligner import CacheAligner
+from .rolling_window import RollingWindow
+from .smart_crusher import SmartCrusher
+from .tool_crusher import ToolCrusher
+
+if TYPE_CHECKING:
+    from ..providers.base import Provider
+
+logger = logging.getLogger(__name__)
+
+
+class TransformPipeline:
+    """
+    Orchestrates multiple transforms in the correct order.
+
+    Transform order:
+    1. Cache Aligner - normalize prefix for cache hits
+    2. Tool Crusher - compress tool outputs
+    3. Rolling Window - enforce token limits
+    """
+
+    def __init__(
+        self,
+        config: HeadroomConfig | None = None,
+        transforms: list[Transform] | None = None,
+        provider: Provider | None = None,
+    ):
+        """
+        Initialize pipeline.
+
+        Args:
+            config: Headroom configuration.
+            transforms: Optional custom transform list (overrides config).
+            provider: Provider for model-specific behavior.
+        """
+        self.config = config or HeadroomConfig()
+        self._provider = provider
+
+        if transforms is not None:
+            self.transforms = transforms
+        else:
+            self.transforms = self._build_default_transforms()
+
+    def _build_default_transforms(self) -> list[Transform]:
+        """Build default transform pipeline from config."""
+        transforms: list[Transform] = []
+
+        # Order matters!
+
+        # 1. Cache Aligner (prefix stabilization)
+        if self.config.cache_aligner.enabled:
+            transforms.append(CacheAligner(self.config.cache_aligner))
+
+        # 2. Tool Output Compression
+        # SmartCrusher (statistical) takes precedence over ToolCrusher (fixed rules)
+        if self.config.smart_crusher.enabled:
+            # Use smart statistical crushing
+            from .smart_crusher import SmartCrusherConfig as SCConfig
+
+            smart_config = SCConfig(
+                enabled=True,
+                min_items_to_analyze=self.config.smart_crusher.min_items_to_analyze,
+                min_tokens_to_crush=self.config.smart_crusher.min_tokens_to_crush,
+                variance_threshold=self.config.smart_crusher.variance_threshold,
+                uniqueness_threshold=self.config.smart_crusher.uniqueness_threshold,
+                similarity_threshold=self.config.smart_crusher.similarity_threshold,
+                max_items_after_crush=self.config.smart_crusher.max_items_after_crush,
+                preserve_change_points=self.config.smart_crusher.preserve_change_points,
+                factor_out_constants=self.config.smart_crusher.factor_out_constants,
+                include_summaries=self.config.smart_crusher.include_summaries,
+            )
+            transforms.append(SmartCrusher(smart_config))
+        elif self.config.tool_crusher.enabled:
+            # Fallback to fixed-rule crushing
+            transforms.append(ToolCrusher(self.config.tool_crusher))
+
+        # 3. Rolling Window (enforce limits last)
+        if self.config.rolling_window.enabled:
+            transforms.append(RollingWindow(self.config.rolling_window))
+
+        return transforms
+
+    def _get_tokenizer(self, model: str) -> Tokenizer:
+        """Get tokenizer for model using provider."""
+        if self._provider is None:
+            raise ValueError(
+                "Provider is required for token counting. "
+                "Pass a provider to TransformPipeline or HeadroomClient."
+            )
+        token_counter = self._provider.get_token_counter(model)
+        return Tokenizer(token_counter, model)
+
+    def apply(
+        self,
+        messages: list[dict[str, Any]],
+        model: str,
+        **kwargs: Any,
+    ) -> TransformResult:
+        """
+        Apply all transforms in sequence.
+
+        Args:
+            messages: List of messages to transform.
+            model: Model name for token counting.
+            **kwargs: Additional arguments passed to transforms.
+                - model_limit: Context limit override.
+                - output_buffer: Output buffer override.
+                - tool_profiles: Per-tool compression profiles.
+                - request_id: Optional request ID for diff artifact.
+
+        Returns:
+            Combined TransformResult.
+        """
+        tokenizer = self._get_tokenizer(model)
+
+        # Get model limit from kwargs (should be set by client)
+        model_limit = kwargs.get("model_limit")
+        if model_limit is None:
+            raise ValueError(
+                "model_limit is required. Provide it via kwargs or "
+                "configure model_context_limits in HeadroomClient."
+            )
+
+        # Start with original tokens
+        tokens_before = tokenizer.count_messages(messages)
+
+        logger.debug(
+            "Pipeline starting: %d messages, %d tokens, model=%s",
+            len(messages),
+            tokens_before,
+            model,
+        )
+
+        # Track all transforms applied
+        all_transforms: list[str] = []
+        all_markers: list[str] = []
+        all_warnings: list[str] = []
+
+        # Track transform diffs if enabled
+        transform_diffs: list[TransformDiff] = []
+        generate_diff = self.config.generate_diff_artifact
+
+        current_messages = deep_copy_messages(messages)
+
+        for transform in self.transforms:
+            # Check if transform should run
+            if not transform.should_apply(current_messages, tokenizer, **kwargs):
+                continue
+
+            # Track tokens before this transform (for diff)
+            tokens_before_transform = tokenizer.count_messages(current_messages)
+
+            # Apply transform
+            result = transform.apply(current_messages, tokenizer, **kwargs)
+
+            # Update messages for next transform
+            current_messages = result.messages
+
+            # Track tokens after this transform (for diff)
+            tokens_after_transform = tokenizer.count_messages(current_messages)
+
+            # Accumulate results
+            all_transforms.extend(result.transforms_applied)
+            all_markers.extend(result.markers_inserted)
+            all_warnings.extend(result.warnings)
+
+            # Log transform results
+            if result.transforms_applied:
+                logger.info(
+                    "Transform %s: %d -> %d tokens (saved %d)",
+                    transform.name,
+                    tokens_before_transform,
+                    tokens_after_transform,
+                    tokens_before_transform - tokens_after_transform,
+                )
+            else:
+                logger.debug("Transform %s: no changes", transform.name)
+
+            # Record diff if enabled
+            if generate_diff:
+                transform_diffs.append(
+                    TransformDiff(
+                        transform_name=transform.name,
+                        tokens_before=tokens_before_transform,
+                        tokens_after=tokens_after_transform,
+                        tokens_saved=tokens_before_transform - tokens_after_transform,
+                        details=", ".join(result.transforms_applied)
+                        if result.transforms_applied
+                        else "",
+                    )
+                )
+
+        # Final token count
+        tokens_after = tokenizer.count_messages(current_messages)
+
+        # Log pipeline summary
+        total_saved = tokens_before - tokens_after
+        if total_saved > 0:
+            logger.info(
+                "Pipeline complete: %d -> %d tokens (saved %d, %.1f%% reduction)",
+                tokens_before,
+                tokens_after,
+                total_saved,
+                (total_saved / tokens_before * 100) if tokens_before > 0 else 0,
+            )
+        else:
+            logger.debug("Pipeline complete: no token savings")
+
+        # Build diff artifact if enabled
+        diff_artifact = None
+        if generate_diff:
+            diff_artifact = DiffArtifact(
+                request_id=kwargs.get("request_id", ""),
+                original_tokens=tokens_before,
+                optimized_tokens=tokens_after,
+                total_tokens_saved=tokens_before - tokens_after,
+                transforms=transform_diffs,
+            )
+
+        return TransformResult(
+            messages=current_messages,
+            tokens_before=tokens_before,
+            tokens_after=tokens_after,
+            transforms_applied=all_transforms,
+            markers_inserted=all_markers,
+            warnings=all_warnings,
+            diff_artifact=diff_artifact,
+        )
+
+    def simulate(
+        self,
+        messages: list[dict[str, Any]],
+        model: str,
+        **kwargs: Any,
+    ) -> TransformResult:
+        """
+        Simulate transforms without modifying messages.
+
+        Same as apply() but returns what WOULD happen.
+
+        Args:
+            messages: List of messages.
+            model: Model name.
+            **kwargs: Additional arguments.
+
+        Returns:
+            TransformResult with simulated changes.
+        """
+        # apply() already works on a copy, so this is safe
+        return self.apply(messages, model, **kwargs)
+
+
+def create_pipeline(
+    tool_crusher_config: ToolCrusherConfig | None = None,
+    cache_aligner_config: CacheAlignerConfig | None = None,
+    rolling_window_config: RollingWindowConfig | None = None,
+) -> TransformPipeline:
+    """
+    Create a pipeline with specific configurations.
+
+    Args:
+        tool_crusher_config: Tool crusher configuration.
+        cache_aligner_config: Cache aligner configuration.
+        rolling_window_config: Rolling window configuration.
+
+    Returns:
+        Configured TransformPipeline.
+    """
+    config = HeadroomConfig()
+
+    if tool_crusher_config is not None:
+        config.tool_crusher = tool_crusher_config
+    if cache_aligner_config is not None:
+        config.cache_aligner = cache_aligner_config
+    if rolling_window_config is not None:
+        config.rolling_window = rolling_window_config
+
+    return TransformPipeline(config)

headroom/transforms/rolling_window.py

@@ -0,0 +1,350 @@
+"""Rolling window transform for Headroom SDK."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from ..config import RollingWindowConfig, TransformResult
+from ..parser import find_tool_units
+from ..tokenizer import Tokenizer
+from ..tokenizers import EstimatingTokenCounter
+from ..utils import create_dropped_context_marker, deep_copy_messages
+from .base import Transform
+
+logger = logging.getLogger(__name__)
+
+
+class RollingWindow(Transform):
+    """
+    Apply rolling window to keep messages within token budget.
+
+    Drop order (deterministic):
+    1. Oldest TOOL UNITS (assistant+tool_calls paired with tool responses)
+    2. Oldest assistant+user pairs
+    3. Oldest RAG blocks (if detectable)
+
+    CRITICAL: Tool calls and tool results are atomic DROP UNITS.
+    Never orphan a tool result.
+
+    Never drops:
+    - System prompt
+    - Stable instructions
+    - Last N conversational turns (configurable)
+    """
+
+    name = "rolling_window"
+
+    def __init__(self, config: RollingWindowConfig | None = None):
+        """
+        Initialize rolling window.
+
+        Args:
+            config: Configuration for window behavior.
+        """
+        self.config = config or RollingWindowConfig()
+
+    def should_apply(
+        self,
+        messages: list[dict[str, Any]],
+        tokenizer: Tokenizer,
+        **kwargs: Any,
+    ) -> bool:
+        """Check if token cap is exceeded."""
+        if not self.config.enabled:
+            return False
+
+        model_limit = kwargs.get("model_limit", 128000)
+        output_buffer = kwargs.get("output_buffer", self.config.output_buffer_tokens)
+
+        current_tokens = tokenizer.count_messages(messages)
+        available = model_limit - output_buffer
+
+        return bool(current_tokens > available)
+
+    def apply(
+        self,
+        messages: list[dict[str, Any]],
+        tokenizer: Tokenizer,
+        **kwargs: Any,
+    ) -> TransformResult:
+        """
+        Apply rolling window to messages.
+
+        Args:
+            messages: List of messages.
+            tokenizer: Tokenizer for counting.
+            **kwargs: Must include 'model_limit', optionally 'output_buffer'.
+
+        Returns:
+            TransformResult with windowed messages.
+        """
+        model_limit = kwargs.get("model_limit", 128000)
+        output_buffer = kwargs.get("output_buffer", self.config.output_buffer_tokens)
+        available = model_limit - output_buffer
+
+        tokens_before = tokenizer.count_messages(messages)
+        result_messages = deep_copy_messages(messages)
+        transforms_applied: list[str] = []
+        markers_inserted: list[str] = []
+        warnings: list[str] = []
+
+        dropped_count = 0
+        tool_units_dropped = 0
+
+        # If already under budget, no changes needed
+        current_tokens = tokens_before
+        if current_tokens <= available:
+            return TransformResult(
+                messages=result_messages,
+                tokens_before=tokens_before,
+                tokens_after=tokens_before,
+                transforms_applied=[],
+                warnings=[],
+            )
+
+        # Identify protected indices
+        protected = self._get_protected_indices(result_messages)
+
+        # Identify tool units
+        tool_units = find_tool_units(result_messages)
+
+        # Create drop candidates with priorities
+        drop_candidates = self._build_drop_candidates(result_messages, protected, tool_units)
+
+        # Drop until under budget
+        indices_to_drop: set[int] = set()
+
+        for candidate in drop_candidates:
+            if current_tokens <= available:
+                break
+
+            # Get indices for this candidate
+            candidate_indices = candidate["indices"]
+
+            # Skip if any are protected
+            if any(idx in protected for idx in candidate_indices):
+                continue
+
+            # Skip if already dropped
+            if any(idx in indices_to_drop for idx in candidate_indices):
+                continue
+
+            # Calculate tokens saved
+            tokens_saved = sum(
+                tokenizer.count_message(result_messages[idx])
+                for idx in candidate_indices
+                if idx < len(result_messages)
+            )
+
+            indices_to_drop.update(candidate_indices)
+            current_tokens -= tokens_saved
+            dropped_count += 1
+
+            if candidate["type"] == "tool_unit":
+                tool_units_dropped += 1
+
+        # Remove dropped messages (in reverse order to preserve indices)
+        for idx in sorted(indices_to_drop, reverse=True):
+            if idx < len(result_messages):
+                del result_messages[idx]
+
+        # Insert marker if we dropped anything
+        if dropped_count > 0:
+            logger.info(
+                "RollingWindow: dropped %d units (%d tool units) to fit budget: %d -> %d tokens",
+                dropped_count,
+                tool_units_dropped,
+                tokens_before,
+                current_tokens,
+            )
+            marker = create_dropped_context_marker("token_cap", dropped_count)
+            markers_inserted.append(marker)
+
+            # Insert marker after system messages
+            insert_idx = 0
+            for i, msg in enumerate(result_messages):
+                if msg.get("role") != "system":
+                    insert_idx = i
+                    break
+            else:
+                insert_idx = len(result_messages)
+
+            result_messages.insert(
+                insert_idx,
+                {
+                    "role": "user",
+                    "content": marker,
+                },
+            )
+
+            transforms_applied.append(f"window_cap:{dropped_count}")
+
+        tokens_after = tokenizer.count_messages(result_messages)
+
+        result = TransformResult(
+            messages=result_messages,
+            tokens_before=tokens_before,
+            tokens_after=tokens_after,
+            transforms_applied=transforms_applied,
+            markers_inserted=markers_inserted,
+            warnings=warnings,
+        )
+
+        return result
+
+    def _get_protected_indices(self, messages: list[dict[str, Any]]) -> set[int]:
+        """Get indices that should never be dropped."""
+        protected: set[int] = set()
+
+        # Protect system messages
+        if self.config.keep_system:
+            for i, msg in enumerate(messages):
+                if msg.get("role") == "system":
+                    protected.add(i)
+
+        # Protect last N turns
+        if self.config.keep_last_turns > 0:
+            # Count turns from end (user+assistant = 1 turn)
+            turns_seen = 0
+            i = len(messages) - 1
+
+            while i >= 0 and turns_seen < self.config.keep_last_turns:
+                msg = messages[i]
+                role = msg.get("role")
+
+                # Protect this message
+                protected.add(i)
+
+                # Count turns
+                if role == "user":
+                    turns_seen += 1
+
+                i -= 1
+
+            # Also protect any tool responses that belong to protected assistant messages
+            for i in list(protected):
+                msg = messages[i]
+                if msg.get("role") == "assistant" and msg.get("tool_calls"):
+                    # Find and protect corresponding tool responses
+                    tool_call_ids = {tc.get("id") for tc in msg.get("tool_calls", [])}
+                    for j, other_msg in enumerate(messages):
+                        if other_msg.get("role") == "tool":
+                            if other_msg.get("tool_call_id") in tool_call_ids:
+                                protected.add(j)
+
+        return protected
+
+    def _build_drop_candidates(
+        self,
+        messages: list[dict[str, Any]],
+        protected: set[int],
+        tool_units: list[tuple[int, list[int]]],
+    ) -> list[dict[str, Any]]:
+        """
+        Build ordered list of drop candidates.
+
+        Returns candidates in drop priority order (first to drop first).
+        """
+        candidates: list[dict[str, Any]] = []
+
+        # Track which indices are part of tool units
+        tool_unit_indices: set[int] = set()
+        for assistant_idx, response_indices in tool_units:
+            tool_unit_indices.add(assistant_idx)
+            tool_unit_indices.update(response_indices)
+
+        # Priority 1: Oldest tool units (all indices as atomic unit)
+        for assistant_idx, response_indices in tool_units:
+            if assistant_idx in protected:
+                continue
+
+            all_indices = [assistant_idx] + response_indices
+            candidates.append(
+                {
+                    "type": "tool_unit",
+                    "indices": all_indices,
+                    "priority": 1,
+                    "position": assistant_idx,  # For sorting by age
+                }
+            )
+
+        # Priority 2: Oldest non-tool messages (user/assistant pairs)
+        i = 0
+        while i < len(messages):
+            msg = messages[i]
+            role = msg.get("role")
+
+            if i in protected or i in tool_unit_indices:
+                i += 1
+                continue
+
+            if role in ("user", "assistant"):
+                # Try to find a pair
+                if role == "user" and i + 1 < len(messages):
+                    next_msg = messages[i + 1]
+                    if next_msg.get("role") == "assistant" and i + 1 not in tool_unit_indices:
+                        candidates.append(
+                            {
+                                "type": "turn",
+                                "indices": [i, i + 1],
+                                "priority": 2,
+                                "position": i,
+                            }
+                        )
+                        i += 2
+                        continue
+
+                # Single message
+                candidates.append(
+                    {
+                        "type": "single",
+                        "indices": [i],
+                        "priority": 2,
+                        "position": i,
+                    }
+                )
+
+            i += 1
+
+        # Sort by priority, then by position (oldest first)
+        candidates.sort(key=lambda c: (c["priority"], c["position"]))
+
+        return candidates
+
+
+def apply_rolling_window(
+    messages: list[dict[str, Any]],
+    model_limit: int,
+    output_buffer: int = 4000,
+    keep_last_turns: int = 2,
+    config: RollingWindowConfig | None = None,
+) -> tuple[list[dict[str, Any]], list[str]]:
+    """
+    Convenience function to apply rolling window.
+
+    Args:
+        messages: List of messages.
+        model_limit: Model's context limit.
+        output_buffer: Tokens to reserve for output.
+        keep_last_turns: Number of recent turns to protect.
+        config: Optional configuration.
+
+    Returns:
+        Tuple of (windowed_messages, dropped_descriptions).
+    """
+    cfg = config or RollingWindowConfig()
+    cfg.output_buffer_tokens = output_buffer
+    cfg.keep_last_turns = keep_last_turns
+
+    window = RollingWindow(cfg)
+    tokenizer = Tokenizer(EstimatingTokenCounter())  # type: ignore[arg-type]
+
+    result = window.apply(
+        messages,
+        tokenizer,
+        model_limit=model_limit,
+        output_buffer=output_buffer,
+    )
+
+    return result.messages, result.transforms_applied