mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-10 14:27:00 -04:00
79baee082f
PR #431 (merged) added text-block compression to support DeepSeek + Cline, but the gate ("skip user/system") leaves assistant text blocks compressible by default. Assistant content is echoed back by the client in subsequent turns and becomes part of the upstream provider's prefix cache (Anthropic explicit cache_control, DeepSeek/OpenAI auto-prefix). Compressing it silently changes the bytes the next turn must match for a cache hit — turning a 90% read discount into a 25% write penalty on Anthropic, or a full prefill on DeepSeek/OpenAI when the in-process result cache evicts or differs across restarts. Re-aligns the design around prefix-cache safety: * Block-level cache_control protection (defense in depth). Any block carrying cache_control is the client's explicit cache breakpoint; never modified, regardless of role or block type. Closes the gap that frozen_message_count alone leaves — that count is a coarse message-level approximation; this is the per-block guarantee. Applies to both tool_result and text paths. * compress_assistant_text_blocks defaults to False (off). Assistant text blocks are skipped by default, restoring pre-#431 cache safety for Anthropic flows. Per-request opt-in via kwargs (or via ContentRouterConfig.compress_assistant_text_blocks for deployment- wide enable) preserves the Cline + DeepSeek goal — only enable when the backend doesn't honor cache_control AND compression is deterministic enough that the auto-prefix cache still hits across eviction/restart. * Unknown roles default-skip too (was: compressed). developer/judge/ custom roles are safer to leave untouched than to compress aggressively without thinking through their cache semantics. * Online streaming usage parser. Replaces the per-stream list[bytes] buffer with a single last_completion_tokens int updated per chunk via a module-level _parse_completion_tokens_from_sse_chunk helper. Streaming memory is now O(1) regardless of stream length — important for 200K-output reasoning models and DeepSeek V4 Pro's 384K max output. * Renames the unused min_tokens parameter to min_chars (the threshold has always been chars, not tokens, in both the tool_result and text paths). Now also wired through ContentRouterConfig .min_chars_for_block_compression so the threshold is configurable per Realignment build constraints. Tests: * 17 new tests in tests/test_transforms_content_router.py covering the role matrix (user / system / assistant / tool / unknown), cache_control protection on both paths, opt-in semantics, the min_chars threshold, and idempotent pinning detection. * 9 new tests in tests/test_streaming_usage_parser.py covering the online parser's success and edge cases (usage frame, [DONE], invalid JSON, multi-frame chunks, zero tokens, non-dict payloads, invalid UTF-8). Trade-off: deployments pointed at non-cache-aware backends (DeepSeek direct, OpenAI direct) lose blanket assistant-text compression by default — they opt in via config. Anthropic flows go back to being prefix-cache-safe out of the box.
60 lines
2.3 KiB
Python
60 lines
2.3 KiB
Python
"""Unit tests for the online SSE usage parser used by the
|
|
OpenAI-via-backend streaming path.
|
|
|
|
These tests pin the per-chunk parsing contract so streaming memory
|
|
stays O(1) regardless of stream length — the prior implementation
|
|
buffered the entire response just to scan the trailing usage frame.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from headroom.proxy.handlers.streaming import _parse_completion_tokens_from_sse_chunk
|
|
|
|
|
|
def test_returns_completion_tokens_from_usage_frame() -> None:
|
|
chunk = b'data: {"id":"x","usage":{"prompt_tokens":10,"completion_tokens":42}}\n\n'
|
|
assert _parse_completion_tokens_from_sse_chunk(chunk) == 42
|
|
|
|
|
|
def test_returns_none_for_content_only_chunk() -> None:
|
|
chunk = b'data: {"choices":[{"delta":{"content":"hi"}}]}\n\n'
|
|
assert _parse_completion_tokens_from_sse_chunk(chunk) is None
|
|
|
|
|
|
def test_returns_none_for_done_marker() -> None:
|
|
assert _parse_completion_tokens_from_sse_chunk(b"data: [DONE]\n\n") is None
|
|
|
|
|
|
def test_returns_none_for_invalid_json() -> None:
|
|
assert _parse_completion_tokens_from_sse_chunk(b"data: not-json\n\n") is None
|
|
|
|
|
|
def test_returns_none_for_empty_chunk() -> None:
|
|
assert _parse_completion_tokens_from_sse_chunk(b"") is None
|
|
|
|
|
|
def test_handles_chunk_with_multiple_frames() -> None:
|
|
# SSE frames can batch across a single chunk write.
|
|
chunk = (
|
|
b'data: {"choices":[{"delta":{"content":"a"}}]}\n\n'
|
|
b'data: {"choices":[{"delta":{"content":"b"}}],"usage":{"completion_tokens":7}}\n\n'
|
|
)
|
|
assert _parse_completion_tokens_from_sse_chunk(chunk) == 7
|
|
|
|
|
|
def test_treats_zero_completion_tokens_as_zero_not_none() -> None:
|
|
chunk = b'data: {"usage":{"completion_tokens":0}}\n\n'
|
|
assert _parse_completion_tokens_from_sse_chunk(chunk) == 0
|
|
|
|
|
|
def test_handles_non_dict_data_payload() -> None:
|
|
# Edge case: a JSON array or scalar where a dict was expected.
|
|
chunk = b"data: [1,2,3]\n\n"
|
|
assert _parse_completion_tokens_from_sse_chunk(chunk) is None
|
|
|
|
|
|
def test_handles_invalid_utf8_bytes_without_crashing() -> None:
|
|
# Leading invalid UTF-8 bytes corrupt the "data: " prefix; parser
|
|
# should skip the malformed line and return None rather than raise.
|
|
chunk = b'\xff\xfedata: {"usage":{"completion_tokens":3}}\n\n'
|
|
assert _parse_completion_tokens_from_sse_chunk(chunk) is None
|