Code Compression

AST-aware compression that preserves imports, signatures, and types while compressing function bodies. Powered by tree-sitter.

Headroom's CodeAwareCompressor uses tree-sitter to parse source code into an AST, then selectively compresses function bodies while preserving the structural elements that LLMs need -- imports, signatures, type annotations, and error handlers.

Why AST-Aware Compression?

Naive truncation breaks code. Cutting a function in half leaves invalid syntax that confuses the LLM. CodeAwareCompressor guarantees:

Syntax validity -- output always parses correctly
Structural preservation -- imports, signatures, types, decorators are kept intact
Lightweight -- ~50MB of tree-sitter parsers, loaded lazily and cached

Supported Languages

Tier	Languages	Support Level
Tier 1	Python, JavaScript, TypeScript	Full AST analysis
Tier 2	Go, Rust, Java, C, C++, C#, PHP	Function body compression

What Gets Preserved vs Compressed

Always preserved:

Import statements
Function and method signatures
Class definitions
Type annotations
Decorators

Compressed:

Function bodies (implementations)
Comments (unless configured to preserve)
Verbose docstrings (configurable: full, first line, or removed)

Example

from headroom.transforms import CodeAwareCompressor, CodeCompressorConfig

# min_tokens_for_compression defaults to 100; a snippet this small needs an
# explicit override or it passes through unmodified (see "Skip small content" below).
compressor = CodeAwareCompressor(CodeCompressorConfig(min_tokens_for_compression=10))

code = '''
import os
from typing import List

def process_items(items: List[str]) -> List[str]:
    """Process a list of items."""
    results = []
    for item in items:
        if not item:
            continue
        processed = item.strip().lower()
        results.append(processed)
    return results
'''

result = compressor.compress(code, language="python")
print(result.compressed)
# import os
# from typing import List
#
# def process_items(items: List[str]) -> List[str]:
#     """Process a list of items."""
#     results = []
#     # [6 lines omitted]
#     pass
# # [32 tokens compressed. Retrieve more: hash=<cache_key>. Expires in 5m.]

print(f"Compression: {result.compression_ratio:.0%}")  # ~57%
print(f"Syntax valid: {result.syntax_valid}")  # True

Configuration

from headroom.transforms import CodeAwareCompressor, CodeCompressorConfig, DocstringMode

config = CodeCompressorConfig(
    preserve_imports=True,              # Always keep imports
    preserve_signatures=True,           # Always keep function signatures
    preserve_type_annotations=True,     # Keep type hints
    preserve_decorators=True,           # Keep decorators
    docstring_mode=DocstringMode.FIRST_LINE,  # FULL, FIRST_LINE, REMOVE
    target_compression_rate=0.2,        # Keep 20% of tokens
    max_body_lines=5,                   # Lines to keep per function body
    min_tokens_for_compression=100,     # Skip small content
    language_hint=None,                 # Auto-detect if None
    fallback_to_kompress=True,          # Use Kompress for unknown langs
)

compressor = CodeAwareCompressor(config)
result = compressor.compress(code)

Configuration Options

Option	Default	Description
`preserve_imports`	`True`	Keep all import statements
`preserve_signatures`	`True`	Keep function/method signatures
`preserve_type_annotations`	`True`	Keep type hints
`preserve_decorators`	`True`	Keep decorators
`docstring_mode`	`FIRST_LINE`	How to handle docstrings: `FULL`, `FIRST_LINE`, `REMOVE`
`target_compression_rate`	`0.2`	Fraction of tokens to keep (0.2 = keep 20%)
`max_body_lines`	`5`	Max lines to keep per function body
`min_tokens_for_compression`	`100`	Skip files smaller than this
`language_hint`	`None`	Override language detection
`fallback_to_kompress`	`True`	Use Kompress for unsupported languages

Before and After

# Before (full source file)
def process_data(items: List[str]) -> Dict[str, int]:
    """Process items and count occurrences."""
    result = {}
    for item in items:
        item = item.strip().lower()
        if item in result:
            result[item] += 1
        else:
            result[item] = 1
    return result

# After (signature preserved, body compressed)
def process_data(items: List[str]) -> Dict[str, int]:
    """Process items and count occurrences."""
    result = {}
    for item in items:
    # ... (5 lines compressed)
    pass

The LLM sees the function's purpose, its input/output types, and the general approach -- enough to reason about the code without needing every implementation line.

Installation

# Install tree-sitter language pack
pip install "headroom-ai[code]"

Memory Management

Tree-sitter parsers are lazy-loaded and cached. You can free memory when done:

from headroom.transforms import is_tree_sitter_available
from headroom.transforms.code_compressor import unload_tree_sitter

# Check if tree-sitter is installed
print(is_tree_sitter_available())  # True

# Free memory when done
unload_tree_sitter()

Performance

Metric	Value
Compression	40-70% token reduction
Speed	~10-50ms per file
Memory	~50MB (tree-sitter parsers)
Syntax validity	Guaranteed

Routing differs between the proxy and the library API

ContentRouter always detects source code automatically, but whether it routes that code to CodeAwareCompressor depends on the entry point. headroom proxy enables code-aware routing by default (HEADROOM_CODE_AWARE_ENABLED defaults to on; disable with --no-code-aware or HEADROOM_CODE_AWARE_ENABLED=0). The one-function compress() API and HeadroomClient, however, build ContentRouterConfig with its own default, enable_code_aware=False -- there, detected code falls back to Kompress unless you construct your own pipeline with ContentRouterConfig(enable_code_aware=True). Direct usage of CodeAwareCompressor, as shown on this page, is unaffected either way. See How Compression Works.

On this page