mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix(compression): protect part-level cache_control rows in compress() too
compress() scores text-only copies of the rows, so a content-part cache_control marker was gone by the time get_protected_indices ran and the pinned row could still be stubbed. Read protection from the original rows, which are index-aligned with the normalized copies, and add a regression test that fails without the change.
This commit is contained in:
parent
dad6e9be4a
commit
931bdb8c0b
2 changed files with 36 additions and 4 deletions
|
|
@ -435,7 +435,7 @@ def compress(
|
|||
combined_scores = bm25_scores
|
||||
|
||||
# Protected messages are never compressed
|
||||
protected_indices: Final = get_protected_indices(normalized_messages)
|
||||
protected_indices: Final = get_protected_indices(original_messages)
|
||||
kept_indices: set[int] = set(protected_indices)
|
||||
|
||||
tool_exchange_spans: list[set[int]] = []
|
||||
|
|
|
|||
|
|
@ -6,7 +6,8 @@ never rewrite. It is consumed by compress() and by the Headroom guardrail, so
|
|||
the two agree on what "never compress this" means.
|
||||
"""
|
||||
|
||||
from litellm.compression.compress import get_protected_indices
|
||||
from litellm.compression.compress import compress, get_protected_indices
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
|
||||
def test_protects_system_last_user_and_last_assistant():
|
||||
|
|
@ -66,13 +67,12 @@ def test_mid_history_cache_control_part_is_protected():
|
|||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "a large cached tool result"},
|
||||
{"type": "text", "text": "a large cached tool result", "cache_control": {"type": "ephemeral"}},
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": "ack"},
|
||||
{"role": "user", "content": "live instruction"},
|
||||
]
|
||||
messages[2]["content"][0]["cache_control"] = {"type": "ephemeral"}
|
||||
|
||||
# index 3 = last assistant, index 4 = last user (both protected by role
|
||||
# regardless), index 2 = the cache_control-marked row itself.
|
||||
|
|
@ -113,3 +113,35 @@ def test_content_that_is_not_a_list_of_mappings_is_not_treated_as_cache_control(
|
|||
]
|
||||
|
||||
assert sorted(get_protected_indices(messages)) == [0, 2]
|
||||
|
||||
|
||||
def test_compress_keeps_part_level_cache_control_row_verbatim():
|
||||
# compress() scores text-only copies of the rows, where a part-level marker
|
||||
# is gone; protection has to read the original rows or the pinned row is stubbed.
|
||||
stale_log = {"role": "user", "content": [{"type": "text", "text": "stale log line " * 2000}]}
|
||||
pinned = {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "cached tool result " * 2000, "cache_control": {"type": "ephemeral"}},
|
||||
],
|
||||
}
|
||||
messages = [
|
||||
stale_log,
|
||||
{"role": "assistant", "content": "old answer"},
|
||||
pinned,
|
||||
{"role": "assistant", "content": "ack"},
|
||||
{"role": "user", "content": "live instruction"},
|
||||
]
|
||||
|
||||
result = compress(
|
||||
messages,
|
||||
model="gpt-4o",
|
||||
call_type=CallTypes.anthropic_messages,
|
||||
compression_trigger=1000,
|
||||
compression_target=500,
|
||||
)
|
||||
|
||||
assert len(result["messages"]) == len(messages)
|
||||
assert result["messages"][2] == pinned
|
||||
assert result["messages"][0] != stale_log
|
||||
assert len(result["cache"]) >= 1
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue