mirror of
https://github.com/usestrix/strix.git
synced 2026-08-21 10:48:59 +02:00
Ahmed Allam
8fd871249c
Head+tail-truncate the serialized history before the summary request so a very large head cannot itself overflow the model context and abort the recovery it is meant to perform.
218 lines
7.8 KiB
Python
218 lines
7.8 KiB
Python
"""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"]
|
|
|
|
|
|
def test_fit_to_tokens_truncates_oversized_text(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr(compaction, "count_tokens", lambda _m, t: len(t))
|
|
text = "x" * 10_000
|
|
|
|
fitted = compaction._fit_to_tokens("m", text, 500)
|
|
|
|
assert len(fitted) <= 500
|
|
assert compaction._HEAD_TRUNCATED_MARKER in fitted
|
|
# Small text is returned untouched.
|
|
assert compaction._fit_to_tokens("m", "short", 500) == "short"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_maybe_compact_bounds_summary_prompt(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
# A tiny window with a huge head must not send an oversized summary request.
|
|
_patch_budget(monkeypatch, keep_tokens=30, window=4_000)
|
|
captured: dict[str, str] = {}
|
|
|
|
async def fake_acompletion(**kwargs: Any) -> Any:
|
|
captured["prompt"] = kwargs["messages"][0]["content"]
|
|
return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content="S"))])
|
|
|
|
monkeypatch.setattr("strix.llm.compaction.litellm.acompletion", fake_acompletion)
|
|
big_turns = [{"role": "user", "content": "y" * 2_000} for _ in range(50)]
|
|
session = FakeSession(big_turns)
|
|
|
|
assert await compaction.maybe_compact(session, model="m") is True
|
|
# count_tokens==len(chars); prompt must fit the model window.
|
|
assert len(captured["prompt"]) <= 4_000
|
|
|
|
|
|
@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
|