mirror of
https://github.com/usestrix/strix.git
synced 2026-08-23 11:22:37 +02:00
feat(context): model-aware conversation compaction for long scans
Restore cumulative history compaction (removed in the SDK migration) so a long scan no longer replays an ever-growing transcript until it overflows the model's context window and the run fails. - strix/llm/context_budget.py: resolve the model's real input/output token limits from LiteLLM metadata (128k gpt-4o, 272k gpt-5, 1M claude-sonnet-4, 131k deepseek), with a large configurable fallback for unmapped models and a chars/4 token-count fallback. - strix/llm/compaction.py: provider-agnostic compaction via litellm. Keeps a security-focused structured summary (objective, findings, credentials, payloads, URLs/paths, work state, dead ends, next move), keeps the most recent turns by token budget, and snaps the summary boundary so no tool call is separated from its result. Both triggers: proactive before each run and reactive compact-and-retry on a real context-overflow error. - Wire both triggers into the agent run loop next to the existing image recovery; add replace_session_items() with restore-on-failure. Env-tunable via STRIX_CONTEXT_* (auto-compact on by default).
This commit is contained in:
@@ -0,0 +1,186 @@
|
||||
"""Tests for provider-agnostic conversation compaction."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import pytest
|
||||
|
||||
from strix.config import ContextSettings
|
||||
from strix.llm import compaction
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agents.memory.session_settings import SessionSettings
|
||||
|
||||
|
||||
class FakeSession:
|
||||
"""Minimal in-memory Session for exercising compaction."""
|
||||
|
||||
session_id = "fake"
|
||||
session_settings: SessionSettings | None = None
|
||||
|
||||
def __init__(self, items: list[Any]) -> None:
|
||||
self._items = list(items)
|
||||
|
||||
async def get_items(self, limit: int | None = None) -> list[Any]:
|
||||
return list(self._items) if limit is None else list(self._items[-limit:])
|
||||
|
||||
async def add_items(self, items: list[Any]) -> None:
|
||||
self._items.extend(items)
|
||||
|
||||
async def clear_session(self) -> None:
|
||||
self._items = []
|
||||
|
||||
async def pop_item(self) -> Any:
|
||||
return self._items.pop() if self._items else None
|
||||
|
||||
|
||||
def _user(text: str) -> dict[str, Any]:
|
||||
return {"role": "user", "content": text}
|
||||
|
||||
|
||||
def _assistant(text: str) -> dict[str, Any]:
|
||||
return {"role": "assistant", "content": text}
|
||||
|
||||
|
||||
def _call(call_id: str, name: str = "exec_command") -> dict[str, Any]:
|
||||
return {"type": "function_call", "call_id": call_id, "name": name, "arguments": "{}"}
|
||||
|
||||
|
||||
def _output(call_id: str, text: str = "done") -> dict[str, Any]:
|
||||
return {"type": "function_call_output", "call_id": call_id, "output": text}
|
||||
|
||||
|
||||
def _turns(n: int) -> list[dict[str, Any]]:
|
||||
items: list[dict[str, Any]] = []
|
||||
for i in range(n):
|
||||
items += [
|
||||
_user(f"task {i}"),
|
||||
_call(f"c{i}"),
|
||||
_output(f"c{i}", f"result {i}"),
|
||||
_assistant(f"ok {i}"),
|
||||
]
|
||||
return items
|
||||
|
||||
|
||||
def _has_orphan_tool_output(items: list[Any]) -> bool:
|
||||
call_ids = {i["call_id"] for i in items if compaction._is_tool_call(i)}
|
||||
return any(i["call_id"] not in call_ids for i in items if compaction._is_tool_output(i))
|
||||
|
||||
|
||||
def test_is_context_overflow_matches_and_excludes() -> None:
|
||||
assert compaction.is_context_overflow(
|
||||
RuntimeError("This model's maximum context length is 8192")
|
||||
)
|
||||
assert compaction.is_context_overflow(ValueError("input is too long for the model"))
|
||||
assert not compaction.is_context_overflow(RuntimeError("rate limit exceeded, retry later"))
|
||||
|
||||
|
||||
def test_select_split_never_orphans_tool_output(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(compaction, "count_tokens", lambda _m, t: len(t))
|
||||
items = _turns(10)
|
||||
split = compaction._select_split("m", items, keep_tokens=25)
|
||||
recent = items[split:]
|
||||
assert recent # something is kept
|
||||
assert not _has_orphan_tool_output(recent)
|
||||
|
||||
|
||||
def test_select_split_handles_parallel_calls(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(compaction, "count_tokens", lambda _m, _t: 1)
|
||||
# Two parallel calls then their two outputs.
|
||||
items = [
|
||||
_user("start"),
|
||||
_call("a"),
|
||||
_call("b"),
|
||||
_output("a"),
|
||||
_output("b"),
|
||||
_assistant("done"),
|
||||
]
|
||||
# keep_tokens picks a boundary that would land between the calls/outputs.
|
||||
split = compaction._select_split("m", items, keep_tokens=2)
|
||||
assert not _has_orphan_tool_output(items[split:])
|
||||
|
||||
|
||||
def _patch_budget(monkeypatch: pytest.MonkeyPatch, *, keep_tokens: int, window: int) -> None:
|
||||
monkeypatch.setattr(compaction, "count_tokens", lambda _m, t: len(t))
|
||||
monkeypatch.setattr(compaction, "context_window", lambda _m: window)
|
||||
monkeypatch.setattr(compaction, "output_limit", lambda _m: 0)
|
||||
context = ContextSettings()
|
||||
context.keep_tokens = keep_tokens
|
||||
context.compact_buffer_tokens = 0
|
||||
context.summary_max_tokens = 64
|
||||
context.auto_compact = True
|
||||
settings = SimpleNamespace(
|
||||
context=context,
|
||||
llm=SimpleNamespace(api_key=None, api_base=None, timeout=1),
|
||||
)
|
||||
monkeypatch.setattr(compaction, "load_settings", lambda: settings)
|
||||
|
||||
|
||||
def _patch_summary(monkeypatch: pytest.MonkeyPatch, text: str) -> None:
|
||||
async def fake_acompletion(**_kwargs: Any) -> Any:
|
||||
message = SimpleNamespace(content=text)
|
||||
return SimpleNamespace(choices=[SimpleNamespace(message=message)])
|
||||
|
||||
monkeypatch.setattr("strix.llm.compaction.litellm.acompletion", fake_acompletion)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_maybe_compact_noop_when_within_budget(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_patch_budget(monkeypatch, keep_tokens=50, window=1_000_000)
|
||||
session = FakeSession(_turns(10))
|
||||
before = await session.get_items()
|
||||
|
||||
assert await compaction.maybe_compact(session, model="m") is False
|
||||
assert await session.get_items() == before
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_maybe_compact_rewrites_and_keeps_pairs(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_patch_budget(monkeypatch, keep_tokens=30, window=50)
|
||||
_patch_summary(monkeypatch, "SUMMARY BODY")
|
||||
session = FakeSession(_turns(12))
|
||||
|
||||
assert await compaction.maybe_compact(session, model="m") is True
|
||||
|
||||
items = await session.get_items()
|
||||
assert items[0]["role"] == "user"
|
||||
assert items[0]["content"].startswith(compaction._CHECKPOINT_TAG)
|
||||
assert "SUMMARY BODY" in items[0]["content"]
|
||||
assert len(items) < len(_turns(12))
|
||||
assert not _has_orphan_tool_output(items)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_maybe_compact_updates_previous_summary(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_patch_budget(monkeypatch, keep_tokens=30, window=50)
|
||||
captured: dict[str, str] = {}
|
||||
|
||||
async def fake_acompletion(**kwargs: Any) -> Any:
|
||||
captured["prompt"] = kwargs["messages"][0]["content"]
|
||||
return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content="NEW"))])
|
||||
|
||||
monkeypatch.setattr("strix.llm.compaction.litellm.acompletion", fake_acompletion)
|
||||
|
||||
prior = compaction._checkpoint_item("OLD SUMMARY TEXT")
|
||||
session = FakeSession([prior, *_turns(12)])
|
||||
|
||||
assert await compaction.maybe_compact(session, model="m") is True
|
||||
assert "OLD SUMMARY TEXT" in captured["prompt"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_maybe_compact_skips_when_summary_fails(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_patch_budget(monkeypatch, keep_tokens=30, window=50)
|
||||
|
||||
async def fake_acompletion(**_kwargs: Any) -> Any:
|
||||
raise RuntimeError("boom")
|
||||
|
||||
monkeypatch.setattr("strix.llm.compaction.litellm.acompletion", fake_acompletion)
|
||||
session = FakeSession(_turns(12))
|
||||
before = await session.get_items()
|
||||
|
||||
assert await compaction.maybe_compact(session, model="m") is False
|
||||
assert await session.get_items() == before
|
||||
@@ -0,0 +1,46 @@
|
||||
"""Tests for model-aware token budgets."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from strix.config import load_settings
|
||||
from strix.llm import context_budget
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import pytest
|
||||
|
||||
|
||||
def test_context_window_known_model() -> None:
|
||||
# gpt-4o is mapped by LiteLLM at 128k input tokens.
|
||||
assert context_budget.context_window("gpt-4o") == 128_000
|
||||
|
||||
|
||||
def test_context_window_strips_provider_prefix() -> None:
|
||||
assert context_budget.context_window("openai/gpt-4o") == 128_000
|
||||
|
||||
|
||||
def test_context_window_unmapped_uses_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
context_budget._model_info.cache_clear()
|
||||
|
||||
def _raise(_model: str) -> dict[str, int]:
|
||||
raise ValueError("This model isn't mapped yet.")
|
||||
|
||||
monkeypatch.setattr("strix.llm.context_budget.litellm.get_model_info", _raise)
|
||||
expected = load_settings().context.fallback_context_tokens
|
||||
assert context_budget.context_window("totally-made-up-model") == expected
|
||||
context_budget._model_info.cache_clear()
|
||||
|
||||
|
||||
def test_count_tokens_fallback_on_error(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def _raise(**_kwargs: object) -> int:
|
||||
raise RuntimeError("no tokenizer")
|
||||
|
||||
monkeypatch.setattr("strix.llm.context_budget.litellm.token_counter", _raise)
|
||||
text = "x" * 400
|
||||
assert context_budget.count_tokens("weird-model", text) == 100
|
||||
|
||||
|
||||
def test_count_tokens_empty_is_zero() -> None:
|
||||
assert context_budget.count_tokens("gpt-4o", "") == 0
|
||||
Reference in New Issue
Block a user