|
|
@@ -102,7 +102,13 @@ def test_chat_session_appends_and_keeps_prefix():
|
|
|
|
|
|
|
|
|
def test_chat_session_trims_in_blocks_and_keeps_system_prompt():
|
|
|
- cfg = LlmConfig(max_prompt_tokens=50, trim_block_turns=1)
|
|
|
+ cfg = LlmConfig(
|
|
|
+ max_context_tokens=1330,
|
|
|
+ max_tokens=1024,
|
|
|
+ compaction_safety_margin=256,
|
|
|
+ trim_block_turns=1,
|
|
|
+ compaction_enabled=False,
|
|
|
+ ) # prompt_budget == 50
|
|
|
client, calls = _make_client(lambda _p: _reply(content="ok"))
|
|
|
session = ChatSession(cfg, client)
|
|
|
|
|
|
@@ -132,3 +138,449 @@ def test_chat_session_drops_reply_when_reset_mid_flight():
|
|
|
|
|
|
assert holder[0].send("hi") == ""
|
|
|
assert holder[0].history == []
|
|
|
+
|
|
|
+
|
|
|
+# ---------------------------------------------------------------------------
|
|
|
+# Phase C tests
|
|
|
+# ---------------------------------------------------------------------------
|
|
|
+
|
|
|
+def _compact_cfg(
|
|
|
+ *,
|
|
|
+ budget_chars=200,
|
|
|
+ trigger_ratio=0.8,
|
|
|
+ keep_turns=1,
|
|
|
+ compaction_max_tokens=64,
|
|
|
+ compaction_timeout_s=5.0,
|
|
|
+ compaction_input_tokens=8192,
|
|
|
+ compaction_enabled=True,
|
|
|
+ trim_block_turns=1,
|
|
|
+):
|
|
|
+ """Return a small LlmConfig whose prompt_budget == budget_chars // 4 tokens.
|
|
|
+
|
|
|
+ prompt_budget = max_context_tokens - max_tokens - compaction_safety_margin
|
|
|
+ We want budget_chars // 4 tokens, so:
|
|
|
+ max_context_tokens = (budget_chars // 4) + max_tokens + compaction_safety_margin
|
|
|
+ """
|
|
|
+ max_tokens = 16
|
|
|
+ compaction_safety_margin = 0
|
|
|
+ max_context_tokens = (budget_chars // 4) + max_tokens + compaction_safety_margin
|
|
|
+ return LlmConfig(
|
|
|
+ max_context_tokens=max_context_tokens,
|
|
|
+ max_tokens=max_tokens,
|
|
|
+ compaction_safety_margin=compaction_safety_margin,
|
|
|
+ compaction_enabled=compaction_enabled,
|
|
|
+ compaction_trigger_ratio=trigger_ratio,
|
|
|
+ compaction_keep_turns=keep_turns,
|
|
|
+ compaction_max_tokens=compaction_max_tokens,
|
|
|
+ compaction_timeout_s=compaction_timeout_s,
|
|
|
+ compaction_input_tokens=compaction_input_tokens,
|
|
|
+ trim_block_turns=trim_block_turns,
|
|
|
+ system_prompt_chat="SYS",
|
|
|
+ system_prompt_command="CMD",
|
|
|
+ compaction_system_prompt="COMPACT",
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+def test_compaction_fires_and_window_shrinks():
|
|
|
+ """Crossing the trigger compacts: window shrinks, summary present in next prompt."""
|
|
|
+ cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
|
|
|
+ # prompt_budget = 200//4 = 50 tokens; trigger at 40 tokens (0.8 * 50)
|
|
|
+
|
|
|
+ summary_text = "summary of old turns"
|
|
|
+ call_index = 0
|
|
|
+
|
|
|
+ def responder(payload):
|
|
|
+ nonlocal call_index
|
|
|
+ msgs = payload["body"]["messages"]
|
|
|
+ call_index += 1
|
|
|
+ # Summarizer calls have the compaction system prompt.
|
|
|
+ if msgs[0]["content"] == "COMPACT":
|
|
|
+ return _reply(content=summary_text)
|
|
|
+ return _reply(content="ok")
|
|
|
+
|
|
|
+ client, calls = _make_client(responder, cfg)
|
|
|
+ session = ChatSession(cfg, client)
|
|
|
+
|
|
|
+ # Send enough turns to exceed the trigger (each "x"*40 = 40 chars = 10 tokens).
|
|
|
+ # After 4 turns (user+assistant pairs = 8 messages * 10 tokens = 80 tokens > 40 trigger).
|
|
|
+ for _ in range(4):
|
|
|
+ session.send("x" * 40)
|
|
|
+
|
|
|
+ # After compaction the window must be smaller than the full history.
|
|
|
+ assert len(session.history) <= 4 # keep_turns=1 → keep=2 messages
|
|
|
+
|
|
|
+ # The summary must be set.
|
|
|
+ assert session.summary != ""
|
|
|
+
|
|
|
+ # The next chat call must include the summary as a second system message,
|
|
|
+ # with the chat system prompt FIRST.
|
|
|
+ session.send("after compaction")
|
|
|
+ chat_calls = [c for c in calls if c["body"]["messages"][0]["content"] == "SYS"]
|
|
|
+ last_chat = chat_calls[-1]["body"]["messages"]
|
|
|
+ assert last_chat[0]["role"] == "system"
|
|
|
+ assert last_chat[0]["content"] == "SYS"
|
|
|
+ assert last_chat[1]["role"] == "system"
|
|
|
+ assert "Earlier conversation summary:" in last_chat[1]["content"]
|
|
|
+ assert summary_text in last_chat[1]["content"]
|
|
|
+
|
|
|
+
|
|
|
+def test_summary_never_appended_to_history():
|
|
|
+ """The summary text must not appear as a history entry."""
|
|
|
+ cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
|
|
|
+ summary_text = "SUMMARY_SENTINEL"
|
|
|
+
|
|
|
+ def responder(payload):
|
|
|
+ msgs = payload["body"]["messages"]
|
|
|
+ if msgs[0]["content"] == "COMPACT":
|
|
|
+ return _reply(content=summary_text)
|
|
|
+ return _reply(content="ok")
|
|
|
+
|
|
|
+ client, _ = _make_client(responder, cfg)
|
|
|
+ session = ChatSession(cfg, client)
|
|
|
+
|
|
|
+ for _ in range(4):
|
|
|
+ session.send("x" * 40)
|
|
|
+
|
|
|
+ for msg in session.history:
|
|
|
+ assert summary_text not in msg["content"]
|
|
|
+
|
|
|
+
|
|
|
+def test_summarizer_failure_falls_back_to_trim():
|
|
|
+ """On summarizer failure the session falls back to _trim() and stays within budget."""
|
|
|
+ cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
|
|
|
+
|
|
|
+ def responder(payload):
|
|
|
+ msgs = payload["body"]["messages"]
|
|
|
+ if msgs[0]["content"] == "COMPACT":
|
|
|
+ return _reply(content="", finish="stop") # empty → failure
|
|
|
+ return _reply(content="ok")
|
|
|
+
|
|
|
+ client, _ = _make_client(responder, cfg)
|
|
|
+ session = ChatSession(cfg, client)
|
|
|
+
|
|
|
+ for _ in range(4):
|
|
|
+ session.send("x" * 40)
|
|
|
+
|
|
|
+ # Summary must remain empty (fallback used trim, not compaction).
|
|
|
+ assert session.summary == ""
|
|
|
+ # History must be bounded within the budget.
|
|
|
+ assert session._estimate_tokens() <= cfg.prompt_budget
|
|
|
+
|
|
|
+
|
|
|
+def test_summarizer_length_finish_reason_falls_back_to_trim():
|
|
|
+ """finish_reason == 'length' on the summarizer is treated as failure."""
|
|
|
+ cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
|
|
|
+
|
|
|
+ def responder(payload):
|
|
|
+ msgs = payload["body"]["messages"]
|
|
|
+ if msgs[0]["content"] == "COMPACT":
|
|
|
+ return _reply(content="partial summary", finish="length")
|
|
|
+ return _reply(content="ok")
|
|
|
+
|
|
|
+ client, _ = _make_client(responder, cfg)
|
|
|
+ session = ChatSession(cfg, client)
|
|
|
+
|
|
|
+ for _ in range(4):
|
|
|
+ session.send("x" * 40)
|
|
|
+
|
|
|
+ assert session.summary == ""
|
|
|
+ assert session._estimate_tokens() <= cfg.prompt_budget
|
|
|
+
|
|
|
+
|
|
|
+def test_summarizer_llmerror_falls_back_to_trim():
|
|
|
+ """LlmError from the summarizer falls back to _trim()."""
|
|
|
+ cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
|
|
|
+
|
|
|
+ def responder(payload):
|
|
|
+ msgs = payload["body"]["messages"]
|
|
|
+ if msgs[0]["content"] == "COMPACT":
|
|
|
+ raise LlmError("network error")
|
|
|
+ return _reply(content="ok")
|
|
|
+
|
|
|
+ client, _ = _make_client(responder, cfg)
|
|
|
+ session = ChatSession(cfg, client)
|
|
|
+
|
|
|
+ for _ in range(4):
|
|
|
+ session.send("x" * 40)
|
|
|
+
|
|
|
+ assert session.summary == ""
|
|
|
+ assert session._estimate_tokens() <= cfg.prompt_budget
|
|
|
+
|
|
|
+
|
|
|
+def test_reset_clears_summary():
|
|
|
+ """reset() clears both history and summary."""
|
|
|
+ cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
|
|
|
+ summary_text = "old summary"
|
|
|
+
|
|
|
+ def responder(payload):
|
|
|
+ msgs = payload["body"]["messages"]
|
|
|
+ if msgs[0]["content"] == "COMPACT":
|
|
|
+ return _reply(content=summary_text)
|
|
|
+ return _reply(content="ok")
|
|
|
+
|
|
|
+ client, _ = _make_client(responder, cfg)
|
|
|
+ session = ChatSession(cfg, client)
|
|
|
+
|
|
|
+ for _ in range(4):
|
|
|
+ session.send("x" * 40)
|
|
|
+
|
|
|
+ assert session.summary != ""
|
|
|
+
|
|
|
+ session.reset()
|
|
|
+ assert session.history == []
|
|
|
+ assert session.summary == ""
|
|
|
+
|
|
|
+
|
|
|
+def test_compaction_disabled_never_calls_summarizer():
|
|
|
+ """compaction_enabled=False never calls the summarizer and still trims."""
|
|
|
+ cfg = _compact_cfg(
|
|
|
+ budget_chars=200,
|
|
|
+ trigger_ratio=0.8,
|
|
|
+ keep_turns=1,
|
|
|
+ compaction_enabled=False,
|
|
|
+ )
|
|
|
+
|
|
|
+ summarizer_called = False
|
|
|
+
|
|
|
+ def responder(payload):
|
|
|
+ nonlocal summarizer_called
|
|
|
+ msgs = payload["body"]["messages"]
|
|
|
+ if msgs[0]["content"] == "COMPACT":
|
|
|
+ summarizer_called = True
|
|
|
+ return _reply(content="should not happen")
|
|
|
+ return _reply(content="ok")
|
|
|
+
|
|
|
+ client, _ = _make_client(responder, cfg)
|
|
|
+ session = ChatSession(cfg, client)
|
|
|
+
|
|
|
+ for _ in range(6):
|
|
|
+ session.send("x" * 40)
|
|
|
+
|
|
|
+ assert not summarizer_called
|
|
|
+ assert session.summary == ""
|
|
|
+ # History must still be bounded (trim ran).
|
|
|
+ assert session._estimate_tokens() <= cfg.prompt_budget
|
|
|
+
|
|
|
+
|
|
|
+def test_nothing_to_evict_skips_summarizer():
|
|
|
+ """When len(history) <= keep, compaction is skipped and trim runs instead."""
|
|
|
+ # keep_turns=10 → keep=20 messages; we only send 2 turns (4 messages).
|
|
|
+ cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=10)
|
|
|
+
|
|
|
+ summarizer_called = False
|
|
|
+
|
|
|
+ def responder(payload):
|
|
|
+ nonlocal summarizer_called
|
|
|
+ msgs = payload["body"]["messages"]
|
|
|
+ if msgs[0]["content"] == "COMPACT":
|
|
|
+ summarizer_called = True
|
|
|
+ return _reply(content="should not happen")
|
|
|
+ return _reply(content="ok")
|
|
|
+
|
|
|
+ client, _ = _make_client(responder, cfg)
|
|
|
+ session = ChatSession(cfg, client)
|
|
|
+
|
|
|
+ # Push over the trigger but keep history shorter than keep=20.
|
|
|
+ for _ in range(4):
|
|
|
+ session.send("x" * 40)
|
|
|
+
|
|
|
+ assert not summarizer_called
|
|
|
+
|
|
|
+
|
|
|
+def test_chunk_and_fold_large_evicted():
|
|
|
+ """Large evicted list triggers chunk-and-fold (more than one summarizer call).
|
|
|
+
|
|
|
+ budget_chars=160 → prompt_budget=40 tokens, trigger at 32 tokens (0.8 * 40).
|
|
|
+ Each turn: user='x'*40 (40 chars) + assistant='ok' (2 chars) = ~10 tokens.
|
|
|
+ Compaction fires at turn 4 with 6 evicted messages.
|
|
|
+ compaction_input_tokens=1 → char_limit=4 chars per chunk, so each 40-char
|
|
|
+ or 2-char message is its own chunk → 6 summarizer calls for 6 chunks.
|
|
|
+ """
|
|
|
+ cfg = _compact_cfg(
|
|
|
+ budget_chars=160,
|
|
|
+ trigger_ratio=0.8,
|
|
|
+ keep_turns=1,
|
|
|
+ compaction_input_tokens=1, # 1 token * 4 = 4 chars per chunk → each msg is a chunk
|
|
|
+ )
|
|
|
+
|
|
|
+ summarizer_calls = []
|
|
|
+
|
|
|
+ def responder(payload):
|
|
|
+ msgs = payload["body"]["messages"]
|
|
|
+ if msgs[0]["content"] == "COMPACT":
|
|
|
+ summarizer_calls.append(payload)
|
|
|
+ return _reply(content="folded")
|
|
|
+ return _reply(content="ok")
|
|
|
+
|
|
|
+ client, _ = _make_client(responder, cfg)
|
|
|
+ session = ChatSession(cfg, client)
|
|
|
+
|
|
|
+ # 4 turns → 8 messages → fires compaction with 6 evicted (keep=2).
|
|
|
+ for _ in range(4):
|
|
|
+ session.send("x" * 40)
|
|
|
+
|
|
|
+ # With compaction_input_tokens=1 (4 chars/chunk) and 40-char messages,
|
|
|
+ # each message is its own chunk → multiple summarizer calls.
|
|
|
+ assert len(summarizer_calls) > 1
|
|
|
+
|
|
|
+
|
|
|
+def test_chat_per_request_max_tokens_override():
|
|
|
+ """LlmClient.chat sends per-request max_tokens when given, cfg default otherwise."""
|
|
|
+ cfg = LlmConfig(max_tokens=1024, timeout_s=60.0)
|
|
|
+ client, calls = _make_client(lambda _p: _reply(content="ok"), cfg)
|
|
|
+
|
|
|
+ # Default: uses cfg.max_tokens.
|
|
|
+ client.chat([{"role": "user", "content": "hi"}])
|
|
|
+ assert calls[0]["body"]["max_tokens"] == 1024
|
|
|
+
|
|
|
+ # Override: uses the supplied value.
|
|
|
+ client.chat([{"role": "user", "content": "hi"}], max_tokens=128)
|
|
|
+ assert calls[1]["body"]["max_tokens"] == 128
|
|
|
+
|
|
|
+
|
|
|
+def test_chat_per_request_timeout_override():
|
|
|
+ """LlmClient.chat uses per-request timeout_s when given, cfg default otherwise."""
|
|
|
+ cfg = LlmConfig(max_tokens=1024, timeout_s=60.0)
|
|
|
+ recorded_timeouts: list[float] = []
|
|
|
+
|
|
|
+ def opener(payload, timeout):
|
|
|
+ recorded_timeouts.append(timeout)
|
|
|
+ return _reply(content="ok")
|
|
|
+
|
|
|
+ client = LlmClient(cfg, opener=opener)
|
|
|
+
|
|
|
+ # Default timeout.
|
|
|
+ client.chat([{"role": "user", "content": "hi"}])
|
|
|
+ assert recorded_timeouts[0] == 60.0
|
|
|
+
|
|
|
+ # Per-request override.
|
|
|
+ client.chat([{"role": "user", "content": "hi"}], timeout_s=5.0)
|
|
|
+ assert recorded_timeouts[1] == 5.0
|
|
|
+
|
|
|
+
|
|
|
+def test_summarizer_uses_compaction_max_tokens_and_timeout():
|
|
|
+ """The summarizer call uses compaction_max_tokens and compaction_timeout_s."""
|
|
|
+ cfg = _compact_cfg(
|
|
|
+ budget_chars=200,
|
|
|
+ trigger_ratio=0.8,
|
|
|
+ keep_turns=1,
|
|
|
+ compaction_max_tokens=42,
|
|
|
+ compaction_timeout_s=7.0,
|
|
|
+ )
|
|
|
+
|
|
|
+ recorded: list[dict] = []
|
|
|
+
|
|
|
+ def opener(payload, timeout):
|
|
|
+ msgs = payload["body"]["messages"]
|
|
|
+ if msgs[0]["content"] == "COMPACT":
|
|
|
+ recorded.append({"max_tokens": payload["body"]["max_tokens"], "timeout": timeout})
|
|
|
+ return _reply(content="summary")
|
|
|
+ return _reply(content="ok")
|
|
|
+
|
|
|
+ client = LlmClient(cfg, opener=opener)
|
|
|
+ session = ChatSession(cfg, client)
|
|
|
+
|
|
|
+ for _ in range(4):
|
|
|
+ session.send("x" * 40)
|
|
|
+
|
|
|
+ assert len(recorded) >= 1
|
|
|
+ assert recorded[0]["max_tokens"] == 42
|
|
|
+ assert recorded[0]["timeout"] == 7.0
|
|
|
+
|
|
|
+
|
|
|
+def test_notify_called_on_compaction():
|
|
|
+ """notify('compacted', turns=N) is called after a successful compaction."""
|
|
|
+ cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
|
|
|
+
|
|
|
+ notifications: list[tuple] = []
|
|
|
+
|
|
|
+ def notify(kind, **fields):
|
|
|
+ notifications.append((kind, fields))
|
|
|
+
|
|
|
+ def responder(payload):
|
|
|
+ msgs = payload["body"]["messages"]
|
|
|
+ if msgs[0]["content"] == "COMPACT":
|
|
|
+ return _reply(content="summary")
|
|
|
+ return _reply(content="ok")
|
|
|
+
|
|
|
+ client = LlmClient(cfg, opener=lambda p, t: responder(p))
|
|
|
+ session = ChatSession(cfg, client, notify=notify)
|
|
|
+
|
|
|
+ for _ in range(4):
|
|
|
+ session.send("x" * 40)
|
|
|
+
|
|
|
+ assert any(kind == "compacted" for kind, _ in notifications)
|
|
|
+ compacted = [(k, f) for k, f in notifications if k == "compacted"]
|
|
|
+ assert compacted[0][1]["turns"] > 0
|
|
|
+
|
|
|
+
|
|
|
+def test_notify_not_called_when_none():
|
|
|
+ """ChatSession works without a notify callback (default None)."""
|
|
|
+ cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
|
|
|
+
|
|
|
+ def responder(payload):
|
|
|
+ msgs = payload["body"]["messages"]
|
|
|
+ if msgs[0]["content"] == "COMPACT":
|
|
|
+ return _reply(content="summary")
|
|
|
+ return _reply(content="ok")
|
|
|
+
|
|
|
+ client = LlmClient(cfg, opener=lambda p, t: responder(p))
|
|
|
+ session = ChatSession(cfg, client) # no notify
|
|
|
+
|
|
|
+ # Must not raise.
|
|
|
+ for _ in range(4):
|
|
|
+ session.send("x" * 40)
|
|
|
+
|
|
|
+
|
|
|
+def test_system_prompt_is_first_message_after_compaction():
|
|
|
+ """After compaction the chat system prompt is always the first message."""
|
|
|
+ cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
|
|
|
+
|
|
|
+ def responder(payload):
|
|
|
+ msgs = payload["body"]["messages"]
|
|
|
+ if msgs[0]["content"] == "COMPACT":
|
|
|
+ return _reply(content="summary text")
|
|
|
+ return _reply(content="ok")
|
|
|
+
|
|
|
+ client, calls = _make_client(responder, cfg)
|
|
|
+ session = ChatSession(cfg, client)
|
|
|
+
|
|
|
+ for _ in range(4):
|
|
|
+ session.send("x" * 40)
|
|
|
+
|
|
|
+ # Find the last chat call (system prompt == "SYS", not "COMPACT").
|
|
|
+ chat_calls = [c for c in calls if c["body"]["messages"][0]["content"] == "SYS"]
|
|
|
+ assert chat_calls, "no chat calls found after compaction"
|
|
|
+ last_msgs = chat_calls[-1]["body"]["messages"]
|
|
|
+ assert last_msgs[0]["role"] == "system"
|
|
|
+ assert last_msgs[0]["content"] == "SYS"
|
|
|
+
|
|
|
+
|
|
|
+def test_chat_session_truncated_empty_reply_notifies_and_is_not_appended():
|
|
|
+ """A length-limited empty reply is reported and never stored as an empty turn."""
|
|
|
+ notifications: list[tuple] = []
|
|
|
+
|
|
|
+ def notify(kind, **fields):
|
|
|
+ notifications.append((kind, fields))
|
|
|
+
|
|
|
+ client, _ = _make_client(lambda _p: _reply(content="", finish="length"))
|
|
|
+ session = ChatSession(LlmConfig(), client, notify=notify)
|
|
|
+
|
|
|
+ assert session.send("hi") == ""
|
|
|
+ assert session.history == [{"role": "user", "content": "hi"}]
|
|
|
+ assert notifications == [("truncated", {"message": "response truncated"})]
|
|
|
+
|
|
|
+
|
|
|
+def test_chat_session_truncated_partial_reply_is_appended_and_notified():
|
|
|
+ """A truncated but non-empty reply is kept and reported as truncated."""
|
|
|
+ notifications: list[tuple] = []
|
|
|
+
|
|
|
+ def notify(kind, **fields):
|
|
|
+ notifications.append((kind, fields))
|
|
|
+
|
|
|
+ client, _ = _make_client(lambda _p: _reply(content="partial answer", finish="length"))
|
|
|
+ session = ChatSession(LlmConfig(), client, notify=notify)
|
|
|
+
|
|
|
+ assert session.send("hi") == "partial answer"
|
|
|
+ assert session.history[-1] == {"role": "assistant", "content": "partial answer"}
|
|
|
+ assert any(kind == "truncated" for kind, _ in notifications)
|