mirrors/hermes-agent
1
0
Fork 0
mirror of https://github.com/NousResearch/hermes-agent.git synced 2026-07-24 16:54:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 6

fix(moa): trim reference messages to fit each model's context window

Browse source 
Reference models may have a smaller context window than the aggregator
(e.g. kimi-k2.7-code @ 262K advising a glm-5.2 @ 1M conversation).
Without context-length protection, a reference whose window is exceeded
gets a hard HTTP 400 from the provider, which _run_reference's
try/except silently converts to a [failed: …] note — the MoA turn
silently degrades to fewer references (#60345).

Redesigned implementation of #60387:
- Estimate AFTER the advisory system prompt is prepended, so the
  request that is actually sent is what gets budgeted.
- Reserve output headroom: the preset's reference_max_tokens when set,
  else an 8192-token constant, plus a 10% estimator-error fraction.
- Trim on advisory-view boundaries (text-only user/assistant turns; no
  tool-result frames to orphan), preserving the system prompt, the
  user-first invariant after every pop (never assistant-first), and the
  trailing synthetic user turn.
- Cache get_model_context_length per (provider, model) in a per-fan-out
  dict shared across the worker threads, so a turn resolves each
  window once instead of probing metadata sources
  per-reference-per-iteration (failures are cached too).

Co-authored-by: webtecnica <75556242+webtecnica@users.noreply.github.com>
This commit is contained in:
Teknium 2026-07-23 12:33:10 -07:00
parent 55f3826224
commit 975eb3a365
2 changed files with 323 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

159
agent/moa_loop.py
View file

@ -430,6 +430,7 @@ def _run_reference(
temperature: float | None = None,
max_tokens: int | None = None,
reference_timeout: float | None = None,
context_length_cache: Any = None,
) -> tuple[str, str, Any]:
"""Call one reference model and return ``(label, text, accounting)``.
@ -465,6 +466,20 @@ def _run_reference(
# trimmed view (_reference_messages) already strips the agent's own
# system prompt, so this is the only system message the reference sees.
messages = [{"role": "system", "content": _REFERENCE_SYSTEM_PROMPT}, *ref_messages]
# Trim to fit THIS reference model's context window. Reference models
# may have a smaller window than the aggregator (e.g. kimi-k2.7-code
# @ 262K advising a glm-5.2 @ 1M conversation); without this trim the
# provider returns a hard HTTP 400 which the except below silently
# converts to a [failed: …] note (issue #60345). Estimated AFTER the
# advisory system prompt is prepended so its tokens count against the
# budget too.
messages = _trim_messages_for_reference(
messages,
slot,
runtime,
reserve_output_tokens=max_tokens,
context_length_cache=context_length_cache,
)
# Apply the same Anthropic-style prompt-caching decoration the main
# agent loop applies (system_and_3 breakpoints). The advisory view is
# append-only across iterations (new turns append before the trailing
@ -567,6 +582,144 @@ def _run_reference(
)
# Output-token headroom reserved inside the reference's context window when
# the preset does not cap advisor output (reference_max_tokens=None). Roughly
# one long-form advisory answer; generous enough for thinking models' visible
# output without starving the input budget.
_REFERENCE_DEFAULT_OUTPUT_RESERVE = 8192
# Additional estimation slack: estimate_messages_tokens_rough is a rough
# chars/4 heuristic and providers tokenize less favorably on code/JSON-heavy
# transcripts, so keep a safety fraction of the window unbudgeted.
_REFERENCE_TRIM_SAFETY_FRACTION = 0.10
def _trim_messages_for_reference(
messages: list[dict[str, Any]],
slot: dict[str, str],
runtime: dict[str, Any],
*,
reserve_output_tokens: int | None = None,
context_length_cache: Any = None,
) -> list[dict[str, Any]]:
"""Trim an advisory request to fit within a reference model's context window.
Reference models may have a smaller context window than the aggregator or
the main conversation. Without this trim, a reference whose window is
exceeded gets a hard HTTP 400 from the provider, which ``_run_reference``'s
try/except silently converts to a ``[failed: …]`` note — the MoA turn
silently degrades to fewer references (issue #60345).
``messages`` is the FULL request as it will be sent — the advisory system
prompt already prepended — so the estimate covers everything the provider
will count. The budget reserves ``reserve_output_tokens`` (the preset's
``reference_max_tokens`` when set, else a sane constant) for the model's
response plus a safety fraction for estimator error.
Trimming drops the OLDEST conversation frames (right after the system
prompt) and preserves two invariants of the advisory view, which is
text-only user/assistant turns (``_reference_messages`` renders tool
calls/results inline, so there are no tool-result frames to orphan):
- the system prompt (index 0) is always kept;
- the first non-system message stays ``user``-first — after each pop,
any now-leading assistant turns are popped too, so no provider ever
sees an assistant-first conversation;
- the trailing user turn (the synthetic judge-the-state marker) and at
least one preceding turn are always kept, even if still over budget —
a too-long-but-recent view beats an empty request.
``context_length_cache`` is an optional per-turn dict keyed by
``(provider, model)`` so one fan-out (and every iteration reusing the
cache) resolves each model's window at most once instead of re-probing
metadata sources per-reference-per-iteration. When the window cannot be
resolved, messages are returned unchanged.
"""
if not messages:
return messages
from agent.model_metadata import (
estimate_messages_tokens_rough,
get_model_context_length,
)
model = str(slot.get("model") or "")
provider = str(runtime.get("provider") or slot.get("provider") or "")
if not model:
return messages
cache_key = (provider, model)
context_length: int | None = None
if isinstance(context_length_cache, dict) and cache_key in context_length_cache:
context_length = context_length_cache[cache_key]
else:
try:
context_length = get_model_context_length(
model=model,
base_url=str(runtime.get("base_url") or ""),
api_key=str(runtime.get("api_key") or ""),
provider=provider,
)
except Exception:
logger.debug(
"MoA reference context-length resolution failed for %s",
_slot_label(slot),
)
context_length = None
if isinstance(context_length_cache, dict):
# Cache failures too (as None) — a flaky metadata source should
# not be re-probed for every reference of every iteration.
context_length_cache[cache_key] = context_length
if not isinstance(context_length, int) or context_length <= 0:
return messages
reserve = (
int(reserve_output_tokens)
if isinstance(reserve_output_tokens, int) and reserve_output_tokens > 0
else _REFERENCE_DEFAULT_OUTPUT_RESERVE
)
budget = int(context_length * (1.0 - _REFERENCE_TRIM_SAFETY_FRACTION)) - reserve
if budget <= 0:
return messages
estimated = estimate_messages_tokens_rough(messages)
if estimated <= budget:
return messages
has_system = bool(messages) and messages[0].get("role") == "system"
head = [messages[0]] if has_system else []
body = list(messages[1:] if has_system else messages)
# Keep the trailing user turn plus at least one preceding turn.
while len(body) > 2 and estimate_messages_tokens_rough(head + body) > budget:
body.pop(0)
# Preserve the user-first invariant: never leave the advisory
# conversation starting on an assistant turn after a pop.
while len(body) > 2 and body[0].get("role") == "assistant":
body.pop(0)
# The loop can stop with two frames left where the first is an
# assistant turn — enforce user-first even then (a lone trailing user
# turn is a valid request; an assistant-first one is not).
while len(body) > 1 and body[0].get("role") == "assistant":
body.pop(0)
trimmed = head + body
dropped = len(messages) - len(trimmed)
if dropped:
logger.info(
"MoA reference %s: estimated %d tokens exceeds budget %d "
"(window %d, output reserve %d); dropped %d oldest message(s).",
_slot_label(slot),
estimated,
budget,
context_length,
reserve,
dropped,
)
return trimmed
_REFERENCE_POLL_INTERVAL_S = 5.0
# Sentinel text for a reference slot whose wait was aborted by a user
@ -637,6 +790,11 @@ def _run_references_parallel(
completed = 0
executor = ThreadPoolExecutor(max_workers=workers)
interrupted = False
# Per-fan-out context-length cache shared by every reference worker, so
# duplicate (provider, model) slots resolve their window once per turn
# instead of re-probing metadata sources per reference (dict get/set is
# GIL-atomic; a rare duplicate probe on a first-use race is harmless).
_ctx_len_cache: dict[tuple[str, str], int | None] = {}
try:
for idx, slot in enumerate(reference_models):
if slot.get("provider") == "moa":
@ -654,6 +812,7 @@ def _run_references_parallel(
temperature=temperature,
max_tokens=max_tokens,
reference_timeout=reference_timeout,
context_length_cache=_ctx_len_cache,
)
] = idx

164
tests/run_agent/test_moa_loop_mode.py
View file

@ -2478,3 +2478,167 @@ moa:
usage2, cost2 = facade.consume_reference_usage()
assert usage2.input_tokens == 0
assert cost2 is None
class _CountingCtxLen:
"""Stub for get_model_context_length that counts resolutions."""
def __init__(self, value):
self.value = value
self.calls = 0
def __call__(self, **kwargs):
self.calls += 1
if isinstance(self.value, Exception):
raise self.value
return self.value
def _trim(messages, *, window=1000, reserve=None, cache=None, counting=None,
monkeypatch=None):
from agent import model_metadata, moa_loop
stub = counting or _CountingCtxLen(window)
monkeypatch.setattr(model_metadata, "get_model_context_length", stub)
return moa_loop._trim_messages_for_reference(
messages,
{"provider": "openrouter", "model": "small-window"},
{"provider": "openrouter", "model": "small-window"},
reserve_output_tokens=reserve,
context_length_cache=cache,
)
def _advisory_view(n_pairs, chunk="x" * 400):
"""A text-only advisory view: system + n user/assistant pairs + trailing user."""
msgs = [{"role": "system", "content": "advisory system prompt"}]
for i in range(n_pairs):
msgs.append({"role": "user", "content": f"u{i} {chunk}"})
msgs.append({"role": "assistant", "content": f"a{i} {chunk}"})
msgs.append({"role": "user", "content": "judge the state above"})
return msgs
def test_reference_trim_untouched_when_within_window(monkeypatch):
msgs = _advisory_view(2)
out = _trim(list(msgs), window=10_000_000, monkeypatch=monkeypatch)
assert out == msgs
def test_reference_trim_drops_oldest_and_keeps_invariants(monkeypatch):
from agent.model_metadata import estimate_messages_tokens_rough
msgs = _advisory_view(30)
out = _trim(list(msgs), window=4000, reserve=100, monkeypatch=monkeypatch)
# Something was dropped and it fits the (window*0.9 - reserve) budget…
assert len(out) < len(msgs)
# System prompt survives at index 0.
assert out[0]["role"] == "system"
# User-first after the system prompt — never assistant-first.
assert out[1]["role"] == "user"
# Trailing synthetic user turn survives.
assert out[-1] == msgs[-1]
# Oldest frames were the ones dropped: the kept body is a contiguous
# suffix of the original body.
assert out[1:] == msgs[len(msgs) - len(out) + 1:]
def test_reference_trim_estimates_after_system_prompt(monkeypatch):
"""The advisory system prompt counts against the budget: a view that fits
without it but not with it must still be trimmed."""
from agent import model_metadata, moa_loop
msgs = _advisory_view(6)
body_tokens = model_metadata.estimate_messages_tokens_rough(msgs[1:])
total_tokens = model_metadata.estimate_messages_tokens_rough(msgs)
assert total_tokens > body_tokens
# Pick a window where (body fits) but (system+body does not):
# budget = window*0.9 - reserve(100) must sit between the two.
window = int((body_tokens + (total_tokens - body_tokens) / 2 + 100) / 0.9)
out = _trim(list(msgs), window=window, reserve=100, monkeypatch=monkeypatch)
assert len(out) < len(msgs)
assert out[0]["role"] == "system"
def test_reference_trim_reserves_output_tokens(monkeypatch):
"""With a huge output reserve the budget shrinks and forces a trim that
a reserve-less estimate would not need."""
msgs = _advisory_view(10)
from agent.model_metadata import estimate_messages_tokens_rough
total = estimate_messages_tokens_rough(msgs)
window = int((total + 500) / 0.9) # fits easily with a small reserve
out_small = _trim(list(msgs), window=window, reserve=100, monkeypatch=monkeypatch)
assert out_small == msgs
out_big = _trim(list(msgs), window=window, reserve=total, monkeypatch=monkeypatch)
assert len(out_big) < len(msgs)
def test_reference_trim_unresolvable_window_is_a_noop(monkeypatch):
msgs = _advisory_view(3)
stub = _CountingCtxLen(RuntimeError("metadata down"))
out = _trim(list(msgs), counting=stub, monkeypatch=monkeypatch)
assert out == msgs
def test_reference_trim_context_length_cache_hits_once(monkeypatch):
"""A shared per-turn cache resolves each (provider, model) window once."""
cache = {}
stub = _CountingCtxLen(10_000_000)
msgs = _advisory_view(2)
for _ in range(4):
_trim(list(msgs), cache=cache, counting=stub, monkeypatch=monkeypatch)
assert stub.calls == 1
assert cache == {("openrouter", "small-window"): 10_000_000}
def test_reference_trim_caches_resolution_failures(monkeypatch):
"""A failing metadata source is probed once, not per reference call."""
cache = {}
stub = _CountingCtxLen(RuntimeError("metadata down"))
msgs = _advisory_view(2)
for _ in range(3):
out = _trim(list(msgs), cache=cache, counting=stub, monkeypatch=monkeypatch)
assert out == msgs
assert stub.calls == 1
assert cache == {("openrouter", "small-window"): None}
def test_run_reference_trims_oversized_view_before_calling(monkeypatch):
"""End-to-end: _run_reference sends a trimmed request for a small-window
model instead of letting the provider 400."""
from agent import model_metadata, moa_loop
sent = {}
def fake_call_llm(**kwargs):
sent.update(kwargs)
return _response_with_usage("advice")
monkeypatch.setattr(moa_loop, "call_llm", fake_call_llm)
monkeypatch.setattr(
moa_loop,
"_slot_runtime",
lambda slot: {"provider": "openrouter", "model": slot["model"]},
)
monkeypatch.setattr(
model_metadata, "get_model_context_length", lambda **k: 3000
)
view = _advisory_view(40)[1:] # advisory view has no system prompt
label, text, _acct = moa_loop._run_reference(
{"provider": "openrouter", "model": "small-window"},
view,
max_tokens=200,
)
assert text == "advice"
sent_messages = sent["messages"]
# System prompt was prepended and survived the trim.
assert sent_messages[0]["role"] == "system"
# The request actually shrank.
assert len(sent_messages) < len(view) + 1
# And ends with the synthetic trailing user turn.
assert sent_messages[-1]["role"] == "user"