diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 86e25fa1e..a97d20f63 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -1414,6 +1414,11 @@ class OpenAIHandlerMixin: if not isinstance(items, list): return payload, False, 0, [], {}, [], 0 try: + from headroom.transforms.compression_batches import ( + CompressionBatchEntry, + build_compression_batches, + compress_batch_with_router, + ) from headroom.transforms.compression_units import ( CompressionUnit, RoutedCompressionUnit, @@ -1447,27 +1452,45 @@ class OpenAIHandlerMixin: ) return payload, False, 0, [], {}, [], 0 - def _slot_text(item: dict[str, Any]) -> tuple[str, tuple[str, int | None]] | None: + def _slot_texts(item: dict[str, Any]) -> list[tuple[str, tuple[str, int | None]]]: # Only tool-output items are eligible for in-place compression. # Message items (user/system/assistant) sit inside the request's # cacheable prefix; mutating them busts prefix caching on every # subsequent turn. Role-level guards in compression_units.py # remain as defense-in-depth. type_tag = item.get("type") - if type_tag in self.OPENAI_RESPONSES_OUTPUT_TYPES: - output_text = _responses_part_text(item.get("output")) - if output_text: - return output_text, ("output", None) - return None + if type_tag not in self.OPENAI_RESPONSES_OUTPUT_TYPES: + return [] + output = item.get("output") + if isinstance(output, str): + return [(output, ("output", None))] + if not isinstance(output, list): + return [] + return [ + (part["text"], ("output_part", index)) + for index, part in enumerate(output) + if isinstance(part, dict) + and part.get("type") in {"input_text", "output_text"} + and isinstance(part.get("text"), str) + ] def _set_slot_text( item: dict[str, Any], slot: tuple[str, int | None], replacement: str, - ) -> None: - kind, _ = slot + ) -> bool: + kind, index = slot if kind == "output": item["output"] = replacement + return True + if kind == "output_part" and isinstance(index, int): + output = item.get("output") + if isinstance(output, list) and 0 <= index < len(output): + part = output[index] + if isinstance(part, dict) and part.get("type") in {"input_text", "output_text"}: + part["text"] = replacement + return True + return False headroom_retrieve_call_ids: set[str] = set() # Map each Responses tool call to its name so that outputs belonging to @@ -1599,25 +1622,25 @@ class OpenAIHandlerMixin: } ) continue - slot = _slot_text(item) - if slot is not None: - text, slot_ref = slot - candidates.append((idx, slot_ref, text)) - if debug_enabled: - extraction_debug.append( - { - "index": idx, - "eligible": True, - "item_type": item_type, - "role": item.get("role"), - "slot": slot_ref, - "text_chars": len(text), - "text_bytes": len(text.encode("utf-8", errors="replace")), - "text_json_shape": _json_shape(text), - "item": item, - "text": text, - } - ) + slots = _slot_texts(item) + if slots: + for text, slot_ref in slots: + candidates.append((idx, slot_ref, text)) + if debug_enabled: + extraction_debug.append( + { + "index": idx, + "eligible": True, + "item_type": item_type, + "role": item.get("role"), + "slot": slot_ref, + "text_chars": len(text), + "text_bytes": len(text.encode("utf-8", errors="replace")), + "text_json_shape": _json_shape(text), + "item": item, + "text": text, + } + ) else: if debug_enabled: extraction_debug.append( @@ -1689,24 +1712,6 @@ class OpenAIHandlerMixin: unit_build_started = time.perf_counter() unit_debug: list[dict[str, Any]] = [] - # Aggregate-then-floor: the Responses payload splits each tool output - # into its own unit, so a per-item size floor would reject every unit - # in a session made of many small tool outputs (e.g. Codex), yielding - # 0% savings even when the combined compressible text is large. The - # Anthropic path compresses the whole message list as one batch and is - # not subject to a per-item floor. Match that: evaluate the floor once - # against the *aggregate* compressible bytes of the extracted group. If - # the group as a whole clears the threshold, disable the per-unit floor - # so small units still reach the router; if the whole group is below - # the threshold, keep the floor so trivially small payloads are skipped. - aggregate_compressible_bytes = sum( - len(text.encode("utf-8", errors="replace")) for _, _, text in candidates - ) - effective_unit_min_bytes = ( - 0 - if aggregate_compressible_bytes >= self.OPENAI_RESPONSES_ROUTER_MIN_BYTES - else self.OPENAI_RESPONSES_ROUTER_MIN_BYTES - ) for item_idx, slot_ref, original_text in candidates: item = items[item_idx] if item_idx < len(items) else {} item_type = item.get("type", "unknown") if isinstance(item, dict) else "unknown" @@ -1719,7 +1724,7 @@ class OpenAIHandlerMixin: item_type=str(item_type), cache_zone="live", mutable=True, - min_bytes=effective_unit_min_bytes, + min_bytes=self.OPENAI_RESPONSES_ROUTER_MIN_BYTES, ) routed_units.append(RoutedCompressionUnit(unit=unit, slot=(item_idx, slot_ref))) if debug_enabled: @@ -1773,9 +1778,25 @@ class OpenAIHandlerMixin: router_total_started = time.perf_counter() routed_results: list[tuple[object, Any, float] | None] = [None] * len(routed_units) + unit_index_by_slot = {routed.slot: unit_idx for unit_idx, routed in enumerate(routed_units)} + small_batch_entries: list[CompressionBatchEntry] = [] + large_unit_indexes: list[int] = [] + for unit_idx, routed in enumerate(routed_units): + text_bytes = len(routed.unit.text.encode("utf-8", errors="replace")) + if text_bytes < routed.unit.min_bytes: + small_batch_entries.append( + CompressionBatchEntry(entry_id=f"u{unit_idx}", routed=routed) + ) + else: + large_unit_indexes.append(unit_idx) + small_batches, small_batch_skipped = build_compression_batches( + small_batch_entries, + min_batch_bytes=self.OPENAI_RESPONSES_ROUTER_MIN_BYTES, + ) cache_misses: list[tuple[int, str, RoutedCompressionUnit]] = [] cache_miss_followers: dict[str, list[int]] = {} - for unit_idx, routed in enumerate(routed_units): + for unit_idx in large_unit_indexes: + routed = routed_units[unit_idx] cache_key = _openai_responses_unit_cache_key( routed.unit, model=model, @@ -1831,6 +1852,39 @@ class OpenAIHandlerMixin: _compress_and_store(unit_idx, cache_key, routed)[2], ) + # Tail batches below the shared 512B floor keep the existing size-floor + # result without entering the router or the unit-result cache. + for entry in small_batch_skipped: + unit_idx = unit_index_by_slot[entry.routed.slot] + routed_results[unit_idx] = _compress_routed_unit(entry.routed) + + def _compress_batch(batch: Any) -> tuple[list[tuple[object, Any]], float]: + batch_started = time.perf_counter() + results = compress_batch_with_router( + batch, + router=router, + tokenizer=tokenizer, + target_ratio=unit_target_ratio, + ) + return results, (time.perf_counter() - batch_started) * 1000.0 + + def _record_batch_result(batch_result: tuple[list[tuple[object, Any]], float]) -> None: + batch_results, elapsed_ms = batch_result + elapsed_per_unit = elapsed_ms / len(batch_results) if batch_results else 0.0 + for slot, result in batch_results: + routed_results[unit_index_by_slot[slot]] = (slot, result, elapsed_per_unit) + + if len(small_batches) > 1 and parallelism > 1: + executor = _openai_responses_unit_executor() + for start in range(0, len(small_batches), parallelism): + batch_group = small_batches[start : start + parallelism] + futures = [executor.submit(_compress_batch, batch) for batch in batch_group] + for future in as_completed(futures): + _record_batch_result(future.result()) + else: + for batch in small_batches: + _record_batch_result(_compress_batch(batch)) + ordered_routed_results = [result for result in routed_results if result is not None] for _, result, elapsed_ms in ordered_routed_results: @@ -1937,12 +1991,12 @@ class OpenAIHandlerMixin: target_item = updated_items[item_idx] if not isinstance(target_item, dict): continue - _set_slot_text(target_item, slot_ref, result.compressed) - modified = True - tokens_saved_total += result.tokens_saved - for transform in result.transforms_applied: - if transform not in transforms: - transforms.append(transform) + if _set_slot_text(target_item, slot_ref, result.compressed): + modified = True + tokens_saved_total += result.tokens_saved + for transform in result.transforms_applied: + if transform not in transforms: + transforms.append(transform) _add_timing("compression_unit_apply_results", apply_started) # Splice byte/data-lossless folds of excluded tool outputs (grep/log/ @@ -1952,8 +2006,8 @@ class OpenAIHandlerMixin: e_target = updated_items[e_idx] if e_idx < len(updated_items) else None if not isinstance(e_target, dict): continue - _set_slot_text(e_target, e_slot, e_folded) - modified = True + if _set_slot_text(e_target, e_slot, e_folded): + modified = True e_before = tokenizer.count_text(e_orig) e_saved = e_before - tokenizer.count_text(e_folded) if e_saved > 0: diff --git a/headroom/transforms/compression_batches.py b/headroom/transforms/compression_batches.py new file mode 100644 index 000000000..4fae8ebeb --- /dev/null +++ b/headroom/transforms/compression_batches.py @@ -0,0 +1,363 @@ +"""Bounded batching for small provider-extracted compression units.""" + +from __future__ import annotations + +import hashlib +import re +from dataclasses import dataclass + +from .compression_units import ( + _CCR_MARKER_RE, + _LOSSY_UNMARKED_STRATEGIES, + CompressionStrategy, + ContentRouter, + RoutedCompressionUnit, + TokenCounterLike, + UnitCompressionResult, + _is_structured_shell_output, +) +from .content_router import RouterCompressionResult +from .tag_protector import protect_tags, restore_tags + +DEFAULT_MAX_BATCH_BYTES = 2048 +DEFAULT_MAX_BATCH_UNITS = 16 + + +@dataclass(frozen=True) +class CompressionBatchEntry: + """One provider slot with a stable batch-local identifier.""" + + entry_id: str + routed: RoutedCompressionUnit + + +@dataclass(frozen=True) +class CompressionBatch: + """Compatible small units that share one future router invocation.""" + + entries: tuple[CompressionBatchEntry, ...] + text_bytes: int + + +def _text_bytes(text: str) -> int: + return len(text.encode("utf-8", errors="replace")) + + +def _compatibility_key(entry: CompressionBatchEntry) -> tuple[object, ...]: + unit = entry.routed.unit + return ( + unit.provider, + unit.endpoint, + unit.role, + unit.cache_zone, + unit.mutable, + unit.context, + unit.question, + unit.bias, + ) + + +def build_compression_batches( + entries: list[CompressionBatchEntry], + *, + min_batch_bytes: int, + max_batch_bytes: int = DEFAULT_MAX_BATCH_BYTES, + max_batch_units: int = DEFAULT_MAX_BATCH_UNITS, +) -> tuple[list[CompressionBatch], list[CompressionBatchEntry]]: + """Greedily group compatible small units and skip under-floor tails. + + Callers retain the skipped entries as normal ``size_floor`` results. The + function deliberately does not turn a unit larger than the configured + batch ceiling into a singleton batch; those units belong to the existing + independent compression path. + """ + + if min_batch_bytes <= 0: + raise ValueError("min_batch_bytes must be positive") + if max_batch_bytes < min_batch_bytes: + raise ValueError("max_batch_bytes must be at least min_batch_bytes") + if max_batch_units <= 0: + raise ValueError("max_batch_units must be positive") + + batches: list[CompressionBatch] = [] + skipped: list[CompressionBatchEntry] = [] + pending: list[CompressionBatchEntry] = [] + pending_bytes = 0 + pending_key: tuple[object, ...] | None = None + + def flush() -> None: + nonlocal pending, pending_bytes, pending_key + if not pending: + return + if pending_bytes >= min_batch_bytes: + batches.append(CompressionBatch(entries=tuple(pending), text_bytes=pending_bytes)) + else: + skipped.extend(pending) + pending = [] + pending_bytes = 0 + pending_key = None + + for entry in entries: + entry_bytes = _text_bytes(entry.routed.unit.text) + entry_key = _compatibility_key(entry) + if entry_bytes >= min_batch_bytes or entry_bytes > max_batch_bytes: + flush() + skipped.append(entry) + continue + if pending and ( + entry_key != pending_key + or len(pending) >= max_batch_units + or pending_bytes + entry_bytes > max_batch_bytes + ): + flush() + pending.append(entry) + pending_bytes += entry_bytes + pending_key = entry_key + if len(pending) == max_batch_units or pending_bytes == max_batch_bytes: + flush() + + flush() + return batches, skipped + + +def _batch_nonce(batch: CompressionBatch) -> str: + digest = hashlib.sha256() + for entry in batch.entries: + digest.update(entry.entry_id.encode("utf-8", errors="replace")) + digest.update(b"\0") + digest.update(entry.routed.unit.text.encode("utf-8", errors="replace")) + digest.update(b"\0") + return digest.hexdigest()[:12] + + +def _batch_envelope(batch: CompressionBatch, nonce: str, texts: list[str]) -> str: + return "\n".join( + ( + f"" + f"{text}" + f"" + ) + for entry, text in zip(batch.entries, texts, strict=True) + ) + + +def _protect_ccr_markers( + batch: CompressionBatch, nonce: str +) -> tuple[list[str], dict[str, tuple[int, str]]]: + """Replace retrieval markers with unique tokens before the single router call.""" + + protected_texts: list[str] = [] + marker_blocks: dict[str, tuple[int, str]] = {} + for entry_index, entry in enumerate(batch.entries): + marker_index = 0 + + def replace_marker(match: re.Match[str], entry_index: int = entry_index) -> str: + nonlocal marker_index + placeholder = f"[[HEADROOM_BATCH_CCR_{nonce}_{entry_index}_{marker_index}]]" + marker_index += 1 + marker_blocks[placeholder] = (entry_index, match.group(0)) + return placeholder + + protected_texts.append(_CCR_MARKER_RE.sub(replace_marker, entry.routed.unit.text)) + return protected_texts, marker_blocks + + +def _parse_batch_envelope( + text: str, + batch: CompressionBatch, + nonce: str, +) -> list[str] | None: + """Return ordered entry bodies only when every expected tag is intact.""" + + cursor = 0 + values: list[str] = [] + for entry in batch.entries: + while cursor < len(text) and text[cursor].isspace(): + cursor += 1 + tag_name = f"headroom-batch-{nonce}-{entry.entry_id}" + pattern = re.compile( + rf"<{re.escape(tag_name)}>(.*?)", + flags=re.DOTALL, + ) + match = pattern.match(text, cursor) + if match is None: + return None + values.append(match.group(1)) + cursor = match.end() + if text[cursor:].strip(): + return None + return values + + +def _passthrough_batch_results( + batch: CompressionBatch, + *, + tokenizer: TokenCounterLike, + reason: str, + router_result: RouterCompressionResult | None = None, +) -> list[tuple[object, UnitCompressionResult]]: + strategy = ( + router_result.strategy_used.value + if router_result + else CompressionStrategy.PASSTHROUGH.value + ) + return [ + ( + entry.routed.slot, + UnitCompressionResult( + original=entry.routed.unit.text, + compressed=entry.routed.unit.text, + modified=False, + tokens_before=tokenizer.count_text(entry.routed.unit.text), + tokens_after=tokenizer.count_text(entry.routed.unit.text), + tokens_saved=0, + transforms_applied=[], + strategy=strategy, + reason=reason, + router_result=router_result, + text_bytes=_text_bytes(entry.routed.unit.text), + min_bytes=entry.routed.unit.min_bytes, + reason_category=reason, + ), + ) + for entry in batch.entries + ] + + +def compress_batch_with_router( + batch: CompressionBatch, + *, + router: ContentRouter, + tokenizer: TokenCounterLike, + target_ratio: float | None = None, +) -> list[tuple[object, UnitCompressionResult]]: + """Compress one tagged batch and split only structurally valid output.""" + + nonce = _batch_nonce(batch) + batch_texts, marker_blocks = _protect_ccr_markers(batch, nonce) + envelope = _batch_envelope(batch, nonce, batch_texts) + protected, protected_blocks = protect_tags(envelope, compress_tagged_content=True) + prior_target_ratio = getattr(router, "_runtime_target_ratio", None) + if target_ratio is not None: + router._runtime_target_ratio = target_ratio + try: + router_result = router.compress( + protected, + context=batch.entries[0].routed.unit.context, + question=batch.entries[0].routed.unit.question, + bias=batch.entries[0].routed.unit.bias, + ) + except Exception: + return _passthrough_batch_results( + batch, + tokenizer=tokenizer, + reason="batch_router_error", + ) + finally: + if target_ratio is not None: + router._runtime_target_ratio = prior_target_ratio + + compressed = router_result.compressed + if not compressed or compressed == protected: + return _passthrough_batch_results( + batch, + tokenizer=tokenizer, + reason="router_no_change", + router_result=router_result, + ) + protected_placeholders = [placeholder for placeholder, _ in protected_blocks] + protected_placeholders.extend(marker_blocks) + if any(compressed.count(placeholder) != 1 for placeholder in protected_placeholders): + return _passthrough_batch_results( + batch, + tokenizer=tokenizer, + reason="batch_invalid", + router_result=router_result, + ) + + restored = restore_tags(compressed, protected_blocks) + replacements = _parse_batch_envelope(restored, batch, nonce) + if replacements is None: + return _passthrough_batch_results( + batch, + tokenizer=tokenizer, + reason="batch_invalid", + router_result=router_result, + ) + if any( + replacements[entry_index].count(placeholder) != 1 + for placeholder, (entry_index, _marker) in marker_blocks.items() + ): + return _passthrough_batch_results( + batch, + tokenizer=tokenizer, + reason="batch_invalid", + router_result=router_result, + ) + for placeholder, (entry_index, marker) in marker_blocks.items(): + replacements[entry_index] = replacements[entry_index].replace(placeholder, marker) + + results: list[tuple[object, UnitCompressionResult]] = [] + strategy = router_result.strategy_used.value + for entry, replacement in zip(batch.entries, replacements, strict=True): + unit = entry.routed.unit + tokens_before = tokenizer.count_text(unit.text) + tokens_after = tokenizer.count_text(replacement) + if ( + unit.role == "tool" + and unit.item_type == "local_shell_call_output" + and _is_structured_shell_output(unit.text) + and strategy in _LOSSY_UNMARKED_STRATEGIES + and not _CCR_MARKER_RE.search(replacement) + ): + result = UnitCompressionResult( + original=unit.text, + compressed=replacement, + modified=False, + tokens_before=tokens_before, + tokens_after=tokens_after, + tokens_saved=0, + transforms_applied=[], + strategy=strategy, + reason="lossy_unrecoverable_tool_output", + router_result=router_result, + text_bytes=_text_bytes(unit.text), + min_bytes=unit.min_bytes, + reason_category="other", + ) + elif tokens_after >= tokens_before: + result = UnitCompressionResult( + original=unit.text, + compressed=replacement, + modified=False, + tokens_before=tokens_before, + tokens_after=tokens_after, + tokens_saved=0, + transforms_applied=[], + strategy=strategy, + reason="rejected_not_smaller", + router_result=router_result, + text_bytes=_text_bytes(unit.text), + min_bytes=unit.min_bytes, + reason_category="rejected_not_smaller", + ) + else: + result = UnitCompressionResult( + original=unit.text, + compressed=replacement, + modified=True, + tokens_before=tokens_before, + tokens_after=tokens_after, + tokens_saved=tokens_before - tokens_after, + transforms_applied=[ + f"router:{unit.provider}:{unit.endpoint}:{unit.item_type}:{strategy}", + strategy, + ], + strategy=strategy, + router_result=router_result, + text_bytes=_text_bytes(unit.text), + min_bytes=unit.min_bytes, + reason_category="applied", + ) + results.append((entry.routed.slot, result)) + return results diff --git a/headroom/transforms/compression_units.py b/headroom/transforms/compression_units.py index 0cfd6757e..c499446c1 100644 --- a/headroom/transforms/compression_units.py +++ b/headroom/transforms/compression_units.py @@ -189,7 +189,7 @@ def _compress_marker_free_text( return text, [], last_router_result leading, core, trailing = boundary.groups() - if len(core) < unit.min_bytes: + if len(core.encode("utf-8", errors="replace")) < unit.min_bytes: return text, [], last_router_result router_result = router.compress( @@ -261,7 +261,7 @@ def compress_unit_with_router( return _with_reason(reason="protected_assistant_message") if unit.cache_zone != "live": return _with_reason(reason=f"cache_zone_{unit.cache_zone}") - if len(unit.text) < unit.min_bytes: + if text_bytes < unit.min_bytes: return _with_reason(reason="below_unit_floor") prior_target_ratio = getattr(router, "_runtime_target_ratio", None) diff --git a/tests/test_compression_batches.py b/tests/test_compression_batches.py new file mode 100644 index 000000000..e6d5a286a --- /dev/null +++ b/tests/test_compression_batches.py @@ -0,0 +1,231 @@ +from __future__ import annotations + +import re + +from headroom.transforms.compression_batches import ( + CompressionBatchEntry, + build_compression_batches, + compress_batch_with_router, +) +from headroom.transforms.compression_units import CompressionUnit, RoutedCompressionUnit +from headroom.transforms.content_router import CompressionStrategy, RouterCompressionResult + + +def _entry(index: int, text: str) -> CompressionBatchEntry: + unit = CompressionUnit( + text=text, + provider="openai", + endpoint="responses", + role="tool", + item_type="local_shell_call_output", + cache_zone="live", + mutable=True, + min_bytes=512, + ) + return CompressionBatchEntry( + entry_id=f"u{index}", + routed=RoutedCompressionUnit(unit=unit, slot=(index, ("output", None))), + ) + + +def test_small_units_over_floor_form_one_batch(): + entries = [_entry(index, "x" * 150) for index in range(4)] + + batches, skipped = build_compression_batches(entries, min_batch_bytes=512) + + assert len(batches) == 1 + assert [entry.entry_id for entry in batches[0].entries] == ["u0", "u1", "u2", "u3"] + assert batches[0].text_bytes == 600 + assert skipped == [] + + +class _CharacterCounter: + def count_text(self, text: str) -> int: + return len(text) + + +class _ShorteningRouter: + def __init__(self) -> None: + self.calls = 0 + + def compress(self, content: str, **_kwargs) -> RouterCompressionResult: + self.calls += 1 + return RouterCompressionResult( + compressed=content.replace("x" * 150, "x"), + original=content, + strategy_used=CompressionStrategy.KOMPRESS, + ) + + +def test_batch_compresses_entries_with_one_router_call(): + entries = [_entry(index, "x" * 150) for index in range(4)] + batches, _ = build_compression_batches(entries, min_batch_bytes=512) + router = _ShorteningRouter() + + results = compress_batch_with_router( + batches[0], + router=router, + tokenizer=_CharacterCounter(), + ) + + assert router.calls == 1 + assert [slot for slot, _ in results] == [entry.routed.slot for entry in entries] + assert [result.compressed for _, result in results] == ["x"] * 4 + assert all(result.modified for _, result in results) + + +def test_under_floor_tail_is_skipped_without_a_batch(): + entries = [_entry(index, "x" * 150) for index in range(3)] + + batches, skipped = build_compression_batches(entries, min_batch_bytes=512) + + assert batches == [] + assert [entry.entry_id for entry in skipped] == ["u0", "u1", "u2"] + + +def test_sixteen_under_floor_entries_are_skipped_before_a_new_batch_starts(): + entries = [_entry(index, "x" * 30) for index in range(17)] + + batches, skipped = build_compression_batches(entries, min_batch_bytes=512) + + assert batches == [] + assert [entry.entry_id for entry in skipped] == [f"u{index}" for index in range(17)] + + +class _CorruptingRouter: + def compress(self, content: str, **_kwargs) -> RouterCompressionResult: + return RouterCompressionResult( + compressed="missing protected tags", + original=content, + strategy_used=CompressionStrategy.KOMPRESS, + ) + + +def test_missing_batch_tags_passes_through_every_entry(): + entries = [_entry(index, "x" * 150) for index in range(4)] + batches, _ = build_compression_batches(entries, min_batch_bytes=512) + + results = compress_batch_with_router( + batches[0], + router=_CorruptingRouter(), + tokenizer=_CharacterCounter(), + ) + + assert [result.compressed for _, result in results] == ["x" * 150] * 4 + assert [result.reason for _, result in results] == ["batch_invalid"] * 4 + assert not any(result.modified for _, result in results) + + +def test_many_small_entries_create_no_more_than_sixteen_per_batch(): + entries = [_entry(index, "x" * 100) for index in range(381)] + + batches, skipped = build_compression_batches(entries, min_batch_bytes=512) + + assert skipped == [] + assert len(batches) <= 24 + assert all(len(batch.entries) <= 16 for batch in batches) + assert all(batch.text_bytes <= 2048 for batch in batches) + + +class _LossyShellRouter: + def compress(self, content: str, **_kwargs) -> RouterCompressionResult: + return RouterCompressionResult( + compressed=content.replace("line alpha beta gamma\n" * 7, "summary"), + original=content, + strategy_used=CompressionStrategy.KOMPRESS, + ) + + +def test_batch_keeps_structured_shell_output_without_a_ccr_marker(): + entries = [_entry(index, "line alpha beta gamma\n" * 7) for index in range(4)] + batches, _ = build_compression_batches(entries, min_batch_bytes=512) + + results = compress_batch_with_router( + batches[0], + router=_LossyShellRouter(), + tokenizer=_CharacterCounter(), + ) + + assert not any(result.modified for _, result in results) + assert [result.reason for _, result in results] == ["lossy_unrecoverable_tool_output"] * 4 + + +class _MarkerStrippingRouter: + def compress(self, content: str, **_kwargs) -> RouterCompressionResult: + return RouterCompressionResult( + compressed=content.replace("word " * 30, "x").replace( + "[100 items compressed to 10. Retrieve more: hash=abc123]", "" + ), + original=content, + strategy_used=CompressionStrategy.KOMPRESS, + ) + + +def test_batch_preserves_ccr_markers_when_router_would_remove_them(): + marker = "[100 items compressed to 10. Retrieve more: hash=abc123]" + original = f"{'word ' * 30}\n{marker}\n" + entries = [_entry(index, original) for index in range(4)] + batches, _ = build_compression_batches(entries, min_batch_bytes=512) + + results = compress_batch_with_router( + batches[0], + router=_MarkerStrippingRouter(), + tokenizer=_CharacterCounter(), + ) + + assert all(result.modified for _, result in results) + assert all(marker in result.compressed for _, result in results) + + +class _MarkerMovingRouter: + def compress(self, content: str, **_kwargs) -> RouterCompressionResult: + placeholders = re.findall(r"\[\[HEADROOM_BATCH_CCR_[^]]+\]\]", content) + moved = content.replace(placeholders[0], "", 1).replace( + placeholders[1], f"{placeholders[1]}{placeholders[0]}", 1 + ) + return RouterCompressionResult( + compressed=moved.replace("word " * 30, "x"), + original=content, + strategy_used=CompressionStrategy.KOMPRESS, + ) + + +def test_batch_rejects_ccr_marker_moved_to_another_entry(): + marker = "[100 items compressed to 10. Retrieve more: hash=abc123]" + original = f"{'word ' * 30}\n{marker}\n" + entries = [_entry(index, original) for index in range(4)] + batches, _ = build_compression_batches(entries, min_batch_bytes=512) + + results = compress_batch_with_router( + batches[0], + router=_MarkerMovingRouter(), + tokenizer=_CharacterCounter(), + ) + + assert [result.compressed for _, result in results] == [original] * 4 + assert [result.reason for _, result in results] == ["batch_invalid"] * 4 + + +class _CjkShorteningRouter: + def compress(self, content: str, **_kwargs) -> RouterCompressionResult: + return RouterCompressionResult( + compressed=content.replace("你" * 150, "短"), + original=content, + strategy_used=CompressionStrategy.KOMPRESS, + ) + + +def test_batch_uses_utf8_bytes_for_cjk_small_units(): + entries = [_entry(index, "你" * 150) for index in range(4)] + + batches, skipped = build_compression_batches(entries, min_batch_bytes=512) + results = compress_batch_with_router( + batches[0], + router=_CjkShorteningRouter(), + tokenizer=_CharacterCounter(), + ) + + assert skipped == [] + assert batches[0].text_bytes == 1800 + assert all(result.modified for _, result in results) + assert [result.compressed for _, result in results] == ["短"] * 4 diff --git a/tests/test_compression_units.py b/tests/test_compression_units.py index 6081a447e..a8e2eac50 100644 --- a/tests/test_compression_units.py +++ b/tests/test_compression_units.py @@ -29,6 +29,29 @@ class Router: ) +class CharacterCounter: + def count_text(self, text: str) -> int: + return len(text) + + +def test_compression_unit_uses_utf8_bytes_for_floor(): + result = compress_unit_with_router( + CompressionUnit( + text="你" * 256, + provider="openai", + endpoint="responses", + role="tool", + item_type="function_call_output", + min_bytes=512, + ), + router=Router("短"), + tokenizer=CharacterCounter(), + ) + + assert result.modified is True + assert result.reason is None + + def test_compression_unit_accepts_token_shrinking_replacement(): result = compress_unit_with_router( CompressionUnit( diff --git a/tests/test_openai_responses_compression_units.py b/tests/test_openai_responses_compression_units.py index c2ae3f294..79761b729 100644 --- a/tests/test_openai_responses_compression_units.py +++ b/tests/test_openai_responses_compression_units.py @@ -176,7 +176,56 @@ def test_openai_responses_adapter_compresses_custom_tool_call_output(): assert strategy_chain == [] -def test_openai_responses_adapter_compresses_output_content_parts(): +def test_openai_responses_adapter_compresses_array_input_text_output(): + router = ContentRouter() + + def compress(self, content: str, **_kwargs): + return RouterCompressionResult( + compressed="custom output summary", + original=content, + strategy_used=CompressionStrategy.KOMPRESS, + ) + + router.compress = MethodType(compress, router) + handler = _handler_with_router(router) + metadata = "Chunk ID: abc\nWall time: 1s" + long_text = " ".join(f"word{i}" for i in range(180)) + image_part = {"type": "input_image", "image_url": "data:image/png;base64,AA=="} + payload = { + "model": "gpt-5", + "input": [ + { + "type": "custom_tool_call_output", + "call_id": "c1", + "output": [ + {"type": "input_text", "text": metadata}, + {"type": "input_text", "text": long_text}, + image_part, + ], + } + ], + } + + new_payload, modified, saved, transforms, units_by_category, strategy_chain, _attempted = ( + handler._compress_openai_responses_live_text_units_with_router( + payload, + model="gpt-5", + request_id="req_test", + ) + ) + + assert modified is True + assert saved > 0 + output = new_payload["input"][0]["output"] + assert output[0]["text"] == metadata + assert output[1]["text"] == "custom output summary" + assert output[2] == image_part + assert "router:openai:responses:custom_tool_call_output:kompress" in transforms + assert units_by_category == {"size_floor": 1, "applied": 1} + assert strategy_chain == [] + + +def test_openai_responses_adapter_compresses_output_text_content_parts(): router = ContentRouter() def compress(self, content: str, **_kwargs): @@ -210,12 +259,154 @@ def test_openai_responses_adapter_compresses_output_content_parts(): assert modified is True assert saved > 0 - assert new_payload["input"][0]["output"] == "content part output summary" + assert new_payload["input"][0]["output"] == [ + {"type": "output_text", "text": "content part output summary"} + ] assert "router:openai:responses:function_call_output:kompress" in transforms assert units_by_category == {"applied": 1} assert strategy_chain == [] +def test_openai_responses_adapter_batches_small_outputs_once(): + router = ContentRouter() + calls: list[str] = [] + floor = OpenAIHandlerMixin.OPENAI_RESPONSES_ROUTER_MIN_BYTES + outputs = [" ".join(f"unit{index}_{token}" for token in range(30)) for index in range(4)] + assert all(len(output.encode("utf-8")) < floor for output in outputs) + assert sum(len(output.encode("utf-8")) for output in outputs) >= floor + + def compress(self, content: str, **_kwargs): + calls.append(content) + compressed = content + for output in outputs: + compressed = compressed.replace(output, "x") + return RouterCompressionResult( + compressed=compressed, + original=content, + strategy_used=CompressionStrategy.KOMPRESS, + ) + + router.compress = MethodType(compress, router) + handler = _handler_with_router(router) + payload = { + "model": "gpt-5", + "input": [ + { + "type": "local_shell_call_output", + "call_id": f"c{index}", + "output": output, + } + for index, output in enumerate(outputs) + ], + } + + new_payload, modified, saved, _, units_by_category, _, attempted = ( + handler._compress_openai_responses_live_text_units_with_router( + payload, + model="gpt-5", + request_id="req_small_batch", + ) + ) + + assert len(calls) == 1 + assert all(output in calls[0] for output in outputs) + assert modified is True + assert saved > 0 + assert attempted == 120 + assert units_by_category == {"applied": 4} + assert [item["output"] for item in new_payload["input"]] == ["x"] * 4 + + +def test_openai_responses_adapter_batches_small_array_parts_without_touching_images(): + router = ContentRouter() + calls = {"count": 0} + + def compress(self, content: str, **_kwargs): + calls["count"] += 1 + return RouterCompressionResult( + compressed=content.replace("word " * 30, "x"), + original=content, + strategy_used=CompressionStrategy.KOMPRESS, + ) + + router.compress = MethodType(compress, router) + handler = _handler_with_router(router) + image_part = {"type": "input_image", "image_url": "data:image/png;base64,AA=="} + payload = { + "model": "gpt-5", + "input": [ + { + "type": "custom_tool_call_output", + "call_id": "c1", + "output": [ + {"type": "input_text", "text": "word " * 30}, + image_part, + {"type": "input_text", "text": "word " * 30}, + {"type": "input_text", "text": "word " * 30}, + {"type": "input_text", "text": "word " * 30}, + ], + } + ], + } + + new_payload, modified, saved, *_ = ( + handler._compress_openai_responses_live_text_units_with_router( + payload, + model="gpt-5", + request_id="req_small_array_batch", + ) + ) + + assert calls["count"] == 1 + assert modified is True + assert saved > 0 + output = new_payload["input"][0]["output"] + assert [output[index]["text"] for index in (0, 2, 3, 4)] == ["x"] * 4 + assert output[1] == image_part + + +def test_openai_responses_adapter_skips_under_floor_small_batch(): + router = ContentRouter() + calls = {"count": 0} + + def compress(self, content: str, **_kwargs): + calls["count"] += 1 + return RouterCompressionResult( + compressed="x", + original=content, + strategy_used=CompressionStrategy.KOMPRESS, + ) + + router.compress = MethodType(compress, router) + handler = _handler_with_router(router) + payload = { + "model": "gpt-5", + "input": [ + { + "type": "function_call_output", + "call_id": f"c{index}", + "output": "word " * 30, + } + for index in range(3) + ], + } + + new_payload, modified, saved, _, units_by_category, _, attempted = ( + handler._compress_openai_responses_live_text_units_with_router( + payload, + model="gpt-5", + request_id="req_under_floor_batch", + ) + ) + + assert calls["count"] == 0 + assert new_payload == payload + assert modified is False + assert saved == 0 + assert attempted == 0 + assert units_by_category == {"size_floor": 3} + + def test_openai_responses_adapter_reuses_exact_tool_output_cache(): router = ContentRouter() calls = {"count": 0} @@ -792,7 +983,7 @@ def test_openai_responses_payload_routes_through_content_router_without_rust( assert any(t.startswith("router:openai:responses:") for t in transforms) -def test_openai_responses_adapter_aggregates_small_tool_outputs_before_floor(): +def test_openai_responses_adapter_batches_small_tool_outputs_before_floor(): """Regression for #2050: many individually-small tool outputs whose combined size clears the floor must still reach the router. @@ -803,10 +994,15 @@ def test_openai_responses_adapter_aggregates_small_tool_outputs_before_floor(): the aggregate of the extracted group, matching the batch (Anthropic) path. """ router = ContentRouter() + calls: list[str] = [] def compress(self, content: str, **_kwargs): + calls.append(content) + compressed = content + for output in outputs: + compressed = compressed.replace(output, "tiny summary") return RouterCompressionResult( - compressed="tiny summary", + compressed=compressed, original=content, strategy_used=CompressionStrategy.KOMPRESS, ) @@ -843,6 +1039,8 @@ def test_openai_responses_adapter_aggregates_small_tool_outputs_before_floor(): ) ) + assert len(calls) == 1 + assert all(output in calls[0] for output in outputs) assert modified is True assert saved > 0 # No unit should be size-floored; every extracted unit is compressed.