"""Code-aware compressor using AST parsing for syntax-preserving compression. This module provides AST-based compression for source code that guarantees valid syntax output. Unlike token-level compression (LLMLingua), this preserves structural elements while compressing function bodies. Key Features: - Syntax validity guaranteed (output always parses) - Preserves imports, signatures, type annotations, error handlers - Compresses function bodies while maintaining structure - Multi-language support via tree-sitter - Data-driven language config (no per-language method duplication) - Thread-safe (no mutable instance state during compression) Supported Languages (Tier 1): - Python, JavaScript, TypeScript Supported Languages (Tier 2): - Go, Rust, Java, C, C++ Compression Strategy: 1. Parse code into AST using tree-sitter 2. Extract and preserve critical structures (imports, signatures, types) 3. Rank functions by importance (using semantic analysis) 4. Compress function bodies while preserving signatures 5. Reassemble into valid code Installation: pip install headroom-ai[code] Usage: >>> from headroom.transforms import CodeAwareCompressor >>> compressor = CodeAwareCompressor() >>> result = compressor.compress(python_code) >>> print(result.compressed) # Valid Python code >>> print(result.syntax_valid) # True Reference: LongCodeZip: Compress Long Context for Code Language Models https://arxiv.org/abs/2510.00446 """ from __future__ import annotations import logging import re import threading from dataclasses import dataclass, field from enum import Enum from typing import Any from ..config import TransformResult from ..tokenizer import Tokenizer from .base import Transform logger = logging.getLogger(__name__) # Lazy import for optional dependency _tree_sitter_available: bool | None = None _tree_sitter_languages: dict[str, Any] = {} _tree_sitter_lock = threading.Lock() def _check_tree_sitter_available() -> bool: """Check if tree-sitter packages are available.""" global _tree_sitter_available if _tree_sitter_available is None: try: import tree_sitter_language_pack # noqa: F401 _tree_sitter_available = True except ImportError: _tree_sitter_available = False return _tree_sitter_available def _get_parser(language: str) -> Any: """Get a tree-sitter parser for the given language. Args: language: Language name (e.g., 'python', 'javascript'). Returns: Configured tree-sitter parser. Raises: ImportError: If tree-sitter is not installed. ValueError: If language is not supported. """ global _tree_sitter_languages if not _check_tree_sitter_available(): raise ImportError( "tree-sitter is not installed. Install with: pip install headroom-ai[code]\n" "This adds ~50MB for tree-sitter grammars." ) with _tree_sitter_lock: if language not in _tree_sitter_languages: try: from tree_sitter_language_pack import get_parser parser = get_parser(language) # type: ignore[arg-type] _tree_sitter_languages[language] = parser logger.debug("Loaded tree-sitter parser for %s", language) except Exception as e: raise ValueError( f"Language '{language}' is not supported by tree-sitter. " f"Supported: python, javascript, typescript, go, rust, java, c, cpp. " f"Error: {e}" ) from e return _tree_sitter_languages[language] def is_tree_sitter_available() -> bool: """Check if tree-sitter is installed and available. Returns: True if tree-sitter-languages package is installed. """ return _check_tree_sitter_available() def is_tree_sitter_loaded() -> bool: """Check if any tree-sitter parsers are currently loaded. Returns: True if parsers are loaded in memory. """ return len(_tree_sitter_languages) > 0 def unload_tree_sitter() -> bool: """Unload all tree-sitter parsers to free memory. Returns: True if parsers were unloaded, False if none were loaded. """ global _tree_sitter_languages with _tree_sitter_lock: if _tree_sitter_languages: count = len(_tree_sitter_languages) _tree_sitter_languages.clear() logger.info("Unloaded %d tree-sitter parsers", count) return True return False class CodeLanguage(Enum): """Supported programming languages.""" PYTHON = "python" JAVASCRIPT = "javascript" TYPESCRIPT = "typescript" GO = "go" RUST = "rust" JAVA = "java" C = "c" CPP = "cpp" UNKNOWN = "unknown" class DocstringMode(Enum): """How to handle docstrings.""" FULL = "full" # Keep entire docstring FIRST_LINE = "first_line" # Keep only first line REMOVE = "remove" # Remove docstrings completely NONE = "none" # Alias for REMOVE (deprecated) # ========================================================================= # Data-driven language configuration # ========================================================================= @dataclass(frozen=True) class LangConfig: """Data-driven configuration for a programming language. Instead of per-language methods, each language declares its AST node types and syntactic conventions. The compressor uses these tables to drive extraction and compression generically. """ # AST node types for structural extraction import_nodes: frozenset[str] function_nodes: frozenset[str] class_nodes: frozenset[str] type_nodes: frozenset[str] body_node_types: frozenset[str] # Node types that represent function/method bodies decorator_node: str | None # e.g. "decorated_definition" for Python # Syntax conventions comment_prefix: str # "#" for Python, "//" for C-family uses_colon_after_signature: bool # Python: True, C-family: False package_node: str | None = None # e.g. "package_clause" for Go # Quick pre-filter hints for language detection (substrings to check) detection_hints: tuple[str, ...] = () _LANG_CONFIGS: dict[CodeLanguage, LangConfig] = { CodeLanguage.PYTHON: LangConfig( import_nodes=frozenset({"import_statement", "import_from_statement"}), function_nodes=frozenset({"function_definition"}), class_nodes=frozenset({"class_definition"}), type_nodes=frozenset({"type_alias_statement"}), body_node_types=frozenset({"block"}), decorator_node="decorated_definition", comment_prefix="#", uses_colon_after_signature=True, detection_hints=("def ", "import ", "from ", "class ", "async def"), ), CodeLanguage.JAVASCRIPT: LangConfig( import_nodes=frozenset({"import_statement", "import_declaration"}), function_nodes=frozenset({"function_declaration", "method_definition"}), class_nodes=frozenset({"class_declaration"}), type_nodes=frozenset(), body_node_types=frozenset({"statement_block"}), decorator_node=None, comment_prefix="//", uses_colon_after_signature=False, detection_hints=("function ", "const ", "let ", "var ", "export ", "require("), ), CodeLanguage.TYPESCRIPT: LangConfig( import_nodes=frozenset({"import_statement", "import_declaration"}), function_nodes=frozenset({"function_declaration", "method_definition"}), class_nodes=frozenset({"class_declaration"}), type_nodes=frozenset({"interface_declaration", "type_alias_declaration"}), body_node_types=frozenset({"statement_block"}), decorator_node=None, comment_prefix="//", uses_colon_after_signature=False, detection_hints=("interface ", "type ", ": string", ": number", ": boolean"), ), CodeLanguage.GO: LangConfig( import_nodes=frozenset({"import_declaration"}), function_nodes=frozenset({"function_declaration", "method_declaration"}), class_nodes=frozenset(), type_nodes=frozenset({"type_declaration"}), body_node_types=frozenset({"block"}), decorator_node=None, comment_prefix="//", uses_colon_after_signature=False, package_node="package_clause", detection_hints=("func ", "package ", "struct {"), ), CodeLanguage.RUST: LangConfig( import_nodes=frozenset({"use_declaration"}), function_nodes=frozenset({"function_item"}), class_nodes=frozenset({"impl_item"}), type_nodes=frozenset({"struct_item", "enum_item", "type_item", "trait_item"}), body_node_types=frozenset({"block"}), decorator_node=None, comment_prefix="//", uses_colon_after_signature=False, detection_hints=("fn ", "struct ", "impl ", "mod ", "use "), ), CodeLanguage.JAVA: LangConfig( import_nodes=frozenset({"import_declaration"}), function_nodes=frozenset({"method_declaration", "constructor_declaration"}), class_nodes=frozenset({"class_declaration", "interface_declaration"}), type_nodes=frozenset({"enum_declaration"}), body_node_types=frozenset({"block"}), decorator_node=None, comment_prefix="//", uses_colon_after_signature=False, package_node="package_declaration", detection_hints=("public ", "private ", "protected ", "class ", "interface "), ), CodeLanguage.C: LangConfig( import_nodes=frozenset({"preproc_include"}), function_nodes=frozenset({"function_definition"}), class_nodes=frozenset(), type_nodes=frozenset({"struct_specifier", "enum_specifier", "type_definition"}), body_node_types=frozenset({"compound_statement"}), decorator_node=None, comment_prefix="//", uses_colon_after_signature=False, detection_hints=("#include", "typedef ", "int main("), ), CodeLanguage.CPP: LangConfig( import_nodes=frozenset({"preproc_include"}), function_nodes=frozenset({"function_definition"}), class_nodes=frozenset({"class_specifier"}), type_nodes=frozenset({"struct_specifier", "enum_specifier", "type_definition"}), body_node_types=frozenset({"compound_statement"}), decorator_node=None, comment_prefix="//", uses_colon_after_signature=False, detection_hints=("#include", "namespace ", "class ", "::"), ), } @dataclass class CodeStructure: """Extracted structure from parsed code.""" imports: list[str] = field(default_factory=list) type_definitions: list[str] = field(default_factory=list) class_definitions: list[str] = field(default_factory=list) function_signatures: list[str] = field(default_factory=list) function_bodies: list[tuple[str, str, int]] = field( default_factory=list ) # (signature, body, line) decorators: list[str] = field(default_factory=list) comments: list[str] = field(default_factory=list) top_level_code: list[str] = field(default_factory=list) other: list[str] = field(default_factory=list) @dataclass class CodeCompressorConfig: """Configuration for code-aware compression. Attributes: preserve_imports: Always keep import statements. preserve_signatures: Always keep function/method signatures. preserve_type_annotations: Keep type hints and annotations. preserve_decorators: Keep decorators on functions/classes. docstring_mode: How to handle docstrings. target_compression_rate: Target compression ratio (0.2 = keep 20%). max_body_lines: Maximum lines to keep per function body. compress_comments: Remove non-docstring comments. min_tokens_for_compression: Minimum tokens to trigger compression. language_hint: Explicit language (None = auto-detect). fallback_to_llmlingua: Use LLMLingua for unknown languages. enable_ccr: Store originals for retrieval. ccr_ttl: TTL for CCR entries in seconds. """ # Preservation settings preserve_imports: bool = True preserve_signatures: bool = True preserve_type_annotations: bool = True preserve_decorators: bool = True docstring_mode: DocstringMode = DocstringMode.FIRST_LINE # Compression settings target_compression_rate: float = 0.2 max_body_lines: int = 5 compress_comments: bool = True # Thresholds min_tokens_for_compression: int = 100 # Language handling language_hint: str | None = None fallback_to_llmlingua: bool = True # Semantic analysis (symbol importance scoring) semantic_analysis: bool = True # CCR integration enable_ccr: bool = True ccr_ttl: int = 300 # 5 minutes @dataclass class CodeCompressionResult: """Result of code-aware compression. Attributes: compressed: The compressed code (guaranteed valid syntax). original: Original code before compression. original_tokens: Token count before compression. compressed_tokens: Token count after compression. compression_ratio: Actual compression ratio achieved. language: Detected or specified language. language_confidence: Confidence in language detection. preserved_imports: Number of import statements preserved. preserved_signatures: Number of function signatures preserved. compressed_bodies: Number of function bodies compressed. syntax_valid: Whether output is syntactically valid. cache_key: CCR cache key if stored. """ compressed: str original: str original_tokens: int compressed_tokens: int compression_ratio: float # Code-specific metadata language: CodeLanguage = CodeLanguage.UNKNOWN language_confidence: float = 0.0 # Structure analysis preserved_imports: int = 0 preserved_signatures: int = 0 compressed_bodies: int = 0 # Validation syntax_valid: bool = True # CCR cache_key: str | None = None # Semantic analysis symbol_scores: dict[str, float] = field(default_factory=dict) @property def tokens_saved(self) -> int: """Number of tokens saved by compression.""" return max(0, self.original_tokens - self.compressed_tokens) @property def savings_percentage(self) -> float: """Percentage of tokens saved.""" if self.original_tokens == 0: return 0.0 return (self.tokens_saved / self.original_tokens) * 100 @property def summary(self) -> str: """Human-readable summary of compression.""" analysis_note = "" if self.symbol_scores: high = sum(1 for s in self.symbol_scores.values() if s >= 0.7) low = sum(1 for s in self.symbol_scores.values() if s < 0.1) if high or low: analysis_note = f" Semantic: {high} high-importance, {low} low-importance." return ( f"Compressed {self.language.value} code: " f"{self.original_tokens:,}→{self.compressed_tokens:,} tokens " f"({self.savings_percentage:.0f}% saved). " f"Kept {self.preserved_imports} imports, " f"{self.preserved_signatures} signatures, " f"compressed {self.compressed_bodies} bodies." f"{analysis_note}" ) # ========================================================================= # Language detection # ========================================================================= # Lightweight pre-filter patterns for language detection. # These are ONLY used as a quick check to avoid parsing with every language. # Actual detection is done by tree-sitter (fewest parse errors wins). _LANGUAGE_PREFILTER: dict[CodeLanguage, list[re.Pattern[str]]] = { CodeLanguage.PYTHON: [ re.compile(r"^\s*(def|class|import|from|async def)\s+\w+", re.MULTILINE), re.compile(r"^\s*@\w+", re.MULTILINE), re.compile(r'^\s*"""', re.MULTILINE), re.compile(r"^\s*if __name__\s*==", re.MULTILINE), ], CodeLanguage.JAVASCRIPT: [ re.compile(r"^\s*(function|const|let|var|class|export)\s+\w+", re.MULTILINE), re.compile(r"^\s*async\s+(function|=>)", re.MULTILINE), re.compile(r"^\s*module\.exports", re.MULTILINE), re.compile(r"^\s*(import|export)\s+.*\s+from\s+['\"]", re.MULTILINE), ], CodeLanguage.TYPESCRIPT: [ re.compile(r"^\s*(interface|type|enum|namespace)\s+\w+", re.MULTILINE), re.compile(r":\s*(string|number|boolean|any|void|Promise)\b", re.MULTILINE), ], CodeLanguage.GO: [ re.compile(r"^\s*(func|type|package|import)\s+", re.MULTILINE), re.compile(r"^\s*func\s+\([^)]+\)\s+\w+", re.MULTILINE), re.compile(r"\bstruct\s*\{", re.MULTILINE), ], CodeLanguage.RUST: [ re.compile(r"^\s*(fn|struct|enum|impl|mod|use|pub)\s+", re.MULTILINE), re.compile(r"^\s*#\[", re.MULTILINE), ], CodeLanguage.JAVA: [ re.compile(r"^\s*(public|private|protected)\s+(class|interface|enum)", re.MULTILINE), re.compile(r"^\s*package\s+[\w.]+;", re.MULTILINE), ], CodeLanguage.C: [ re.compile(r"^\s*#include\s*[<\"]", re.MULTILINE), re.compile(r"^\s*(int|void|char|float|double)\s+\w+\s*\(", re.MULTILINE), re.compile(r"^\s*typedef\s+", re.MULTILINE), ], CodeLanguage.CPP: [ re.compile(r"^\s*#include\s*[<\"]", re.MULTILINE), re.compile(r"\bnamespace\s+\w+", re.MULTILINE), re.compile(r"::\w+", re.MULTILINE), ], } def _count_error_nodes(node: Any) -> int: """Count ERROR and MISSING nodes in a tree-sitter AST.""" count = 0 if node.type == "ERROR" or node.is_missing: count += 1 for child in node.children: count += _count_error_nodes(child) return count def detect_language(code: str) -> tuple[CodeLanguage, float]: """Detect the programming language of code. Uses tree-sitter AST parsing when available (most accurate), with a regex pre-filter to avoid parsing with all languages. Falls back to regex-only scoring when tree-sitter is unavailable. Args: code: Source code to analyze. Returns: Tuple of (detected language, confidence score 0.0-1.0). """ if not code or not code.strip(): return CodeLanguage.UNKNOWN, 0.0 sample = code[:5000] # Phase 1: Pre-filter — find candidate languages using quick regex candidates: dict[CodeLanguage, int] = {} for lang, patterns in _LANGUAGE_PREFILTER.items(): score = 0 for pattern in patterns: matches = len(pattern.findall(sample)) score += matches if score > 0: candidates[lang] = score if not candidates: return CodeLanguage.UNKNOWN, 0.0 # Disambiguation: TypeScript superset of JavaScript if CodeLanguage.TYPESCRIPT in candidates and CodeLanguage.JAVASCRIPT in candidates: if candidates[CodeLanguage.TYPESCRIPT] >= 2: candidates[CodeLanguage.JAVASCRIPT] = 0 # Disambiguation: C++ superset of C if CodeLanguage.CPP in candidates and CodeLanguage.C in candidates: if candidates[CodeLanguage.CPP] >= 2: candidates[CodeLanguage.C] = 0 # Phase 2: If tree-sitter available, parse with candidates and pick fewest errors if _check_tree_sitter_available(): best_lang = CodeLanguage.UNKNOWN min_errors = float("inf") best_node_count = 0 code_bytes = bytes(code[:10000], "utf-8") # Sort candidates by pre-filter score (try most likely first) sorted_candidates = sorted(candidates.items(), key=lambda x: x[1], reverse=True) for lang, _prefilter_score in sorted_candidates: if lang == CodeLanguage.UNKNOWN or candidates.get(lang, 0) == 0: continue try: parser = _get_parser(lang.value) tree = parser.parse(code_bytes) error_count = _count_error_nodes(tree.root_node) node_count = tree.root_node.child_count # Prefer: fewest errors, then most top-level nodes (richer parse) if error_count < min_errors or ( error_count == min_errors and node_count > best_node_count ): min_errors = error_count best_lang = lang best_node_count = node_count except (ValueError, ImportError): continue if best_lang != CodeLanguage.UNKNOWN: # Confidence based on error ratio total_lines = max(1, len(code.strip().split("\n"))) error_ratio = min_errors / total_lines confidence = max(0.3, min(1.0, 1.0 - error_ratio)) return best_lang, confidence # Phase 3: Fallback — regex-only scoring (no tree-sitter) best_lang = max(candidates, key=lambda k: candidates[k]) best_score = candidates[best_lang] if best_score == 0: return CodeLanguage.UNKNOWN, 0.0 confidence = min(1.0, 0.3 + (best_score * 0.1)) return best_lang, confidence # ========================================================================= # Symbol importance analysis # ========================================================================= @dataclass class _SymbolAnalysis: """Result of intra-file symbol importance analysis. All dicts are keyed by qualified name (e.g., 'ClassName.method') to avoid collisions between identically-named methods in different classes. """ scores: dict[str, float] = field(default_factory=dict) calls: dict[str, set[str]] = field(default_factory=dict) ref_counts: dict[str, int] = field(default_factory=dict) body_line_counts: dict[str, int] = field(default_factory=dict) bare_names: dict[str, str] = field(default_factory=dict) # qname -> short_name class CodeAwareCompressor(Transform): """AST-preserving compression for source code. This compressor uses tree-sitter to parse code into an AST, then selectively compresses function bodies while preserving structure. The output is guaranteed to be syntactically valid. Key advantages over token-level compression: - Syntax validity guaranteed - Preserves imports, signatures, types - Better compression ratios for code (5-8x vs 3-5x) - Lower latency (~20-50ms vs 50-200ms for LLMLingua) - Smaller memory footprint (~50MB vs ~1GB) - Thread-safe (no mutable instance state during compression) Example: >>> compressor = CodeAwareCompressor() >>> result = compressor.compress(''' ... import os ... from typing import List ... ... def process_data(items: List[str]) -> List[str]: ... \"\"\"Process a list of items.\"\"\" ... results = [] ... for item in items: ... # Validate item ... if not item: ... continue ... # Process valid item ... processed = item.strip().lower() ... results.append(processed) ... return results ... ''') >>> print(result.compressed) import os from typing import List def process_data(items: List[str]) -> List[str]: \"\"\"Process a list of items.\"\"\" # ... (body compressed: 10 lines → 2 lines) pass """ name: str = "code_aware_compressor" def __init__(self, config: CodeCompressorConfig | None = None): """Initialize code-aware compressor. Args: config: Compression configuration. If None, uses defaults. Note: Tree-sitter parsers are loaded lazily on first use to avoid startup overhead when the compressor isn't used. """ self.config = config or CodeCompressorConfig() # ========================================================================= # Token estimation # ========================================================================= @staticmethod def _estimate_tokens(text: str, tokenizer: Tokenizer | None = None) -> int: """Count or estimate tokens for text. Uses real tokenizer when available; falls back to chars/4 which is a much closer approximation for code than word count. """ if tokenizer is not None: return tokenizer.count_text(text) # chars/4 is a reasonable approximation for code tokens # (code has lots of punctuation that tokenizes separately) return max(1, len(text) // 4) # ========================================================================= # Symbol importance analysis # ========================================================================= def _analyze_symbol_importance( self, root: Any, code: str, language: CodeLanguage, context: str = "", ) -> _SymbolAnalysis: """Analyze symbol importance using distribution-based scoring. Collects raw signals (reference count, fan-out, visibility, context match, convention importance) per symbol, then normalizes using min-max scaling so scores are relative within the file. This adapts to any file structure: utility libraries, test files, orchestrators, etc. Returns _SymbolAnalysis with normalized scores (0.0-1.0) per symbol. """ if not self.config.semantic_analysis: return _SymbolAnalysis() lang_config = _LANG_CONFIGS.get(language) if not lang_config: return _SymbolAnalysis() all_definition_types = lang_config.function_nodes | lang_config.class_nodes # Use qualified keys (ClassName.method) to avoid collisions definitions: dict[str, Any] = {} # qualified_name -> node bare_names: dict[str, str] = {} # qualified_name -> short_name all_identifiers: dict[str, int] = {} # short_name -> count function_calls: dict[str, set[str]] = {} def collect_definitions(node: Any, parent_name: str = "") -> None: if node.type in all_definition_types: short_name = _get_definition_name(node) if short_name: qualified = f"{parent_name}.{short_name}" if parent_name else short_name definitions[qualified] = node bare_names[qualified] = short_name for child in node.children: collect_definitions(child, parent_name=qualified) return # Also check for decorated definitions if lang_config.decorator_node and node.type == lang_config.decorator_node: for child in node.children: if child.type in all_definition_types: short_name = _get_definition_name(child) if short_name: qualified = f"{parent_name}.{short_name}" if parent_name else short_name definitions[qualified] = child bare_names[qualified] = short_name for grandchild in child.children: collect_definitions(grandchild, parent_name=qualified) return for child in node.children: collect_definitions(child, parent_name) def collect_identifiers(node: Any) -> None: if node.type in ("identifier", "property_identifier", "type_identifier"): text = node.text name = text.decode("utf-8") if isinstance(text, bytes) else str(text) all_identifiers[name] = all_identifiers.get(name, 0) + 1 for child in node.children: collect_identifiers(child) def collect_calls_in_function(func_node: Any, func_qname: str) -> None: func_short = bare_names[func_qname] defined_short_names = set(bare_names.values()) calls: set[str] = set() def walk(node: Any) -> None: if node.type in ("identifier", "property_identifier"): text = node.text name = text.decode("utf-8") if isinstance(text, bytes) else str(text) if name in defined_short_names and name != func_short: calls.add(name) for child in node.children: walk(child) walk(func_node) function_calls[func_qname] = calls # Pass 1: Collect definitions with qualified names collect_definitions(root) if not definitions: return _SymbolAnalysis() # Pass 2: Collect all identifiers collect_identifiers(root) # Pass 3: Collect call relationships and body sizes body_line_counts: dict[str, int] = {} for qname, node in definitions.items(): collect_calls_in_function(node, qname) node_text = code[node.start_byte : node.end_byte] body_line_counts[qname] = max(1, len(node_text.split("\n")) - 2) # Reference counts: subtract definition occurrences short_name_def_count: dict[str, int] = {} for short in bare_names.values(): short_name_def_count[short] = short_name_def_count.get(short, 0) + 1 ref_counts: dict[str, int] = {} for qname in definitions: short = bare_names[qname] count = all_identifiers.get(short, 0) ref_counts[qname] = max(0, count - short_name_def_count.get(short, 1)) # Raw importance signals per symbol context_lower = context.lower() if context else "" context_words = set(re.split(r"[\s,;:.()\[\]{}\"']+", context_lower)) if context else set() context_words.discard("") raw_signals: dict[str, float] = {} for qname in definitions: short = bare_names[qname] refs = ref_counts.get(qname, 0) fan_out = len(function_calls.get(qname, set())) is_public = _is_public_symbol(short, language) raw = float(refs) raw += 1.0 if is_public else 0.0 raw += fan_out * 0.5 # Convention importance (language-specific) if language == CodeLanguage.PYTHON: if short.startswith("__") and short.endswith("__"): raw += 2.0 elif language == CodeLanguage.GO: if short and short[0].isupper(): raw += 1.0 # Context boost if context_words: name_lower = short.lower() if name_lower in context_words or ( len(name_lower) > 3 and name_lower in context_lower ): raw += 3.0 raw_signals[qname] = raw # Normalize to 0-1 using min-max scaling values = list(raw_signals.values()) min_val = min(values) max_val = max(values) range_val = max_val - min_val if range_val > 0: scores = {name: round((v - min_val) / range_val, 3) for name, v in raw_signals.items()} else: scores = dict.fromkeys(raw_signals, 0.5) return _SymbolAnalysis( scores=scores, calls=function_calls, ref_counts=ref_counts, body_line_counts=body_line_counts, bare_names=bare_names, ) def _allocate_body_budget(self, analysis: _SymbolAnalysis, code: str) -> dict[str, int]: """Allocate body line budget across functions using target_compression_rate. Returns dict mapping symbol name to max body lines to keep. """ if not analysis.scores or not analysis.body_line_counts: return {} scores = analysis.scores body_sizes = analysis.body_line_counts target_rate = self.config.target_compression_rate total_lines = len(code.strip().split("\n")) total_body_lines = sum(body_sizes.values()) fixed_lines = max(0, total_lines - total_body_lines) target_total = total_lines * target_rate body_budget = max(0.0, target_total - fixed_lines) if total_body_lines == 0: return {} score_floor = 0.05 weights: dict[str, float] = {} for name in scores: score = max(scores.get(name, 0.5), score_floor) size = body_sizes.get(name, 0) weights[name] = score * size total_weight = sum(weights.values()) if total_weight == 0: per_func = max(0, int(body_budget / max(len(scores), 1))) return {name: min(per_func, body_sizes.get(name, 0)) for name in scores} limits: dict[str, int] = {} for qname in scores: allocation = body_budget * weights[qname] / total_weight max_lines = body_sizes.get(qname, 0) limit = min(int(round(allocation)), max_lines) limits[qname] = limit # Also store by short name so _get_body_limit can find it. short = analysis.bare_names.get(qname, qname) if short not in limits or limit > limits[short]: limits[short] = limit return limits # ========================================================================= # Core compression # ========================================================================= def compress( self, code: str, language: str | None = None, context: str = "", tokenizer: Tokenizer | None = None, ) -> CodeCompressionResult: """Compress code while preserving syntax validity. Args: code: Source code to compress. language: Language name (e.g., 'python'). Auto-detected if None. context: Optional context for relevance-aware compression. tokenizer: Optional tokenizer for accurate token counting. Returns: CodeCompressionResult with compressed code and metadata. """ if not code or not code.strip(): return CodeCompressionResult( compressed=code, original=code, original_tokens=0, compressed_tokens=0, compression_ratio=1.0, syntax_valid=True, ) original_tokens = self._estimate_tokens(code, tokenizer) # Skip small content if original_tokens < self.config.min_tokens_for_compression: return CodeCompressionResult( compressed=code, original=code, original_tokens=original_tokens, compressed_tokens=original_tokens, compression_ratio=1.0, syntax_valid=True, ) # Detect or use specified language if language: detected_lang = CodeLanguage(language.lower()) confidence = 1.0 elif self.config.language_hint: detected_lang = CodeLanguage(self.config.language_hint.lower()) confidence = 1.0 else: detected_lang, confidence = detect_language(code) # If language unknown and fallback enabled, try LLMLingua if detected_lang == CodeLanguage.UNKNOWN: if self.config.fallback_to_llmlingua: return self._fallback_compress(code, original_tokens) else: return CodeCompressionResult( compressed=code, original=code, original_tokens=original_tokens, compressed_tokens=original_tokens, compression_ratio=1.0, language=CodeLanguage.UNKNOWN, language_confidence=0.0, syntax_valid=True, ) # Check if tree-sitter is available if not _check_tree_sitter_available(): logger.warning("tree-sitter not available. Install with: pip install headroom-ai[code]") if self.config.fallback_to_llmlingua: return self._fallback_compress(code, original_tokens) return CodeCompressionResult( compressed=code, original=code, original_tokens=original_tokens, compressed_tokens=original_tokens, compression_ratio=1.0, language=detected_lang, language_confidence=confidence, syntax_valid=True, ) # Parse and compress try: compressed, structure, symbol_scores = self._compress_with_ast( code, detected_lang, context, tokenizer ) compressed_tokens = self._estimate_tokens(compressed, tokenizer) # Verify syntax validity (checks both ERROR and MISSING nodes) syntax_valid = self._verify_syntax(compressed, detected_lang) # If syntax invalid, fall back to original if not syntax_valid: logger.warning("Compression produced invalid syntax, returning original") return CodeCompressionResult( compressed=code, original=code, original_tokens=original_tokens, compressed_tokens=original_tokens, compression_ratio=1.0, language=detected_lang, language_confidence=confidence, syntax_valid=True, ) ratio = compressed_tokens / max(original_tokens, 1) # Guard against over-aggressive compression (data loss) if ratio < 0.05: logger.warning( "Code compression too aggressive (ratio=%.3f), returning original", ratio, ) return CodeCompressionResult( compressed=code, original=code, original_tokens=original_tokens, compressed_tokens=original_tokens, compression_ratio=1.0, language=detected_lang, language_confidence=confidence, syntax_valid=True, ) # Store in CCR if significant compression cache_key = None if self.config.enable_ccr and ratio < 0.8: cache_key = self._store_in_ccr(code, compressed, original_tokens) if cache_key: from .compression_summary import summarize_compressed_code code_summary = summarize_compressed_code( structure.function_bodies, len(structure.function_bodies), ) summary_str = f" {code_summary}." if code_summary else "" # Use the actual config attribute (not the wrong name) ttl_min = max(1, self.config.ccr_ttl // 60) compressed += ( f"\n# [{original_tokens - compressed_tokens} tokens compressed." f"{summary_str}" f" Retrieve more: hash={cache_key}." f" Expires in {ttl_min}m.]" ) return CodeCompressionResult( compressed=compressed, original=code, original_tokens=original_tokens, compressed_tokens=compressed_tokens, compression_ratio=ratio, language=detected_lang, language_confidence=confidence, preserved_imports=len(structure.imports), preserved_signatures=len(structure.function_signatures), compressed_bodies=len(structure.function_bodies), syntax_valid=syntax_valid, cache_key=cache_key, symbol_scores=symbol_scores, ) except Exception as e: logger.warning("AST compression failed: %s, falling back", e) if self.config.fallback_to_llmlingua: return self._fallback_compress(code, original_tokens) return CodeCompressionResult( compressed=code, original=code, original_tokens=original_tokens, compressed_tokens=original_tokens, compression_ratio=1.0, language=detected_lang, language_confidence=confidence, syntax_valid=True, ) def _compress_with_ast( self, code: str, language: CodeLanguage, context: str, tokenizer: Tokenizer | None = None, ) -> tuple[str, CodeStructure, dict[str, float]]: """Compress code using AST parsing with symbol importance analysis. Thread-safe: all mutable state is passed through parameters, not stored on self. Args: code: Source code. language: Detected language. context: User context for relevance. tokenizer: Optional tokenizer for accurate token counting. Returns: Tuple of (compressed code, extracted structure, symbol scores). """ parser = _get_parser(language.value) tree = parser.parse(bytes(code, "utf-8")) root = tree.root_node # Analyze symbol importance and allocate compression budget analysis = self._analyze_symbol_importance(root, code, language, context) body_limits = self._allocate_body_budget(analysis, code) # Extract structure using data-driven language config lang_config = _LANG_CONFIGS.get(language) if lang_config: structure = self._extract_structure( root, code, language, lang_config, body_limits, analysis ) else: structure = self._extract_generic_structure(root, code) # Assemble compressed code compressed = self._assemble_compressed(structure, language) # Expose scores with short names for the public API symbol_scores: dict[str, float] = {} if analysis.scores: for qname, score in analysis.scores.items(): short = analysis.bare_names.get(qname, qname) if short not in symbol_scores or score > symbol_scores[short]: symbol_scores[short] = score return compressed, structure, symbol_scores # ========================================================================= # Unified structure extraction (data-driven, replaces per-language methods) # ========================================================================= def _extract_structure( self, root: Any, code: str, language: CodeLanguage, lang_config: LangConfig, body_limits: dict[str, int], analysis: _SymbolAnalysis, ) -> CodeStructure: """Extract structure from AST using data-driven language config. A single visitor handles all languages by checking node types against the LangConfig tables. No per-language extraction methods needed. """ structure = CodeStructure() captured_byte_ranges: list[tuple[int, int]] = [] def visit(node: Any) -> None: node_type = node.type # Package declarations (Go, Java) if lang_config.package_node and node_type == lang_config.package_node: structure.imports.insert(0, _get_node_text(node, code)) captured_byte_ranges.append((node.start_byte, node.end_byte)) return # Import statements if node_type in lang_config.import_nodes: structure.imports.append(_get_node_text(node, code)) captured_byte_ranges.append((node.start_byte, node.end_byte)) return # Export statements (JS/TS) — may contain functions or re-exports if node_type == "export_statement": text = _get_node_text(node, code) # Check if this export wraps a function or class has_func_or_class = False for child in node.children: if ( child.type in lang_config.function_nodes or child.type in lang_config.class_nodes ): has_func_or_class = True compressed = self._compress_function_ast( child, code, language, lang_config, body_limits, analysis ) # Reconstruct export with compressed inner definition export_prefix = code[node.start_byte : child.start_byte] export_suffix = code[child.end_byte : node.end_byte] structure.function_signatures.append( export_prefix + compressed + export_suffix ) break if not has_func_or_class: structure.imports.append(text) captured_byte_ranges.append((node.start_byte, node.end_byte)) return # Decorated definitions (Python) if lang_config.decorator_node and node_type == lang_config.decorator_node: decorator_text = [] definition_compressed = None for child in node.children: if child.type == "decorator": decorator_text.append(_get_node_text(child, code)) elif child.type in lang_config.function_nodes: definition_compressed = self._compress_function_ast( child, code, language, lang_config, body_limits, analysis ) elif child.type in lang_config.class_nodes: definition_compressed = self._compress_class_ast( child, code, language, lang_config, body_limits, analysis ) if decorator_text and definition_compressed: full_def = "\n".join(decorator_text) + "\n" + definition_compressed # Route to correct list based on inner definition type for child in node.children: if child.type in lang_config.class_nodes: structure.class_definitions.append(full_def) break else: structure.function_signatures.append(full_def) elif definition_compressed: structure.function_signatures.append(definition_compressed) captured_byte_ranges.append((node.start_byte, node.end_byte)) return # Function/method definitions if node_type in lang_config.function_nodes: compressed = self._compress_function_ast( node, code, language, lang_config, body_limits, analysis ) structure.function_signatures.append(compressed) captured_byte_ranges.append((node.start_byte, node.end_byte)) return # Class definitions — compress each method individually if node_type in lang_config.class_nodes: compressed = self._compress_class_ast( node, code, language, lang_config, body_limits, analysis ) structure.class_definitions.append(compressed) captured_byte_ranges.append((node.start_byte, node.end_byte)) return # Type definitions if node_type in lang_config.type_nodes: structure.type_definitions.append(_get_node_text(node, code)) captured_byte_ranges.append((node.start_byte, node.end_byte)) return # Recurse into children for child in node.children: visit(child) visit(root) # Capture top-level code that wasn't handled by any of the above. # This preserves global variables, constants, if __name__ blocks, # module-level assignments, etc. for child in root.children: child_range = (child.start_byte, child.end_byte) if child_range not in captured_byte_ranges: text = _get_node_text(child, code).strip() if text: structure.top_level_code.append(text) return structure # ========================================================================= # Unified function/class compression (data-driven) # ========================================================================= def _compress_function_ast( self, node: Any, code: str, language: CodeLanguage, lang_config: LangConfig, body_limits: dict[str, int], analysis: _SymbolAnalysis, ) -> str: """Compress a function/class/impl block using AST body detection. Uses the AST to find the body node directly instead of string-scanning for '{' or ':'. Works for all languages via lang_config.body_node_types. Key insight: tree-sitter byte offsets may not include leading whitespace on the first line. We use LINE-based slicing from the original code to preserve indentation faithfully. """ # Use line-based slicing from original code (not byte offsets) to # preserve indentation. This is critical for nested definitions # (methods inside classes). code_lines = code.split("\n") start_row = node.start_point[0] end_row = node.end_point[0] node_lines = code_lines[start_row : end_row + 1] node_text = "\n".join(node_lines) func_name = _get_definition_name(node) body_limit = _get_body_limit(func_name, body_limits, self.config.max_body_lines) # Small enough to keep as-is if len(node_lines) <= body_limit + 2: return node_text # Find the body node using AST (not string scanning) body_node = None for child in node.children: if child.type in lang_config.body_node_types: body_node = child break if body_node is None: return node_text # Use line numbers to slice: this preserves original indentation. # tree-sitter gives 0-based row numbers. node_start_line = node.start_point[0] body_start_line = body_node.start_point[0] body_end_line = body_node.end_point[0] # Lines within the node (0-indexed relative to node start) sig_end = body_start_line - node_start_line # exclusive body_end_rel = body_end_line - node_start_line + 1 # inclusive # Handle case where signature and body start on the SAME line # (common in brace languages: `function foo(arg) { ... }`) if sig_end == 0 and not lang_config.uses_colon_after_signature: # Signature and body on same line: `function foo(arg) { ... }` # Keep them together: sig includes up to and including `{` first_line = node_lines[0] # Include the opening brace in the signature line sig_with_brace = first_line.rstrip() signature_lines = [sig_with_brace] # Body lines are everything between { and } (inner content only) body_lines = node_lines[1:body_end_rel] after_lines = node_lines[body_end_rel:] # We've already included { in signature, so mark it _brace_in_signature = True else: signature_lines = node_lines[:sig_end] body_lines = node_lines[sig_end:body_end_rel] after_lines = node_lines[body_end_rel:] _brace_in_signature = False # For brace languages, detect opening/closing braces in the body lines. opening_brace_line = None closing_brace_line = None if not lang_config.uses_colon_after_signature: if _brace_in_signature: # Opening brace already in signature line — just find closing pass elif body_lines and body_lines[0].strip().startswith("{"): opening_brace_line = body_lines[0] body_lines = body_lines[1:] if body_lines and body_lines[-1].strip().endswith("}"): closing_brace_line = body_lines[-1] body_lines = body_lines[:-1] # Handle Python docstrings via AST docstring_text = "" ds_skip_lines = 0 if language == CodeLanguage.PYTHON and body_node.child_count > 0: first_child = body_node.children[0] # tree-sitter Python may represent docstrings as: # - bare `string` node directly in block, OR # - `expression_statement` containing a `string` node ds_node = None if first_child.type == "string": ds_node = first_child elif first_child.type == "expression_statement" and first_child.child_count > 0: if first_child.children[0].type == "string": ds_node = first_child if ds_node is not None: ds_text = _get_node_text(ds_node, code) ds_lines_count = ds_node.end_point[0] - ds_node.start_point[0] + 1 ds_start_rel = ds_node.start_point[0] - body_node.start_point[0] if self.config.docstring_mode == DocstringMode.FULL: # Keep entire docstring as-is (preserve indentation from body_lines) docstring_text = "\n".join( body_lines[ds_start_rel : ds_start_rel + ds_lines_count] ) elif self.config.docstring_mode == DocstringMode.FIRST_LINE: full_ds = ds_text.strip() ds_first_line_indent = ( body_lines[ds_start_rel][ : len(body_lines[ds_start_rel]) - len(body_lines[ds_start_rel].lstrip()) ] if body_lines[ds_start_rel:] else "" ) ds_stripped_lines = full_ds.split("\n") if len(ds_stripped_lines) == 1: docstring_text = body_lines[ds_start_rel] else: first_stripped = ds_stripped_lines[0].strip() if first_stripped in ('"""', "'''", 'r"""', "r'''"): # Opening quote on its own line if len(ds_stripped_lines) > 1: docstring_text = ( ds_first_line_indent + '"""' + ds_stripped_lines[1].strip() + '"""' ) elif first_stripped.endswith('"""'): # Single-line docstring: """text""" docstring_text = body_lines[ds_start_rel] else: # Multi-line: """text\n...\n""" → """text...""" docstring_text = ds_first_line_indent + first_stripped + '..."""' # elif REMOVE: docstring_text stays empty ds_skip_lines = ds_start_rel + ds_lines_count # Get content lines after docstring content_lines = body_lines[ds_skip_lines:] total_body = len(content_lines) keep_lines = min(body_limit, total_body) # Detect indentation from actual body code (preserves whatever the file uses) indent = _detect_indent(body_lines) if body_lines else " " # Build compressed output preserving original indentation result_parts: list[str] = [] # Signature lines (may be multi-line) if signature_lines: result_parts.extend(signature_lines) else: # Single-line: signature and body on same line — extract up to colon/brace sig_text = code[node.start_byte : body_node.start_byte].rstrip() result_parts.append(sig_text) if opening_brace_line is not None: result_parts.append(opening_brace_line) if docstring_text and self.config.docstring_mode not in ( DocstringMode.NONE, DocstringMode.REMOVE, ): result_parts.append(docstring_text) if keep_lines > 0: result_parts.extend(content_lines[:keep_lines]) if total_body > keep_lines: omitted = total_body - keep_lines result_parts.append( _make_omitted_comment( func_name, omitted, indent, lang_config.comment_prefix, analysis ) ) if lang_config.uses_colon_after_signature: result_parts.append(f"{indent}pass") if closing_brace_line is not None: result_parts.append(closing_brace_line) elif after_lines: result_parts.extend(after_lines) return "\n".join(result_parts) def _compress_class_ast( self, node: Any, code: str, language: CodeLanguage, lang_config: LangConfig, body_limits: dict[str, int], analysis: _SymbolAnalysis, ) -> str: """Compress a class by individually compressing each method. Preserves class-level attributes, type annotations, and decorators while compressing method bodies individually. This ensures correct indentation for each method's omitted-body comment. """ # Use line-based extraction to preserve indentation code_lines = code.split("\n") start_row = node.start_point[0] end_row = node.end_point[0] node_lines = code_lines[start_row : end_row + 1] node_text = "\n".join(node_lines) # Find the body node body_node = None for child in node.children: if child.type in lang_config.body_node_types: body_node = child break if body_node is None: return node_text # Class header (signature) — everything before the body node_start_line = node.start_point[0] body_start_line = body_node.start_point[0] sig_end = body_start_line - node_start_line header_lines = node_lines[:sig_end] if sig_end > 0 else [node_lines[0]] # Process each child of the class body individually body_parts: list[str] = [] processed_ranges: list[tuple[int, int]] = [] for child in body_node.children: # Use line-based extraction for children too child_start = child.start_point[0] child_end = child.end_point[0] child_text = "\n".join(code_lines[child_start : child_end + 1]) # Methods/functions inside the class — compress individually if child.type in lang_config.function_nodes: compressed = self._compress_function_ast( child, code, language, lang_config, body_limits, analysis ) body_parts.append(compressed) processed_ranges.append((child.start_byte, child.end_byte)) # Decorated methods elif lang_config.decorator_node and child.type == lang_config.decorator_node: decorator_lines = [] method_compressed = None for deco_child in child.children: if deco_child.type == "decorator": decorator_lines.append(_get_node_text(deco_child, code)) elif deco_child.type in lang_config.function_nodes: method_compressed = self._compress_function_ast( deco_child, code, language, lang_config, body_limits, analysis ) if decorator_lines and method_compressed: body_parts.append("\n".join(decorator_lines) + "\n" + method_compressed) elif method_compressed: body_parts.append(method_compressed) else: body_parts.append(child_text) processed_ranges.append((child.start_byte, child.end_byte)) # Nested classes — recurse elif child.type in lang_config.class_nodes: compressed = self._compress_class_ast( child, code, language, lang_config, body_limits, analysis ) body_parts.append(compressed) processed_ranges.append((child.start_byte, child.end_byte)) else: # Class-level attributes, type annotations, docstrings, etc. # Keep them as-is with original indentation if child_text.strip(): body_parts.append(child_text) processed_ranges.append((child.start_byte, child.end_byte)) # Reconstruct class with proper indentation result_parts = list(header_lines) for part in body_parts: result_parts.append(part) # Handle closing brace for brace-delimited languages body_end_line = body_node.end_point[0] body_end_rel = body_end_line - node_start_line + 1 after_lines = node_lines[body_end_rel:] if after_lines: result_parts.extend(after_lines) elif not lang_config.uses_colon_after_signature: # Ensure closing brace last_body_line = node_lines[-1] if node_lines else "" if last_body_line.strip() == "}": result_parts.append(last_body_line) return "\n".join(result_parts) def _extract_generic_structure(self, root: Any, code: str) -> CodeStructure: """Extract structure from generic/unknown code. For languages without a LangConfig, we can't reliably separate imports from other code. Just preserve everything in 'other'. """ structure = CodeStructure() structure.other = code.split("\n") return structure def _assemble_compressed( self, structure: CodeStructure, language: CodeLanguage, ) -> str: """Assemble compressed code from structure.""" parts: list[str] = [] # Imports first if structure.imports: parts.extend(structure.imports) parts.append("") # Type definitions if structure.type_definitions: parts.extend(structure.type_definitions) parts.append("") # Class definitions if structure.class_definitions: parts.extend(structure.class_definitions) parts.append("") # Function signatures/definitions if structure.function_signatures: parts.extend(structure.function_signatures) parts.append("") # Top-level code (global variables, constants, if __name__, etc.) if structure.top_level_code: parts.extend(structure.top_level_code) parts.append("") # Other content (used by generic extraction) if structure.other: parts.extend(structure.other) # Remove trailing empty lines while parts and not parts[-1].strip(): parts.pop() return "\n".join(parts) def _verify_syntax(self, code: str, language: CodeLanguage) -> bool: """Verify that code is syntactically valid. Checks for both ERROR nodes (parse failures) and MISSING nodes (tokens the parser expected but didn't find). """ try: parser = _get_parser(language.value) tree = parser.parse(bytes(code, "utf-8")) return not _has_syntax_issues(tree.root_node) except Exception: return False def _fallback_compress(self, code: str, original_tokens: int) -> CodeCompressionResult: """Fall back to LLMLingua compression.""" try: from .llmlingua_compressor import LLMLinguaCompressor, _check_llmlingua_available if _check_llmlingua_available(): compressor = LLMLinguaCompressor() result = compressor.compress(code, content_type="code") return CodeCompressionResult( compressed=result.compressed, original=code, original_tokens=result.original_tokens, compressed_tokens=result.compressed_tokens, compression_ratio=result.compression_ratio, language=CodeLanguage.UNKNOWN, language_confidence=0.0, # LLMLingua does NOT guarantee syntax validity syntax_valid=False, ) except ImportError: pass # No fallback available, return original return CodeCompressionResult( compressed=code, original=code, original_tokens=original_tokens, compressed_tokens=original_tokens, compression_ratio=1.0, language=CodeLanguage.UNKNOWN, language_confidence=0.0, syntax_valid=True, ) def _store_in_ccr( self, original: str, compressed: str, original_tokens: 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_tokens=original_tokens, compressed_tokens=self._estimate_tokens(compressed), compression_strategy="code_aware", ) except ImportError: return None except Exception as e: logger.debug("CCR storage failed: %s", e) return None def apply( self, messages: list[dict[str, Any]], tokenizer: Tokenizer, **kwargs: Any, ) -> TransformResult: """Apply code-aware compression to messages. Handles both string content and Anthropic content block format (list of {"type": "text", "text": "..."} dicts). Args: messages: List of message dicts to transform. tokenizer: Tokenizer for accurate token counting. **kwargs: Additional arguments (e.g., 'context'). Returns: TransformResult with compressed messages and metadata. """ tokens_before = sum(tokenizer.count_text(str(m.get("content", ""))) for m in messages) context = kwargs.get("context", "") transformed_messages = [] transforms_applied: list[str] = [] warnings: list[str] = [] for message in messages: content = message.get("content", "") # Handle content blocks (Anthropic format) if isinstance(content, list): new_blocks = [] for block in content: if isinstance(block, dict) and block.get("type") == "text": text = block.get("text", "") compressed_text = self._try_compress_text( text, context, tokenizer, transforms_applied ) new_blocks.append({**block, "text": compressed_text}) else: new_blocks.append(block) transformed_messages.append({**message, "content": new_blocks}) continue # Handle string content if not content or not isinstance(content, str): transformed_messages.append(message) continue compressed_content = self._try_compress_text( content, context, tokenizer, transforms_applied ) if compressed_content != content: transformed_messages.append({**message, "content": compressed_content}) else: transformed_messages.append(message) tokens_after = sum( tokenizer.count_text(str(m.get("content", ""))) for m in transformed_messages ) if not _check_tree_sitter_available(): warnings.append( "tree-sitter not installed. Install with: pip install headroom-ai[code]" ) return TransformResult( messages=transformed_messages, tokens_before=tokens_before, tokens_after=tokens_after, transforms_applied=transforms_applied if transforms_applied else ["code_aware:noop"], warnings=warnings, ) def _try_compress_text( self, text: str, context: str, tokenizer: Tokenizer, transforms_applied: list[str], ) -> str: """Try to compress a text string if it contains code.""" from .content_detector import ContentType, detect_content_type if not text: return text detection = detect_content_type(text) if detection.content_type == ContentType.SOURCE_CODE: language = detection.metadata.get("language") result = self.compress(text, language=language, context=context, tokenizer=tokenizer) if result.compression_ratio < 0.9: transforms_applied.append( f"code_aware:{result.language.value}:{result.compression_ratio:.2f}" ) return result.compressed return text def should_apply( self, messages: list[dict[str, Any]], tokenizer: Tokenizer, **kwargs: Any, ) -> bool: """Check if code-aware compression should be applied. Returns True if: - tree-sitter is available, AND - Content contains detected source code Args: messages: Messages to check. tokenizer: Tokenizer for counting. **kwargs: Additional arguments. Returns: True if compression should be applied. """ if not _check_tree_sitter_available(): return False from .content_detector import ContentType, detect_content_type for message in messages: content = message.get("content", "") # Handle string content if content and isinstance(content, str): detection = detect_content_type(content) if detection.content_type == ContentType.SOURCE_CODE: return True # Handle content blocks elif isinstance(content, list): for block in content: if isinstance(block, dict) and block.get("type") == "text": text = block.get("text", "") if text: detection = detect_content_type(text) if detection.content_type == ContentType.SOURCE_CODE: return True return False # ========================================================================= # Module-level helper functions (stateless, used by the class) # ========================================================================= def _get_node_text(node: Any, code: str) -> str: """Extract text from AST node.""" return code[node.start_byte : node.end_byte] def _get_definition_name(node: Any) -> str | None: """Extract the name identifier from a definition AST node.""" for child in node.children: if child.type in ("identifier", "name", "type_identifier", "property_identifier"): text = child.text return text.decode("utf-8") if isinstance(text, bytes) else str(text) return None def _is_public_symbol(name: str, language: CodeLanguage) -> bool: """Heuristic for whether a symbol is public/exported.""" if not name: return False if language == CodeLanguage.GO: return name[0].isupper() return not name.startswith("_") def _get_body_limit( func_name: str | None, body_limits: dict[str, int], max_body_lines: int, ) -> int: """Look up the allocated body line limit for a function. Falls back to max_body_lines if no budget allocation was computed. max_body_lines always acts as a hard cap. """ if body_limits and func_name and func_name in body_limits: return min(body_limits[func_name], max_body_lines) return max_body_lines def _make_omitted_comment( func_name: str | None, omitted_count: int, indent: str, comment_prefix: str, analysis: _SymbolAnalysis | None, ) -> str: """Build omitted comment with call information from analysis.""" calls_info = "" if analysis and func_name: for key in ( func_name, *(k for k in analysis.calls if k.endswith(f".{func_name}")), ): if key in analysis.calls: called = analysis.calls[key] if called: sorted_calls = sorted(called)[:5] calls_info = "; calls: " + ", ".join(sorted_calls) if len(called) > 5: calls_info += f" +{len(called) - 5} more" break return f"{indent}{comment_prefix} [{omitted_count} lines omitted{calls_info}]" def _detect_indent(lines: list[str]) -> str: """Detect the indentation used in a list of code lines.""" for line in lines: if line.strip(): return line[: len(line) - len(line.lstrip())] return " " def _has_syntax_issues(node: Any) -> bool: """Check if AST contains ERROR or MISSING nodes.""" if node.type == "ERROR" or node.is_missing: return True for child in node.children: if _has_syntax_issues(child): return True return False def compress_code( code: str, language: str | None = None, target_rate: float = 0.2, context: str = "", ) -> str: """Convenience function for one-off code compression. Args: code: Source code to compress. language: Language hint (auto-detected if None). target_rate: Target compression rate (0.2 = keep 20%). context: Optional context for relevance. Returns: Compressed code string. Example: >>> compressed = compress_code(large_python_file) >>> print(compressed) # Valid Python code """ config = CodeCompressorConfig( target_compression_rate=target_rate, language_hint=language, ) compressor = CodeAwareCompressor(config) result = compressor.compress(code, language=language, context=context) return result.compressed