mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
4e30dde2ac
## Description
`ContentRouter` counts section tokens with `len(content.split())`. On
compact machine-generated JSON — the default output of
`json.dumps(separators=(",", ":"))`, `JSON.stringify`, and boto3 — there
are no spaces, so a large payload counts as ~1 "token". Every section
compression ratio then computes as ~1.0 and the `min_ratio` acceptance
gate silently rejects the compressor's real output: the router logs
`router:noop` while SmartCrusher separately logs `was_modified=true`.
Compression effectively no-ops on the most common agent payload type
(tool results returning JSON), on every provider.
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)
## Changes Made
- Add `_estimate_tokens(text)` — a size-proportional estimate
(`len(text) // 4`, floored at 1), monotone in content size for any
format.
- Replace the decision-relevant `len(...split())` counts in
`ContentRouter` (section original/compressed token counts feeding the
ratio gates, plus the debug estimates) with `_estimate_tokens(...)`.
- Add `tests/test_content_router_compact_json.py`.
## Testing
- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed
### Test Output
```text
$ pytest tests/test_content_router_compact_json.py -q
2 passed, 1 warning
$ ruff check headroom/transforms/content_router.py tests/test_content_router_compact_json.py
All checks passed!
$ mypy headroom/transforms/content_router.py --ignore-missing-imports
Success: no issues found in 1 source file
```
## Real Behavior Proof
- Environment: `ContentRouter` invoked directly on a 150-item
ECS-service JSON tool_result, Python 3.13, estimator tokenizer.
- Exact command / steps: run the same payload two ways — compact
(`json.dumps(..., separators=(",", ":"))`) and the identical data with
spaces (`separators=(", ", ": ")`) — through
`ContentRouter(ContentRouterConfig(skip_user_messages=False))`.
- Observed result: before this change, compact JSON saved 0.0%
(`router:noop`) while the identical data with spaces saved 43.3%
(`router:tool_result:smart_crusher`) — same data, same compressor, only
whitespace differed. After this change, compact JSON compresses
equivalently to the spaced form.
- Not tested: no behavior change expected for content that already
tokenizes with whitespace (prose, code); those counts move from
word-count to chars/4 but the ratio comparison is self-consistent (both
sides use the same estimator).
## Review Readiness
- [x] I have performed a self-review
- [x] This PR is ready for human review
## Checklist
- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [ ] I have updated the CHANGELOG.md if applicable
## Additional Notes
Scope kept deliberately narrow: only the counts that feed
compression-acceptance decisions are changed. Non-decision `.split()`
uses elsewhere are left alone. Happy to add a CHANGELOG entry if you'd
like one.
---------
Co-authored-by: Tejas Chopra <chopratejas@gmail.com>
71 lines
2.4 KiB
Python
71 lines
2.4 KiB
Python
"""Compact machine-generated JSON must not evade compression (token-estimate bug).
|
|
|
|
Whitespace-split token counting made a compact JSON payload (no spaces —
|
|
the default output of json.dumps with separators, JSON.stringify, boto3)
|
|
count as ~1 token, so compression ratios computed as ~1.0 and the
|
|
min_ratio gate rejected SmartCrusher's real output.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import random
|
|
|
|
import pytest
|
|
|
|
from headroom.transforms.content_router import (
|
|
ContentRouter,
|
|
ContentRouterConfig,
|
|
_estimate_tokens,
|
|
)
|
|
|
|
|
|
def test_estimate_tokens_monotone_on_compact_json() -> None:
|
|
small = json.dumps([{"a": 1}] * 5, separators=(",", ":"))
|
|
large = json.dumps([{"a": 1}] * 500, separators=(",", ":"))
|
|
assert _estimate_tokens(large) > _estimate_tokens(small) > 1
|
|
|
|
|
|
def test_compact_json_tool_result_compresses() -> None:
|
|
tokenizer = pytest.importorskip("headroom.tokenizers.estimator")
|
|
from headroom.tokenizer import Tokenizer
|
|
|
|
tok = Tokenizer(tokenizer.EstimatingTokenCounter())
|
|
random.seed(42)
|
|
rows = [
|
|
{
|
|
"serviceArn": f"arn:aws:ecs:us-east-1:123456789012:service/x/svc-{s:03d}",
|
|
"serviceName": f"svc-{s:03d}",
|
|
"status": "ACTIVE",
|
|
"desiredCount": random.randint(1, 6),
|
|
"runningCount": random.randint(0, 6),
|
|
}
|
|
for s in range(150)
|
|
]
|
|
payload = json.dumps(rows, separators=(",", ":"))
|
|
assert " " not in payload[:200] # genuinely compact
|
|
messages = [
|
|
{"role": "user", "content": "Investigate."},
|
|
{
|
|
"role": "assistant",
|
|
"content": [
|
|
{"type": "text", "text": "Checking."},
|
|
{"type": "tool_use", "id": "toolu_1", "name": "list_services", "input": {}},
|
|
],
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": payload}],
|
|
},
|
|
{"role": "user", "content": "Summarize in one sentence."},
|
|
]
|
|
router = ContentRouter(ContentRouterConfig(skip_user_messages=False))
|
|
before = tok.count_messages(messages)
|
|
result = router.apply(
|
|
[json.loads(json.dumps(m)) for m in messages],
|
|
tok,
|
|
context="Summarize",
|
|
frozen_message_count=0,
|
|
)
|
|
after = tok.count_messages(result.messages)
|
|
assert after < before * 0.9, (before, after, result.transforms_applied[:5])
|