Studio: handle keep_recent=0 + anchored pair cleanup
Two related corner cases on the compaction module, both regression- tested: 1) keep_recent=0 collected ONE group instead of zero. The recent- window loop tested its limit AFTER appending, so the first iteration always added one group and broke. Flip the check to BEFORE appending so a caller asking for the minimal-anchor regime actually gets only the anchor set (system + first-user + multimodal) and drops everything else. 2) Pair cleanup could drop an ANCHORED message. When all tool-role messages paired with an assistant tool_call message were dropped, the cleanup rule forced the assistant to drop too, even if that assistant was an anchor (multimodal carrying tool_calls, rare but valid for vision-capable tool-calling models). Mirror invariant: if the assistant is anchored leak an orphan tool_call shape rather than violate "multimodal turns are never dropped". Same for the inverse direction: a dropped assistant tool_call message should not drag anchored tool-role partners into the drop set. Tests: 19 -> 21. Add TestKeepRecentZero and TestAnchoredMultimodalPairCleanup classes covering the two fixes.
This commit is contained in:
parent
12270ca9fa
commit
41192f0452
2 changed files with 103 additions and 6 deletions
|
|
@ -195,17 +195,22 @@ class SlidingWindowCompact(CompactStrategy):
|
|||
group_id[ti] = g
|
||||
|
||||
# The "recent window" is the last ``keep_recent`` distinct
|
||||
# groups encountered scanning from the end.
|
||||
# groups encountered scanning from the end. ``keep_recent == 0``
|
||||
# must collect ZERO groups so the caller can drop everything
|
||||
# outside the anchor set (system + first user + multimodal).
|
||||
# The pre-fix loop tested the limit AFTER appending and so
|
||||
# always preserved at least one group even when keep_recent
|
||||
# was 0; flip the check to BEFORE appending so the bound holds.
|
||||
recent_groups: list[int] = []
|
||||
seen_groups: set[int] = set()
|
||||
for i in range(len(messages) - 1, -1, -1):
|
||||
if len(recent_groups) >= self.keep_recent:
|
||||
break
|
||||
g = group_id[i]
|
||||
if g in seen_groups:
|
||||
continue
|
||||
seen_groups.add(g)
|
||||
recent_groups.append(g)
|
||||
if len(recent_groups) >= self.keep_recent:
|
||||
break
|
||||
recent_groups_set = set(recent_groups)
|
||||
|
||||
# Decide drop set: every index whose group is NOT in the recent
|
||||
|
|
@ -230,13 +235,23 @@ class SlidingWindowCompact(CompactStrategy):
|
|||
# When dropping an assistant-with-tool-calls message we must
|
||||
# also drop the matching tool-role messages (and vice versa)
|
||||
# so the chat template stays valid. Iterate pair_map once.
|
||||
# Anchor indices stay regardless: dragging an anchored
|
||||
# multimodal assistant or first-user message into the drop
|
||||
# set just because its tool-pair partner was dropped would
|
||||
# violate the structural invariant the anchor set exists to
|
||||
# enforce, and llama-server would 400 on the resulting
|
||||
# template (a tool message whose tool_call_id has no
|
||||
# surviving assistant tool_calls entry).
|
||||
for asst_idx, tool_idxs in pair_map.items():
|
||||
if asst_idx in dropped:
|
||||
dropped.update(tool_idxs)
|
||||
dropped.update(t for t in tool_idxs if t not in anchor_idx)
|
||||
elif tool_idxs and tool_idxs <= dropped:
|
||||
# All matching tool messages were dropped: drop the
|
||||
# assistant tool-call message too.
|
||||
dropped.add(asst_idx)
|
||||
# assistant tool-call message too -- unless it's
|
||||
# anchored, in which case we'd rather leak an
|
||||
# orphan tool_call shape than violate the invariant.
|
||||
if asst_idx not in anchor_idx:
|
||||
dropped.add(asst_idx)
|
||||
|
||||
return [m for i, m in enumerate(messages) if i not in dropped]
|
||||
|
||||
|
|
|
|||
|
|
@ -277,6 +277,88 @@ class TestStrategyRegistry:
|
|||
assert isinstance(get_strategy("totally-made-up"), NoCompact)
|
||||
|
||||
|
||||
class TestKeepRecentZero:
|
||||
"""Regression: keep_recent=0 must collect ZERO recent groups, not 1.
|
||||
|
||||
Pre-fix, the loop tested the limit AFTER appending, so the first
|
||||
iteration always added one group and broke. Trim with keep_recent=0
|
||||
should drop everything outside the anchor set (system + first-user
|
||||
+ multimodal). This matters when a caller intentionally wants only
|
||||
the anchors to survive (extreme-pressure regime).
|
||||
"""
|
||||
|
||||
def test_keep_recent_zero_keeps_only_anchors(self):
|
||||
msgs = [
|
||||
_msg("system", "sys"),
|
||||
_msg("user", "task"),
|
||||
_long("assistant", 4000),
|
||||
_long("user", 4000),
|
||||
_long("assistant", 4000),
|
||||
_long("user", 4000),
|
||||
]
|
||||
out = SlidingWindowCompact(keep_recent = 0).compact(
|
||||
msgs, budget_tokens = 50
|
||||
)
|
||||
roles = [m["role"] for m in out]
|
||||
# System + first-user only.
|
||||
assert roles == ["system", "user"]
|
||||
assert out[1]["content"].startswith("task")
|
||||
|
||||
|
||||
class TestAnchoredMultimodalPairCleanup:
|
||||
"""Regression: an anchored multimodal assistant whose paired tool
|
||||
messages all get dropped must survive pair-cleanup. Without this
|
||||
guard the pair_map's "all tool messages dropped -> drop the
|
||||
assistant too" rule would yank a multimodal anchor out from under
|
||||
the structural invariant that says multimodal turns are never
|
||||
dropped, and llama-server would 400 the resulting template.
|
||||
"""
|
||||
|
||||
def test_anchored_multimodal_assistant_survives_pair_cleanup(self):
|
||||
# Multimodal assistant carrying tool_calls (rare but valid:
|
||||
# vision models can call tools while emitting image parts).
|
||||
multimodal_asst = {
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "describe and call tool"},
|
||||
{"type": "image_url", "image_url": {"url": "x"}},
|
||||
],
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_mm",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "web_search",
|
||||
"arguments": '{"q":"y"}',
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
msgs = [
|
||||
_msg("system", "sys"),
|
||||
_msg("user", "task"),
|
||||
multimodal_asst,
|
||||
_msg(
|
||||
"tool",
|
||||
content = "result for y" * 200,
|
||||
tool_call_id = "call_mm",
|
||||
name = "web_search",
|
||||
),
|
||||
_long("assistant", 4000),
|
||||
_long("user", 4000),
|
||||
_long("assistant", 4000),
|
||||
_long("user", 4000),
|
||||
]
|
||||
out = SlidingWindowCompact(keep_recent = 1).compact(
|
||||
msgs, budget_tokens = 50
|
||||
)
|
||||
# Anchored multimodal assistant must still be there.
|
||||
assert any(
|
||||
m.get("role") == "assistant" and isinstance(m.get("content"), list)
|
||||
for m in out
|
||||
), "multimodal assistant got dropped by pair_map cleanup"
|
||||
|
||||
|
||||
class TestConstructorValidation:
|
||||
def test_negative_keep_recent_raises(self):
|
||||
import pytest
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue