mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(compression): correct JSON array item counting and entropy gate (#887)
## Description
Two bugs in `JSONStructureHandler` that jointly defeated the "keep first
N array items fully" design. (1) Every comma under `array_depth > 0` was
counted as an array item separator — including commas *between keys
inside objects* — so for arrays of objects the first record's own keys
exhausted `max_array_items_full` and dropped values belonging to item 0.
(2) Fixing that unmasked a second bug: self-normalized Shannon entropy
scores English prose at 0.90+, above the 0.85 "identifier" threshold, so
every long description was preserved as a fake high-entropy identifier.
Closes # <!-- found during a compression-handler review -->
## 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
- `headroom/compression/handlers/json_handler.py`: replace depth-keyed
comma counting with a container stack so only commas whose immediate
enclosing container is an array advance that array's item index.
- `headroom/compression/handlers/json_handler.py`: gate the entropy
preservation check on a no-spaces identifier signal, so UUIDs/hashes
still pass but prose compresses.
- `tests/test_compression/test_json_handler.py`: regression tests for
object-comma counting, items past the threshold, and prose-vs-identifier
entropy.
## 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_compression/test_json_handler.py -q
32 passed
```
## Real Behavior Proof
- Environment: local macOS, Python 3.11, branch
`fix/json-array-item-count`.
- Exact command / steps: `pytest
tests/test_compression/test_json_handler.py -q` plus an empirical mask
dump on `[{"a":1,"b":2,...}]`.
- Observed result: Values inside array item 0 are now preserved; long
prose values compress while UUIDs are retained (prose scored
0.906-0.929, UUID 0.956 — the threshold alone could not separate them).
- Not tested: End-to-end through the live proxy pipeline (the handler is
not yet wired into the proxy hot path).
## 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
## Screenshots (if applicable)
N/A — library/compression change with no UI. See Test Output.
## Additional Notes
First of a 7-PR compression-handler review series.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
parent
b70fccbe17
commit
d6f0f0f642
2 changed files with 105 additions and 18 deletions
|
|
@ -138,30 +138,44 @@ class JSONStructureHandler(BaseStructureHandler):
|
||||||
# Build character-level mask
|
# Build character-level mask
|
||||||
mask = [False] * len(content)
|
mask = [False] * len(content)
|
||||||
|
|
||||||
# Track array depth for selective preservation
|
# Track containers so commas inside objects are not counted as
|
||||||
array_depth = 0
|
# array item separators — only commas whose immediate enclosing
|
||||||
array_item_counts: dict[int, int] = {} # depth -> count
|
# container is an array advance that array's item index. Values
|
||||||
|
# nested in an object that is itself an array item inherit the
|
||||||
|
# innermost enclosing array's index via array_item_stack[-1].
|
||||||
|
container_stack: list[str] = [] # "[" or "{" per open container
|
||||||
|
array_item_stack: list[int] = [] # item count per open array
|
||||||
|
|
||||||
for token in json_tokens:
|
for token in json_tokens:
|
||||||
# Track array items
|
# Track containers
|
||||||
if token.token_type == JSONTokenType.BRACKET:
|
if token.token_type == JSONTokenType.BRACKET:
|
||||||
if token.text == "[":
|
if token.text in "{[":
|
||||||
array_depth += 1
|
container_stack.append(token.text)
|
||||||
array_item_counts[array_depth] = 0
|
if token.text == "[":
|
||||||
|
array_item_stack.append(0)
|
||||||
|
elif token.text == "}":
|
||||||
|
if container_stack and container_stack[-1] == "{":
|
||||||
|
container_stack.pop()
|
||||||
elif token.text == "]":
|
elif token.text == "]":
|
||||||
if array_depth in array_item_counts:
|
if container_stack and container_stack[-1] == "[":
|
||||||
del array_item_counts[array_depth]
|
container_stack.pop()
|
||||||
array_depth = max(0, array_depth - 1)
|
if array_item_stack:
|
||||||
|
array_item_stack.pop()
|
||||||
|
|
||||||
# Count array items at commas
|
# Count array items only at the array's own commas
|
||||||
if token.token_type == JSONTokenType.COMMA and array_depth > 0:
|
if (
|
||||||
array_item_counts[array_depth] = array_item_counts.get(array_depth, 0) + 1
|
token.token_type == JSONTokenType.COMMA
|
||||||
|
and container_stack
|
||||||
|
and container_stack[-1] == "["
|
||||||
|
and array_item_stack
|
||||||
|
):
|
||||||
|
array_item_stack[-1] += 1
|
||||||
|
|
||||||
# Determine if this token should be preserved
|
# Determine if this token should be preserved
|
||||||
preserve = self._should_preserve_token(
|
preserve = self._should_preserve_token(
|
||||||
token,
|
token,
|
||||||
array_depth,
|
len(array_item_stack),
|
||||||
array_item_counts.get(array_depth, 0),
|
array_item_stack[-1] if array_item_stack else 0,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Mark in mask
|
# Mark in mask
|
||||||
|
|
@ -225,9 +239,13 @@ class JSONStructureHandler(BaseStructureHandler):
|
||||||
if self.preserve_high_entropy:
|
if self.preserve_high_entropy:
|
||||||
# Strip quotes for entropy calculation
|
# Strip quotes for entropy calculation
|
||||||
value = token.text.strip('"')
|
value = token.text.strip('"')
|
||||||
score = EntropyScore.compute(value, self.entropy_threshold)
|
# Entropy targets identifiers (UUIDs, hashes, API keys).
|
||||||
if score.should_preserve:
|
# Self-normalized entropy also scores English prose >0.85,
|
||||||
return True
|
# so gate on the cheapest identifier signal: no spaces.
|
||||||
|
if " " not in value:
|
||||||
|
score = EntropyScore.compute(value, self.entropy_threshold)
|
||||||
|
if score.should_preserve:
|
||||||
|
return True
|
||||||
|
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -153,6 +153,75 @@ class TestJSONStructureHandler:
|
||||||
assert result.mask.mask[first_id] is True
|
assert result.mask.mask[first_id] is True
|
||||||
assert result.mask.mask[second_id] is True
|
assert result.mask.mask[second_id] is True
|
||||||
|
|
||||||
|
def test_object_commas_do_not_advance_array_item_index(self):
|
||||||
|
"""Commas between keys inside an object must not count as array
|
||||||
|
item separators.
|
||||||
|
|
||||||
|
Regression: depth-keyed comma counting treated every comma under
|
||||||
|
array_depth > 0 as an item boundary, so an array's FIRST object
|
||||||
|
exhausted max_array_items_full within its own keys and later
|
||||||
|
values were dropped despite being in item 0.
|
||||||
|
"""
|
||||||
|
handler = JSONStructureHandler(
|
||||||
|
preserve_short_values=True,
|
||||||
|
short_value_threshold=20,
|
||||||
|
max_array_items_full=3,
|
||||||
|
)
|
||||||
|
# Single array item (index 0) with 5 short values — all must
|
||||||
|
# stay eligible for preservation.
|
||||||
|
content = '[{"a": "valuea", "b": "valueb", "c": "valuec", "d": "valued", "e": "valuee"}]'
|
||||||
|
result = handler.get_mask(content)
|
||||||
|
|
||||||
|
for marker in ('"valuea"', '"valueb"', '"valuec"', '"valued"', '"valuee"'):
|
||||||
|
start = content.index(marker)
|
||||||
|
for i in range(start, start + len(marker)):
|
||||||
|
assert result.mask.mask[i] is True, f"{marker} in array item 0 should be preserved"
|
||||||
|
|
||||||
|
def test_array_items_past_threshold_compressed(self):
|
||||||
|
"""Items at index >= max_array_items_full are still compressed."""
|
||||||
|
handler = JSONStructureHandler(
|
||||||
|
preserve_short_values=True,
|
||||||
|
short_value_threshold=20,
|
||||||
|
max_array_items_full=2,
|
||||||
|
)
|
||||||
|
content = '[{"v": "itemzero"}, {"v": "itemone"}, {"v": "itemtwo"}]'
|
||||||
|
result = handler.get_mask(content)
|
||||||
|
|
||||||
|
# Items 0 and 1 preserved
|
||||||
|
for marker in ('"itemzero"', '"itemone"'):
|
||||||
|
start = content.index(marker)
|
||||||
|
assert result.mask.mask[start + 1] is True, f"{marker} should be preserved"
|
||||||
|
|
||||||
|
# Item 2 is past the threshold — value chars compressed
|
||||||
|
start = content.index('"itemtwo"')
|
||||||
|
inner = range(start + 1, start + len('"itemtwo"') - 1)
|
||||||
|
assert not any(result.mask.mask[i] for i in inner), (
|
||||||
|
"values in array items past max_array_items_full should be compressible"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_prose_not_preserved_by_entropy(self):
|
||||||
|
"""Long natural-language strings must not pass the entropy gate.
|
||||||
|
|
||||||
|
Regression: self-normalized Shannon entropy scores English prose
|
||||||
|
above the 0.85 threshold, so every description survived
|
||||||
|
compression. Identifiers (no spaces) are still preserved.
|
||||||
|
"""
|
||||||
|
handler = JSONStructureHandler(short_value_threshold=20)
|
||||||
|
prose = "Context optimization layer for LLM applications with caching"
|
||||||
|
uuid = "8f14e45f-ceea-4123-8f14-e45fceea4123"
|
||||||
|
content = f'{{"description": "{prose}", "id": "{uuid}"}}'
|
||||||
|
result = handler.get_mask(content)
|
||||||
|
|
||||||
|
prose_start = content.index(prose)
|
||||||
|
assert not any(result.mask.mask[i] for i in range(prose_start, prose_start + len(prose))), (
|
||||||
|
"long prose value should be compressible"
|
||||||
|
)
|
||||||
|
|
||||||
|
uuid_start = content.index(uuid)
|
||||||
|
assert all(result.mask.mask[i] for i in range(uuid_start, uuid_start + len(uuid))), (
|
||||||
|
"UUID value should be preserved via entropy"
|
||||||
|
)
|
||||||
|
|
||||||
def test_metadata_contains_key_count(self, handler):
|
def test_metadata_contains_key_count(self, handler):
|
||||||
"""Test that metadata includes key count."""
|
"""Test that metadata includes key count."""
|
||||||
content = '{"a": 1, "b": 2, "c": 3}'
|
content = '{"a": 1, "b": 2, "c": 3}'
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue