dongarandClaude Opus 5.5
cb638fe687
Before each driver call, over compact_at (0.6) of the prompt budget, the same model is asked for notes over the very request it was about to get, plus one message, so the server's cache covers all but the notes. The notes, with code's list of the calls they replace, become a compaction event that renders in place of its range, keeping word for word what started the work and the last few turns. Rare, large steps keep the prompt's start unchanged between them. A later compaction folds in the earlier; one from an earlier run's view is hidden. Failed notes are logged and fitting takes over. Fitting now applies compactions before it measures, records what it cuts (Trim) and logs each cut as context.trimmed. Wired into both agents' drivers and the benchmark; nisse-dev's helper loop is a follow-up. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01J5Utyfzw7gF7UhjRDEtjTk
343 lines
12 KiB
Python
343 lines
12 KiB
Python
"""The context budget (architecture L6)."""
|
|
|
|
import io
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from nisse_core.log.canonical import (
|
|
AssistantText,
|
|
CanonicalEvent,
|
|
Compaction,
|
|
StoredEvent,
|
|
ToolCall,
|
|
ToolResult,
|
|
UserMessage,
|
|
)
|
|
from nisse_core.loop.driver import Driver
|
|
from nisse_core.models.catalog import ModelCatalog, UnknownLimits
|
|
from nisse_core.models.context import (
|
|
Budget,
|
|
BudgetedClient,
|
|
ContextOverflow,
|
|
estimate_tokens,
|
|
fit,
|
|
head_and_tail,
|
|
)
|
|
from nisse_core.oplog import OpLog
|
|
from nisse_core.state.store import StateStore
|
|
|
|
|
|
def stored(events: list[CanonicalEvent]) -> list[StoredEvent]:
|
|
return [StoredEvent(id=i + 1, ts="t", event=e) for i, e in enumerate(events)]
|
|
|
|
|
|
def turn(n: int, output: str) -> list[CanonicalEvent]:
|
|
"""One executor turn: a command and its (possibly long) output."""
|
|
return [
|
|
ToolCall(call_id=f"c{n}", tool="run_command", arguments={"command": "make"}),
|
|
ToolResult(call_id=f"c{n}", result=output, success=True),
|
|
]
|
|
|
|
|
|
def transcript(turns: int, output_chars: int) -> list[StoredEvent]:
|
|
events: list[CanonicalEvent] = [UserMessage(text="start milestone 1", source="t")]
|
|
for n in range(turns):
|
|
events += turn(n, f"output {n} " + "x" * output_chars)
|
|
return stored(events)
|
|
|
|
|
|
def budget(window: int, keep: int = 2) -> Budget:
|
|
return Budget(
|
|
window=window, max_output=500, keep_recent_turns=keep, safety_margin=0.1
|
|
)
|
|
|
|
|
|
def results(events: list[StoredEvent]) -> dict[str, str]:
|
|
return {
|
|
e.event.call_id: e.event.result
|
|
for e in events
|
|
if isinstance(e.event, ToolResult)
|
|
}
|
|
|
|
|
|
# --- keeping the end of long output ---------------------------------------
|
|
|
|
|
|
def test_head_and_tail_keeps_the_failure_summary_at_the_end() -> None:
|
|
output = (
|
|
"exit code: 1\nstdout:\n"
|
|
+ "collected 400 items\n" * 500
|
|
+ "FAILED tests/test_app.py::test_greet - AssertionError\n"
|
|
+ "==== 1 failed, 399 passed ====\n"
|
|
)
|
|
kept = head_and_tail(output, 1000)
|
|
assert len(kept) < 1100
|
|
assert kept.startswith("exit code: 1") # the head keeps the exit code
|
|
assert "FAILED tests/test_app.py::test_greet" in kept
|
|
assert "1 failed, 399 passed" in kept
|
|
assert "characters omitted" in kept
|
|
|
|
|
|
def test_short_text_is_untouched() -> None:
|
|
assert head_and_tail("fine", 1000) == "fine"
|
|
|
|
|
|
# --- fitting, step by step -------------------------------------------------
|
|
|
|
|
|
def test_a_transcript_that_fits_is_sent_unchanged() -> None:
|
|
events = transcript(turns=3, output_chars=100)
|
|
assert fit(events, "system", None, budget(100_000)) == events
|
|
|
|
|
|
def test_step_one_cuts_older_tool_results_and_keeps_recent_ones_whole() -> None:
|
|
events = transcript(turns=6, output_chars=3000)
|
|
fitted = fit(events, "system", None, budget(4000, keep=2))
|
|
|
|
kept = results(fitted)
|
|
assert len(kept) == 6 # nothing dropped yet
|
|
for older in ("c0", "c1", "c2", "c3"):
|
|
assert "characters omitted" in kept[older]
|
|
for recent in ("c4", "c5"):
|
|
assert "characters omitted" not in kept[recent]
|
|
|
|
|
|
def test_step_two_drops_the_oldest_turns_and_says_so() -> None:
|
|
events = transcript(turns=20, output_chars=3000)
|
|
fitted = fit(events, "system", None, budget(4000, keep=2))
|
|
|
|
first = fitted[0].event
|
|
assert isinstance(first, UserMessage)
|
|
assert "omitted to fit the context window" in first.text
|
|
kept = results(fitted)
|
|
assert "c0" not in kept # the oldest went
|
|
assert {"c18", "c19"} <= kept.keys() # the recent turns stayed
|
|
assert "characters omitted" not in kept["c19"]
|
|
# Every surviving result still has its call: a provider rejects an
|
|
# answer to a call it never saw.
|
|
calls = {e.event.call_id for e in fitted if isinstance(e.event, ToolCall)}
|
|
assert kept.keys() <= calls
|
|
assert estimate_tokens(fitted, "system", None) <= budget(4000).for_prompt
|
|
|
|
|
|
def test_step_three_refuses_what_cannot_fit() -> None:
|
|
# Even the two most recent turns are bigger than the whole window.
|
|
events = transcript(turns=3, output_chars=20_000)
|
|
with pytest.raises(ContextOverflow, match="last 2 turns"):
|
|
fit(events, "system", None, budget(2000, keep=2))
|
|
|
|
|
|
def test_a_huge_system_prompt_alone_overflows() -> None:
|
|
events = transcript(turns=1, output_chars=10)
|
|
with pytest.raises(ContextOverflow):
|
|
fit(events, "s" * 50_000, None, budget(4096))
|
|
|
|
|
|
def test_the_output_ceiling_and_margin_come_out_of_the_window() -> None:
|
|
b = Budget(window=4096, max_output=1024, keep_recent_turns=6, safety_margin=0.1)
|
|
assert b.for_prompt == int(4096 * 0.9) - 1024
|
|
|
|
|
|
# --- the client wrapper ----------------------------------------------------
|
|
|
|
|
|
class RecordingClient:
|
|
def __init__(self, reported_input: int = 0) -> None:
|
|
self.calls: list[dict[str, Any]] = []
|
|
self.last_usage = {"input": reported_input, "output": 5}
|
|
|
|
def complete(self, events, tools, model, system=None, *, max_tokens=None, **kwargs):
|
|
self.calls.append({"events": events, "max_tokens": max_tokens, **kwargs})
|
|
return [AssistantText(text="ok", model=model)]
|
|
|
|
|
|
def test_the_ceiling_is_the_smaller_of_config_and_model() -> None:
|
|
inner = RecordingClient()
|
|
client = BudgetedClient(inner, lambda m: (4096, 1024), max_output_tokens=16384)
|
|
client.complete(transcript(1, 10), None, "llama3.2")
|
|
assert inner.calls[0]["max_tokens"] == 1024
|
|
|
|
client = BudgetedClient(inner, lambda m: (200_000, 64_000), max_output_tokens=16384)
|
|
client.complete(transcript(1, 10), None, "claude-haiku-4-5")
|
|
assert inner.calls[1]["max_tokens"] == 16384
|
|
|
|
|
|
def test_the_wrapper_sends_the_fitted_events() -> None:
|
|
inner = RecordingClient()
|
|
client = BudgetedClient(
|
|
inner, lambda m: (4000, 500), keep_recent_turns=2, safety_margin=0.1
|
|
)
|
|
client.complete(transcript(20, 3000), None, "m", system="system")
|
|
sent = inner.calls[0]["events"]
|
|
assert "omitted to fit" in sent[0].event.text
|
|
|
|
|
|
def test_reported_usage_only_ever_raises_the_estimate() -> None:
|
|
events = transcript(1, 10)
|
|
estimate = estimate_tokens(events, None, None)
|
|
|
|
under = RecordingClient(reported_input=estimate * 2) # we undercounted
|
|
client = BudgetedClient(under, lambda m: (100_000, None))
|
|
client.complete(events, None, "m")
|
|
assert client._factors["m"] == pytest.approx(2.0, rel=0.05)
|
|
|
|
over = RecordingClient(reported_input=estimate // 2) # we overcounted
|
|
under_then_over = BudgetedClient(over, lambda m: (100_000, None))
|
|
under_then_over._factors["m"] = 2.0
|
|
under_then_over.complete(events, None, "m")
|
|
assert under_then_over._factors["m"] == 2.0 # never loosened
|
|
|
|
|
|
def test_a_request_that_fits_is_sent_whole_and_nothing_is_logged() -> None:
|
|
stream = io.StringIO()
|
|
client = BudgetedClient(
|
|
RecordingClient(), lambda m: (100_000, None), oplog=OpLog("t", stream)
|
|
)
|
|
events = transcript(3, 100)
|
|
client.complete(events, None, "m")
|
|
assert client.last_trim is None
|
|
assert stream.getvalue() == ""
|
|
|
|
|
|
def test_a_cut_is_recorded_and_logged() -> None:
|
|
stream = io.StringIO()
|
|
client = BudgetedClient(
|
|
RecordingClient(),
|
|
lambda m: (4000, 500),
|
|
keep_recent_turns=2,
|
|
oplog=OpLog("t", stream),
|
|
)
|
|
client.complete(transcript(20, 3000), None, "m", system="system")
|
|
assert client.last_trim is not None and client.last_trim.step == 2
|
|
assert '"event": "context.trimmed"' in stream.getvalue()
|
|
|
|
|
|
def test_a_compaction_counts_for_what_it_replaced() -> None:
|
|
events = transcript(20, 3000)
|
|
last = events[-1].id
|
|
compaction = StoredEvent(
|
|
id=last + 1,
|
|
ts="t",
|
|
event=Compaction(summary="notes", model="m", replaces=[2, last - 4]),
|
|
)
|
|
inner = RecordingClient()
|
|
client = BudgetedClient(inner, lambda m: (8000, 500), keep_recent_turns=2)
|
|
client.complete([*events, compaction], None, "m")
|
|
assert client.last_trim is None # fits once the summary stands in
|
|
sent = inner.calls[0]["events"]
|
|
assert isinstance(sent[1].event, Compaction)
|
|
|
|
|
|
def test_last_usage_passes_through() -> None:
|
|
client = BudgetedClient(RecordingClient(reported_input=7), lambda m: (100, None))
|
|
assert client.last_usage["input"] == 7
|
|
|
|
|
|
# --- the limits come from LiteLLM's catalog --------------------------------
|
|
|
|
|
|
CATALOG = ModelCatalog.from_payload(
|
|
{
|
|
"data": [
|
|
{
|
|
"model_name": "claude-haiku-4-5",
|
|
"litellm_params": {"model": "anthropic/claude-haiku-4-5"},
|
|
"model_info": {"max_input_tokens": 200000, "max_output_tokens": 64000},
|
|
},
|
|
{
|
|
"model_name": "llama3.2",
|
|
"litellm_params": {"model": "ollama/llama3.2"},
|
|
"model_info": {"max_input_tokens": 4096, "max_output_tokens": 1024},
|
|
},
|
|
{
|
|
"model_name": "undeclared",
|
|
"litellm_params": {"model": "ollama/qwen3"},
|
|
"model_info": {"max_input_tokens": None},
|
|
},
|
|
]
|
|
}
|
|
)
|
|
|
|
|
|
def test_the_catalog_reads_each_models_limits() -> None:
|
|
facts = CATALOG.facts("llama3.2")
|
|
assert (facts.max_input_tokens, facts.max_output_tokens) == (4096, 1024)
|
|
assert CATALOG.facts("undeclared").max_input_tokens is None
|
|
|
|
|
|
def test_a_role_on_a_model_without_a_window_is_refused() -> None:
|
|
CATALOG.require_limits({"driver": "llama3.2", "helper": "claude-haiku-4-5"})
|
|
with pytest.raises(UnknownLimits, match="driver='undeclared'"):
|
|
CATALOG.require_limits({"driver": "undeclared", "helper": "llama3.2"})
|
|
|
|
|
|
def test_a_request_that_cannot_fit_ends_the_run(tmp_path: Path) -> None:
|
|
# Never sent: the driver's run fails and the item waits on the operator,
|
|
# like any failed call, rather than a server getting a prompt it would
|
|
# silently cut.
|
|
store = StateStore(tmp_path / "state.db", oplog=OpLog("t", io.StringIO()))
|
|
store.transcript("t", "operator").append(UserMessage(text="go", source="t"))
|
|
inner = RecordingClient()
|
|
client = BudgetedClient(inner, lambda m: (4096, 1024))
|
|
outcome = Driver(client, "m").run(
|
|
store.transcript("t", "driver"),
|
|
system="s" * 20_000, # a system prompt too big to fit
|
|
tools=[],
|
|
waiting=(),
|
|
workspace=None,
|
|
)
|
|
|
|
assert outcome.kind == "failed"
|
|
assert "fit" in outcome.note
|
|
assert inner.calls == []
|
|
|
|
|
|
# --- tool_choice is only asked of servers that enforce it -------------------
|
|
|
|
FINISH = [{"type": "function", "function": {"name": "finish", "parameters": {}}}]
|
|
|
|
|
|
def test_required_is_asked_only_of_models_declared_to_enforce_it() -> None:
|
|
# Measured: qwen3-14b under "required" cannot stop and rambles to the
|
|
# output ceiling, so a model that is not declared is never asked.
|
|
inner = RecordingClient()
|
|
client = BudgetedClient(
|
|
inner, lambda m: (32768, 8192), enforces_tool_choice=lambda m: m == "local-fast"
|
|
)
|
|
client.complete(transcript(1, 10), FINISH, "local-fast")
|
|
client.complete(transcript(1, 10), FINISH, "local-coder")
|
|
client.complete(transcript(1, 10), None, "local-fast")
|
|
assert [c["require_tool_call"] for c in inner.calls] == [True, False, False]
|
|
|
|
|
|
def test_the_catalog_reads_supports_tool_choice() -> None:
|
|
catalog = ModelCatalog.from_payload(
|
|
{
|
|
"data": [
|
|
{"model_name": "a", "model_info": {"supports_tool_choice": True}},
|
|
{"model_name": "b", "model_info": {"supports_tool_choice": False}},
|
|
{"model_name": "c", "model_info": {}},
|
|
]
|
|
}
|
|
)
|
|
assert [catalog.facts(m).supports_tool_choice for m in "abc"] == [True, False, None]
|
|
|
|
|
|
def test_a_fallback_answer_is_recognised() -> None:
|
|
facts = ModelCatalog.from_payload(
|
|
{
|
|
"data": [
|
|
{
|
|
"model_name": "claude-haiku-4-5",
|
|
"litellm_params": {"model": "anthropic/claude-haiku-4-5"},
|
|
}
|
|
]
|
|
}
|
|
).facts("claude-haiku-4-5")
|
|
assert facts.answered("claude-haiku-4-5") # LiteLLM's alias
|
|
assert facts.answered("claude-haiku-4-5-20251001") # the provider's own name
|
|
assert facts.answered(None) # nothing reported
|
|
assert not facts.answered("claude-sonnet-5") # somebody else answered
|