TokenShrinkWrapper

Shrink text you send to LLMs: fewer tokens via a fast local pipeline, code-friendly handling of fenced blocks and preprocessor-style # lines, optional OpenAI-compatible distillation, and tiktoken counts for budgeting.

PyPI · Source

What it’s for

• Agent / RAG dumps — long markdown design notes mixed with pasted code; you want smaller prompts without mangling snippets.
• Repo context exports — AGENTS.md, architecture docs, or generated context files before a model reads them.
• Cost control — measure before/after with the same tokenizer you use for API estimates (tiktoken).
• Not a substitute for learned compressors (LLMLingua‑style stacks trade weight and infra for heavier compression). TokenShrinkWrapper stays tiny: typer + tiktoken (+ stdlib HTTP if you distill).

How it works (local shrink)

Documents are split into prose vs code-like chunks:

Kind	Detection	Local transforms
Code	Inside Markdown triple‑backtick fences, or a short heuristic for naked source (shebang, #include, def / import …, directive-like #, etc.)	Strip safely skippable full-line #////single-line /*…*/, cap blank runs — never caveman strip
Prose	Everything else	Drop # headings, collapse whitespace, caveman (stop‑word shred), aggressive mode adds extractive line sampling (25% retention)

Modes: light → balanced → aggressive (progressively tighter on prose; code paths differ mainly by blank-line caps).

Optional --distill runs an OpenAI‑compatible chat completion on the (optionally) prepassed text (tokenshrink.toml → [llm]), then compares token counts vs the original — counts API tokens for your bill; distillation itself uses the model.

Installation

From PyPI:

pip install tokenshrinkwrapper

From git (contributors):

pip install -e ".[dev]"

CLI

# Shrink file to stdout or --output
tokenshrink shrink input.md -o shrunk.md --mode balanced

# Token estimate
tokenshrink stats input.md

# Compare local modes (before/after/saved/ratio via tiktoken)
tokenshrink benchmark sample.md

# LLM distill (needs OPENAI_API_KEY or configured env)
tokenshrink shrink input.md -o out.md --distill

Configuration

Optional tokenshrink.toml (or --config). Example shipped as tokenshrink.toml in the repo.

Python API

from tokenshrinkwrapper import shrink_text, estimate_tokens, benchmark_modes
from tokenshrinkwrapper.config import LLMSection
from tokenshrinkwrapper.llm import OpenAICompatibleClient

text = "## Long blob\n...\n"

r = shrink_text(text, mode="balanced")
print(r.token_count_before, "->", r.token_count_after)
print(r.compressed_text)

estimate_tokens(text, model="gpt-4o-mini")

for row in benchmark_modes(text):
    mode, before, after, ratio = row
    print(mode, before, after, ratio)

llm_section = LLMSection()  # or load via resolve_config / merge_app_config from TOML
llm = OpenAICompatibleClient.from_env(llm_section)
distilled = shrink_text(
    text,
    mode="balanced",
    distill=True,
    distill_prepass="light",
    llm=llm,
)

Benchmarks (representative fixtures)

Measured with gpt-4o-mini tiktoken. Δtokens below is legacy output − current output — positive means the current pipeline produced fewer tokens than a replay of the older uniform pipeline (same headings + whitespace + caveman + extractive 0.4 on the entire blob).

Corpus	Input tok	balanced Δ	aggressive Δ	Notes
Prose chunks (~14k chars)	~1 425	0	0	Mostly one long line after caveman → extractive 0.25 vs 0.4 rarely differs
Mixed MD + fenced Python (×8)	~768	+64	+64	Current ~11% fewer tokens vs legacy output; legacy drops for in fenced code
Bare Python heuristic	~460	−10	−10	Current more conservative on pure code (fewer removals) → higher token counts by design
C-style header snippet	~195	−50	−50	Current preserves #include / #define; legacy strips them (integrity over ratio)

Median wall-clock (25 runs, aggressive): mixed MD ~−29 % vs legacy; bare-Python fixture ~−60 %; tiny C blob ~−43 %; prose-only slightly ~+5 % (fence/heuristic overhead).

Reproduce

git clone https://github.com/dferlemann/TokenShrinkWrapper.git
cd TokenShrinkWrapper
pip install -e .
python benchmarks/run.py

The same file ships inside the sdist on PyPI (not in the wheel) if you download the source archive.

CLI spot check:

tokenshrink benchmark path/to/your-context.md

Limitations

• Heuristic segmentation — best results when code is in Markdown triple-backtick fenced blocks. Naked blobs use a heuristic; edge cases exist.
• No full lexer — full-line # removal can be wrong inside odd formats; preprocessor / coding-cookie-ish # lines are kept via patterns.
• Distill quality depends on provider, prompt, cost, and latency.
• stats/benchmark/local shrink measure with tiktoken; other providers tokenize differently.

Related PyPI tooling

Rough axes: llmlingua — model-heavy compression · tiktoken — counting · smaller heuristic packages (token-diet, llm-token-optimizer, etc.) vary in maturity.

Publishing to PyPI (maintainers)

1. Bump version in pyproject.toml, update changelog notes in git tag message if desired.
2. Build and inspect:

pip install -e ".[publish]"
rm -rf dist/ build/
python -m build
python -m twine check dist/*

3. Upload (use API token, not password):

python -m twine upload dist/*

4. First-time name: ensure tokenshrinkwrapper is free on PyPI; twine will error if taken.

First upload may require trusted publishing or username __token__ with --repository pypi.

Development

python -m venv .venv
source .venv/bin/activate
pip install -e ".[dev]"
pytest
python benchmarks/run.py

License

MIT