ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: sliding-window context compaction module

Browse source 
Adds a standalone context-compaction module so long Studio Chat
sessions can trim older turns from the prompt sent to the model
without losing them from the persisted transcript. This PR ships the
module and tests only; the wire-up into the agentic and chat loops
lands in a follow-up so behavior on main is unchanged.

Two strategies:

* NoCompact: passthrough, returns a copy.
* SlidingWindowCompact: keeps the system message, the first user
  message, and the last N "groups" (an assistant tool-call message
  plus its matching tool-role responses count as one group).

Invariants preserved by every strategy:

* System message and first user message are never dropped.
* Tool-call <-> tool-result pair linkage stays valid: an assistant
  message that carries tool_calls is kept or dropped together with
  every tool-role message matching its tool_call_ids.
* Multimodal turns (list-typed content) are never dropped (no tested
  compacted-media representation today).
* The input list is never mutated; the strategy returns a new list.

Token estimation uses a 4-char-per-token char-count heuristic;
tokenizer-aware estimates are a follow-up. Adapted from forge
(https://github.com/antoinezambelli/forge, MIT).
This commit is contained in:
Daniel Han 2026-05-22 14:46:27 +00:00
commit 91b2fa57a8
2 changed files with 523 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

262
studio/backend/core/inference/context_compaction.py Normal file
View file

@ -0,0 +1,262 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
#
# SlidingWindowCompact is adapted from forge
# (https://github.com/antoinezambelli/forge), Copyright (c) 2025-2026
# Antoine Zambelli, used under the MIT License.
"""Context-window compaction for OpenAI-style chat messages.
Long-running Studio Chat sessions can outgrow a model's context window.
This module ships a strategy that trims older messages from the prompt
sent to the model while preserving the persisted transcript shown in
the UI. The strategy returns a NEW list; ``messages`` is never mutated.
Invariants preserved by all strategies:
1. The system message (when present at index 0) is never dropped.
2. The first user message (when present) is never dropped: it carries
the task prompt the rest of the conversation references.
3. Tool-call <-> tool-result pair linkage stays valid. An assistant
message that carries ``tool_calls`` and the matching tool-role
messages are kept or dropped as a unit. Dropping one side leaves
the OpenAI chat template invalid and llama-server returns 400.
4. Multimodal turns (any message whose ``content`` is a list of parts)
are treated as non-droppable. There is no tested compacted-media
representation today; the strategies leave such turns intact.
Only one strategy is shipped here. ``TieredCompact`` from forge depends
on per-message metadata tags Studio's flat message dicts do not carry;
it will land in a follow-up once the message model gets the tags.
"""
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import Any
# Char-to-token heuristic. Conservative on the high side so the
# compactor triggers earlier rather than later. Tokenizer-aware
# estimates may land in a follow-up.
_CHARS_PER_TOKEN = 4
def estimate_tokens(messages: list[dict]) -> int:
"""Rough token count for ``messages``. Uses a 4-char-per-token
char-count heuristic on the visible ``content`` (str) and on
serialized ``tool_calls`` arguments. Multimodal parts (list-typed
content) contribute only their text parts.
"""
total_chars = 0
for m in messages:
content = m.get("content")
if isinstance(content, str):
total_chars += len(content)
elif isinstance(content, list):
for part in content:
if isinstance(part, dict):
t = part.get("text")
if isinstance(t, str):
total_chars += len(t)
tcs = m.get("tool_calls")
if isinstance(tcs, list):
for tc in tcs:
args = (tc.get("function") or {}).get("arguments")
if isinstance(args, str):
total_chars += len(args)
return total_chars // _CHARS_PER_TOKEN
def _is_multimodal(msg: dict) -> bool:
return isinstance(msg.get("content"), list)
def _is_tool_message(msg: dict) -> bool:
return msg.get("role") == "tool"
def _assistant_tool_call_ids(msg: dict) -> set[str]:
"""Return the set of ``id`` values from an assistant message's
``tool_calls``. Empty set when the message has no tool calls.
"""
out: set[str] = set()
tcs = msg.get("tool_calls")
if isinstance(tcs, list):
for tc in tcs:
tcid = tc.get("id")
if isinstance(tcid, str) and tcid:
out.add(tcid)
return out
def _pair_linked_indices(messages: list[dict]) -> dict[int, set[int]]:
"""Map an assistant-message index to the indices of its tool-role
follow-ups (matching ``tool_call_id``). Used so the compactor drops
or keeps an assistant+tool group as a unit.
"""
out: dict[int, set[int]] = {}
# Walk in order so the next tool-role messages after an assistant
# call are the natural matches. A tool message is matched to the
# most recent prior assistant whose ``tool_calls`` contain that id.
pending_ids: dict[str, int] = {}
for i, m in enumerate(messages):
if m.get("role") == "assistant":
ids = _assistant_tool_call_ids(m)
out.setdefault(i, set())
for tid in ids:
pending_ids[tid] = i
elif _is_tool_message(m):
tcid = m.get("tool_call_id")
if isinstance(tcid, str) and tcid in pending_ids:
out.setdefault(pending_ids[tcid], set()).add(i)
return out
class CompactStrategy(ABC):
"""Interface for context-compaction strategies."""
@abstractmethod
def compact(
self, messages: list[dict], budget_tokens: int
) -> list[dict]:
"""Return a (possibly shorter) list of messages within
``budget_tokens``. Returns ``messages`` unchanged when no
compaction is needed or possible.
"""
...
class NoCompact(CompactStrategy):
"""Passthrough strategy. Returns ``messages`` unchanged."""
def compact(
self, messages: list[dict], budget_tokens: int
) -> list[dict]:
return list(messages)
class SlidingWindowCompact(CompactStrategy):
"""Keep the system message, the first user message, and the last
``keep_recent`` non-droppable turns. Multimodal turns are never
dropped (no compacted-media representation today). Assistant
messages with ``tool_calls`` are grouped with their matching
tool-role responses and treated as one unit.
The strategy is a no-op when the estimated token count is already
within ``budget_tokens`` or when there is nothing left to drop.
"""
def __init__(
self, keep_recent: int = 2, compact_threshold: float = 0.85
) -> None:
if keep_recent < 0:
raise ValueError("keep_recent must be >= 0")
if not (0.0 < compact_threshold <= 1.0):
raise ValueError("compact_threshold must be in (0, 1]")
self.keep_recent = keep_recent
self.compact_threshold = compact_threshold
def compact(
self, messages: list[dict], budget_tokens: int
) -> list[dict]:
if budget_tokens <= 0 or not messages:
return list(messages)
threshold = int(budget_tokens * self.compact_threshold)
if estimate_tokens(messages) <= threshold:
return list(messages)
# Anchor indices that may never be dropped.
anchor_idx: set[int] = set()
# System message at index 0.
if messages and messages[0].get("role") == "system":
anchor_idx.add(0)
# First user message (the task prompt).
for i, m in enumerate(messages):
if m.get("role") == "user":
anchor_idx.add(i)
break
# Multimodal turns.
for i, m in enumerate(messages):
if _is_multimodal(m):
anchor_idx.add(i)
# Tool-call / tool-result grouping. An assistant tool-call
# message and its matching tool-role responses get the same
# group id so we keep or drop them as a unit.
pair_map = _pair_linked_indices(messages)
group_id: list[int] = list(range(len(messages)))
next_g = len(messages)
for asst_idx, tool_idxs in pair_map.items():
if not tool_idxs:
continue
g = next_g
next_g += 1
group_id[asst_idx] = g
for ti in tool_idxs:
group_id[ti] = g
# The "recent window" is the last ``keep_recent`` distinct
# groups encountered scanning from the end.
recent_groups: list[int] = []
seen_groups: set[int] = set()
for i in range(len(messages) - 1, -1, -1):
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
# window AND that is not an anchor.
droppable: list[int] = []
for i in range(len(messages)):
if i in anchor_idx:
continue
if group_id[i] in recent_groups_set:
continue
droppable.append(i)
# Drop oldest-first until under threshold or nothing left.
dropped: set[int] = set()
for i in droppable:
kept = [
m for j, m in enumerate(messages)
if j not in dropped and j != i
]
if estimate_tokens(kept) <= threshold:
dropped.add(i)
break
dropped.add(i)
# 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.
for asst_idx, tool_idxs in pair_map.items():
if asst_idx in dropped:
dropped.update(tool_idxs)
elif tool_idxs and tool_idxs <= dropped:
# All matching tool messages were dropped: drop the
# assistant tool-call message too.
dropped.add(asst_idx)
return [m for i, m in enumerate(messages) if i not in dropped]
_STRATEGIES: dict[str, CompactStrategy] = {
"none": NoCompact(),
"sliding": SlidingWindowCompact(),
}
def get_strategy(name: str) -> CompactStrategy:
"""Return the compaction strategy registered under ``name``.
Falls back to ``NoCompact`` for unknown names so a misconfigured
request degrades to no-op rather than raising.
"""
return _STRATEGIES.get(name, _STRATEGIES["none"])

261
studio/backend/tests/test_context_compaction.py Normal file
View file

@ -0,0 +1,261 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for the context-compaction module.
Coverage focuses on the structural invariants the strategies promise:
* System and first-user messages are never dropped.
* Tool-call <-> tool-result pair linkage stays valid.
* Multimodal turns are never dropped.
* Persisted history is not mutated; the strategy returns a new list.
* ``NoCompact`` is a true passthrough.
* ``SlidingWindowCompact`` is a no-op when the budget already fits.
"""
from core.inference.context_compaction import (
NoCompact,
SlidingWindowCompact,
estimate_tokens,
get_strategy,
)
def _msg(role, content = "", tool_calls = None, tool_call_id = None, name = None):
m = {"role": role}
if content is not None:
m["content"] = content
if tool_calls is not None:
m["tool_calls"] = tool_calls
if tool_call_id is not None:
m["tool_call_id"] = tool_call_id
if name is not None:
m["name"] = name
return m
def _long(role, length, *, content_prefix = "x"):
return _msg(role, content_prefix * length)
class TestEstimateTokens:
def test_empty_returns_zero(self):
assert estimate_tokens([]) == 0
def test_simple_string_content(self):
# 4 chars per token: "abcd" -> 1.
msgs = [_msg("user", "abcd")]
assert estimate_tokens(msgs) == 1
def test_multimodal_text_part_counts(self):
msgs = [
{
"role": "user",
"content": [
{"type": "text", "text": "abcdefgh"},
{"type": "image_url", "image_url": {"url": "x"}},
],
}
]
# 8 chars text, image URL ignored.
assert estimate_tokens(msgs) == 2
def test_tool_call_arguments_count(self):
msgs = [
_msg(
"assistant",
content = "",
tool_calls = [{
"id": "c1",
"type": "function",
"function": {"name": "web_search", "arguments": '{"q":"hi"}'},
}],
)
]
# 10 char arguments -> 2 tokens.
assert estimate_tokens(msgs) >= 2
class TestNoCompact:
def test_passthrough_returns_copy(self):
msgs = [_msg("user", "hi")]
out = NoCompact().compact(msgs, budget_tokens = 10)
assert out == msgs
# Mutating the result does not affect the original.
out.append(_msg("assistant", "new"))
assert len(msgs) == 1
class TestSlidingWindowUnderBudget:
def test_no_op_when_already_fits(self):
msgs = [_msg("system", "be concise"), _msg("user", "hi"), _msg("assistant", "bye")]
out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 1000)
assert out == msgs
def test_no_op_when_budget_is_zero(self):
# A zero/negative budget collapses to no-op (defensive).
msgs = [_msg("user", "abcd" * 1000)]
out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 0)
assert out == msgs
class TestSlidingWindowInvariants:
def _make_long_chat(self, n_turns, length_per_turn = 1000):
msgs = [_msg("system", "system prompt")]
msgs.append(_msg("user", "the original task: " + "x" * length_per_turn))
# Alternating assistant/user follow-ups.
for i in range(n_turns):
msgs.append(_long("assistant", length_per_turn))
msgs.append(_long("user", length_per_turn))
return msgs
def test_keeps_system_and_first_user(self):
msgs = self._make_long_chat(n_turns = 20)
out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 200)
assert out[0]["role"] == "system"
# First user message must survive.
assert any(
m.get("role") == "user" and "original task" in m.get("content", "")
for m in out
)
def test_keeps_last_n_turns(self):
msgs = self._make_long_chat(n_turns = 20)
out = SlidingWindowCompact(keep_recent = 3).compact(msgs, budget_tokens = 200)
# The last assistant and user pair must survive.
assert out[-1] == msgs[-1]
assert out[-2] == msgs[-2]
def test_does_not_mutate_input(self):
msgs = self._make_long_chat(n_turns = 10)
snap = [dict(m) for m in msgs]
_ = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 100)
assert msgs == snap
def test_multimodal_turn_never_dropped(self):
msgs = [
_msg("system", "sys"),
_msg("user", "first task"),
_long("assistant", 1000),
_long("user", 1000),
{
"role": "user",
"content": [
{"type": "text", "text": "look at this"},
{"type": "image_url", "image_url": {"url": "x"}},
],
},
_long("assistant", 1000),
_long("user", 1000),
]
out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 200)
# The multimodal turn must still be present.
assert any(isinstance(m.get("content"), list) for m in out)
class TestSlidingWindowToolPairs:
def _make_chat_with_tool_pair(self):
return [
_msg("system", "sys"),
_msg("user", "first task"),
_long("assistant", 800),
_long("user", 800),
# The tool pair we want to test linkage on.
_msg(
"assistant",
content = "",
tool_calls = [{
"id": "call_42",
"type": "function",
"function": {
"name": "web_search",
"arguments": '{"q":"x"}',
},
}],
),
_msg("tool", content = "result for x", tool_call_id = "call_42", name = "web_search"),
_long("assistant", 800),
_long("user", 800),
_long("assistant", 800),
_long("user", 800),
]
def test_tool_pair_dropped_together(self):
msgs = self._make_chat_with_tool_pair()
# Force aggressive compaction.
out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 50)
kept_assistant_with_calls = [
m for m in out
if m.get("role") == "assistant" and isinstance(m.get("tool_calls"), list)
]
kept_tool_msgs = [m for m in out if m.get("role") == "tool"]
# If the assistant tool-call message survives, every matching
# tool-role message must also survive (and vice versa).
kept_ids = set()
for m in kept_assistant_with_calls:
for tc in m.get("tool_calls", []):
if isinstance(tc, dict) and tc.get("id"):
kept_ids.add(tc["id"])
for m in kept_tool_msgs:
assert m.get("tool_call_id") in kept_ids
def test_tool_pair_kept_when_in_recent_window(self):
# Build a chat where the tool pair sits inside the recent
# window so it must survive even when the rest is far over
# budget. Layout: system, first-user, long pair x2, tool pair,
# one final user turn. keep_recent=2 catches the tool pair as
# the second-to-last group and the final user turn as the last.
msgs = [
_msg("system", "sys"),
_msg("user", "first task"),
_long("assistant", 800),
_long("user", 800),
_long("assistant", 800),
_long("user", 800),
_msg(
"assistant",
content = "",
tool_calls = [{
"id": "call_42",
"type": "function",
"function": {
"name": "web_search",
"arguments": '{"q":"x"}',
},
}],
),
_msg("tool", content = "result", tool_call_id = "call_42", name = "web_search"),
_msg("user", "thanks"),
]
out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 50)
ids = [
(m.get("role"), m.get("tool_call_id") or
(m.get("tool_calls") and m["tool_calls"][0].get("id")))
for m in out
]
assert ("assistant", "call_42") in ids
assert ("tool", "call_42") in ids
class TestStrategyRegistry:
def test_get_strategy_known(self):
assert isinstance(get_strategy("none"), NoCompact)
assert isinstance(get_strategy("sliding"), SlidingWindowCompact)
def test_get_strategy_unknown_falls_back_to_none(self):
# Unknown strategy names must degrade to no-op rather than raise.
assert isinstance(get_strategy("totally-made-up"), NoCompact)
class TestConstructorValidation:
def test_negative_keep_recent_raises(self):
import pytest
with pytest.raises(ValueError):
SlidingWindowCompact(keep_recent = -1)
def test_invalid_threshold_raises(self):
import pytest
with pytest.raises(ValueError):
SlidingWindowCompact(compact_threshold = 0.0)
with pytest.raises(ValueError):
SlidingWindowCompact(compact_threshold = 1.5)