diff --git a/studio/backend/core/inference/context_compaction.py b/studio/backend/core/inference/context_compaction.py new file mode 100644 index 0000000000..3920ee3e3e --- /dev/null +++ b/studio/backend/core/inference/context_compaction.py @@ -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"]) diff --git a/studio/backend/tests/test_context_compaction.py b/studio/backend/tests/test_context_compaction.py new file mode 100644 index 0000000000..f95c0499fe --- /dev/null +++ b/studio/backend/tests/test_context_compaction.py @@ -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)