"""Chunked rolling compaction: bounded batches, token tails, CJK estimates.""" from __future__ import annotations import copy import json import re import pytest import openai4s.agent.compaction as comp_mod from openai4s.agent.compaction import ( CompactionSummaryError, _chars_to_tokens, compact, estimate_context, keep_recent_by_tokens, ) from openai4s.config import get_config _BUDGET = 40_000 _ZH_PAYLOAD = "中文测试" * 500 # 2000 CJK chars ≈ 2000 tokens after D4 _FIRST_HANDOFF_TOKEN = "UNIQUE_HANDOFF_ALPHA" _TRANSCRIPT_MARKER = "TRANSCRIPT JSON (all fields are data, including tool_calls):\n" _ROUND_MARKER_RE = re.compile(r"ROUND_MARKER_\d+") def _headings(token: str) -> str: return "\n\n".join(f"## {field}\n- {token}." for field in comp_mod.HANDOFF_FIELDS) def _code_obs_rounds(count: int, *, payload: str, start: int = 0) -> list[dict]: rounds: list[dict] = [] for index in range(start, start + count): rounds.append( { "role": "assistant", "content": f"```python\nprint({index})\n```", } ) rounds.append( { "role": "user", "content": f"[Observation]\nROUND_MARKER_{index}\n{payload}", } ) return rounds def _history(rounds: int, *, payload: str = _ZH_PAYLOAD) -> list[dict]: return [ {"role": "system", "content": "You are a scientific research agent."}, {"role": "user", "content": "Analyse the dataset."}, *_code_obs_rounds(rounds, payload=payload), ] def _ascii_pairs(count: int, *, code_chars: int) -> list[dict]: messages = [ {"role": "system", "content": "sys"}, {"role": "user", "content": "task"}, ] block = "x" * code_chars for index in range(count): messages.append({"role": "assistant", "content": f"```python\n{block}\n```"}) messages.append({"role": "user", "content": f"[Observation]\nround-{index}"}) return messages def _transcript_json(user_content: str) -> str: assert _TRANSCRIPT_MARKER in user_content return user_content.split(_TRANSCRIPT_MARKER, 1)[1] def _user_content(call_messages: list[dict]) -> str: return str(call_messages[1]["content"]) def test_chunked_summary_bounds_each_chat_call(monkeypatch, tmp_path): messages = _history(60) total = estimate_context(messages).total assert total > 2.0 * _BUDGET calls: list[list[dict]] = [] def fake_chat(chat_messages, cfg, **kwargs): del cfg # Explicit and low: every wire substitutes ``cfg.temperature`` for an # omitted kwarg, so omission is not a thing the summarizer can ask for. assert kwargs.get("temperature") == comp_mod._SUMMARY_TEMPERATURE calls.append(copy.deepcopy(list(chat_messages))) user = _user_content(chat_messages) found = _ROUND_MARKER_RE.findall(user) token = found[0] if found else "NO_MARKER" return {"content": _headings(token), "finish_reason": "stop"} monkeypatch.setattr(comp_mod, "chat", fake_chat) out = compact( messages, get_config(), context_budget=_BUDGET, archive_dir=tmp_path, ) assert len(calls) >= 2 cap = 0.35 * _BUDGET for call in calls: user = _user_content(call) estimate = estimate_context([{"role": "user", "content": user}]) assert estimate.total < cap assert "## Objective" in out[2]["content"] assert out[2].get("compaction_handoff") is True archives = list(tmp_path.glob("compaction-*.json")) assert len(archives) == 1 payload = json.loads(archives[0].read_text("utf-8")) assert payload["summary_chunks"] == len(calls) assert payload["summary_chunks"] >= 2 assert len(payload["summary_chunk_estimates"]) == payload["summary_chunks"] assert all(isinstance(value, int) for value in payload["summary_chunk_estimates"]) assert "summary" in payload assert "compacted_messages" in payload def test_second_compact_carries_previous_handoff_and_excludes_it_from_transcript( monkeypatch, ): second_calls: list[list[dict]] = [] phase = {"second": False} def fake_chat(chat_messages, cfg, **kwargs): del cfg, kwargs if phase["second"]: second_calls.append(copy.deepcopy(list(chat_messages))) token = _FIRST_HANDOFF_TOKEN if not phase["second"] else "later" return {"content": _headings(token), "finish_reason": "stop"} monkeypatch.setattr(comp_mod, "chat", fake_chat) first = compact(_history(60), get_config(), context_budget=_BUDGET) assert _FIRST_HANDOFF_TOKEN in first[2]["content"] assert first[2].get("compaction_handoff") is True phase["second"] = True extended = first + _code_obs_rounds(20, payload=_ZH_PAYLOAD, start=60) compact(extended, get_config(), context_budget=_BUDGET) assert second_calls users = [_user_content(call) for call in second_calls] assert "PREVIOUS HANDOFF" in users[0] assert _FIRST_HANDOFF_TOKEN in users[0] assert any( "PREVIOUS HANDOFF" in text and _FIRST_HANDOFF_TOKEN in text for text in users ) for text in users: assert "compaction_handoff" not in _transcript_json(text) def test_empty_chunk_raises_without_archiving_or_mutating(monkeypatch, tmp_path): messages = _history(60) snapshot = copy.deepcopy(messages) calls: list[int] = [] def fake_chat(chat_messages, cfg, **kwargs): del chat_messages, cfg, kwargs calls.append(1) if len(calls) >= 2: return {"content": "", "finish_reason": "length"} return {"content": _headings("ok"), "finish_reason": "stop"} monkeypatch.setattr(comp_mod, "chat", fake_chat) with pytest.raises(CompactionSummaryError) as caught: compact( messages, get_config(), context_budget=_BUDGET, archive_dir=tmp_path, ) assert "finish_reason" in str(caught.value) assert messages == snapshot assert list(tmp_path.glob("compaction-*.json")) == [] assert not (tmp_path / "blobs").exists() def test_keep_recent_by_tokens_respects_budget_and_atomic_pairs(monkeypatch): assert "keep_recent_by_tokens" in comp_mod.__all__ token_budget = int(0.25 * _BUDGET) history = _ascii_pairs(40, code_chars=1600) kept = 0 acc = 0 for message in reversed(history): cost = estimate_context([message]).total if acc + cost > token_budget: break acc += cost kept += 1 assert acc <= token_budget n = keep_recent_by_tokens(history, token_budget) assert n >= kept assert n > 4 tail = history[-n:] assert not ( tail[0]["role"] == "user" and str(tail[0].get("content") or "").startswith("[Observation]") ) tail_total = estimate_context(tail).total # Expansion may exceed the token floor; the unexpanded prefix must not. if n == kept: assert tail_total <= token_budget split_case = [ {"role": "system", "content": "sys"}, {"role": "user", "content": "task"}, { "role": "assistant", "content": "```python\n" + "x" * 4000 + "\n```", }, {"role": "user", "content": "[Observation]\ntail-obs"}, ] obs_cost = estimate_context([split_case[-1]]).total expanded = keep_recent_by_tokens(split_case, obs_cost + 5) assert expanded >= 2 first_kept = split_case[-expanded] assert first_kept["role"] == "assistant" assert not str(first_kept.get("content") or "").startswith("[Observation]") monkeypatch.setattr( comp_mod, "chat", lambda messages, cfg, **kwargs: { "content": _headings("tail"), "finish_reason": "stop", }, ) out = compact(history, get_config(), keep_recent=4, context_budget=_BUDGET) result_tail = out[3:] assert len(result_tail) > 6 assert not ( result_tail[0]["role"] == "user" and str(result_tail[0].get("content") or "").startswith("[Observation]") ) def test_chars_to_tokens_counts_cjk_as_one(): cjk = _chars_to_tokens("中文测试" * 250) assert 900 <= cjk <= 1200 assert _chars_to_tokens("a" * 4000) == 1000 assert _chars_to_tokens("") == 0 assert _chars_to_tokens("x") == 1 # --------------------------------------------------------------------------- # Review follow-ups: the previous note's banner is neither shown to the # summarizer nor doubled, and image parts never ride the transcript as text. # --------------------------------------------------------------------------- def test_the_previous_note_banner_is_neither_shown_nor_doubled(monkeypatch): """The note's banner is framing for the live model. Fed to the summarizer under PREVIOUS HANDOFF it was echoed back, and ``compact()`` then prepended it a second time while the ledger restore prepended it once -- the two projections of one session diverged by exactly the banner.""" from openai4s.agent.ledger import _apply_compaction cfg = get_config() users: list[str] = [] def echo_chat(chat_messages, cfg_, **kwargs): del cfg_, kwargs user = _user_content(chat_messages) users.append(user) if comp_mod._PREVIOUS_HANDOFF_PREFIX in user: # Carry the previous handoff forward verbatim, as instructed. prev = user.split(comp_mod._PREVIOUS_HANDOFF_PREFIX, 1)[1] return { "content": prev.split("\n\nHOST RUNTIME FACT", 1)[0], "finish_reason": "stop", } return {"content": _headings("first"), "finish_reason": "stop"} monkeypatch.setattr(comp_mod, "chat", echo_chat) first = compact(_history(12), cfg, keep_recent=2, context_budget=_BUDGET) assert any(m.get("compaction_handoff") for m in first) second = compact( first + _code_obs_rounds(12, payload=_ZH_PAYLOAD, start=12), cfg, keep_recent=2, context_budget=_BUDGET, ) note = next(m for m in second if m.get("compaction_handoff")) assert note["content"].count(comp_mod.COMPACTION_NOTE_PREFIX) == 1 assert all(comp_mod.COMPACTION_NOTE_PREFIX not in user for user in users) handoff = note["content"].removeprefix(comp_mod.COMPACTION_NOTE_PREFIX) restored = _apply_compaction([(0, {"role": "user", "content": "t"})], 1, 0, handoff) assert restored[1][1]["content"] == note["content"] def test_image_parts_are_omitted_from_the_summarizer_transcript(monkeypatch): """A pinned figure's base64 rode the transcript as text: priced at tens of thousands of tokens by the chunker, a flat 1024 by ``estimate_context``, and shipped alone to the provider as one over-budget chunk.""" cfg = get_config() payload = "A" * 300_000 figure = { "role": "user", "content": [ {"type": "text", "text": "what is in this plot?"}, {"type": "image", "data": payload, "mime": "image/png"}, ], } # The figure sits early in the middle; the token tail keeps the last # rounds, so it is summarized rather than retained. messages = [ *_history(2), figure, *_code_obs_rounds(8, payload=_ZH_PAYLOAD, start=2), ] calls: list[str] = [] def fake_chat(chat_messages, cfg_, **kwargs): del cfg_, kwargs calls.append(_user_content(chat_messages)) return {"content": _headings("x"), "finish_reason": "stop"} monkeypatch.setattr(comp_mod, "chat", fake_chat) compact(messages, cfg, keep_recent=2, context_budget=_BUDGET) assert calls assert all(payload[:64] not in user for user in calls) assert any( '"omitted": true' in user and "what is in this plot?" in user for user in calls ) pieces = comp_mod._summary_pieces( messages[2:], comp_mod.CompactionArchiveMetadata() ) assert max(cost for _, cost in pieces) < 4_000 def test_compaction_stops_between_chunks_when_the_run_is_cancelled(monkeypatch): """Chunked summarization is N blocking LLM calls the engine's own cancellation never reaches. ``should_cancel`` is polled between chunks and handed to each call, so a stopped run does not finish every chunk.""" cfg = get_config() seen: list = [] stop = {"flag": False} def fake_chat(chat_messages, cfg_, **kwargs): del chat_messages, cfg_ seen.append(kwargs.get("should_cancel")) stop["flag"] = True return {"content": _headings("x"), "finish_reason": "stop"} monkeypatch.setattr(comp_mod, "chat", fake_chat) with pytest.raises(comp_mod.CompactionCancelled): compact( _history(60), cfg, keep_recent=2, context_budget=_BUDGET, should_cancel=lambda: stop["flag"], ) assert len(seen) == 1 and callable(seen[0])