"""Transform pipeline orchestration for Headroom SDK.""" from __future__ import annotations import logging from typing import TYPE_CHECKING, Any from ..config import ( CacheAlignerConfig, DiffArtifact, HeadroomConfig, IntelligentContextConfig, RollingWindowConfig, ToolCrusherConfig, TransformDiff, TransformResult, ) from ..tokenizer import Tokenizer from ..utils import deep_copy_messages from .base import Transform from .cache_aligner import CacheAligner from .intelligent_context import IntelligentContextManager 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. Context Management (enforce limits last) # IntelligentContextManager takes precedence over RollingWindow when enabled if self.config.intelligent_context.enabled: # Use semantic-aware context management with scoring transforms.append(IntelligentContextManager(self.config.intelligent_context)) logger.info( "Pipeline using IntelligentContextManager with strategies: " "COMPRESS_FIRST -> SUMMARIZE -> DROP_BY_SCORE" ) elif self.config.rolling_window.enabled: # Fallback to position-based rolling window 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, intelligent_context_config: IntelligentContextConfig | 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. intelligent_context_config: Intelligent context configuration. When provided with enabled=True, replaces RollingWindow with semantic-aware context management. 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 if intelligent_context_config is not None: config.intelligent_context = intelligent_context_config return TransformPipeline(config)