Spaces:
Build error
Build error
| """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) | |