From 4c2e34f07d69ab3b180b2ffc1e541f1bca458871 Mon Sep 17 00:00:00 2001 From: Jakub Wolniewicz <4850809+frizikk@users.noreply.github.com> Date: Thu, 16 Jul 2026 08:04:36 +0200 Subject: [PATCH] fix(moa): measure advisor guidance before compression --- agent/conversation_loop.py | 35 +++++ agent/moa_loop.py | 180 +++++++++++++++++--------- tests/run_agent/test_moa_loop_mode.py | 114 ++++++++++++++++ 3 files changed, 266 insertions(+), 63 deletions(-) diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index 4cea0eb989a0..df81c20be211 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -706,6 +706,10 @@ def run_conversation( # reused as the final response — not merely because any interim was # streamed. (#65919 review: response-loss blocker) _pending_verification_response_previewed = False + # If pre-API compression fires after MoA advisors have produced guidance, + # retain that ephemeral output and rebase it onto the compacted transcript + # on the next loop iteration. This prevents a second advisor fan-out. + pending_moa_prepared_request = None # Per-turn tally of consecutive successful credential-pool token refreshes, # keyed by (provider, pool-entry-id). A persistent upstream 401 lets @@ -1089,6 +1093,29 @@ def run_conversation( # the OpenAI SDK. Sanitizing here prevents the 3-retry cycle. _sanitize_messages_surrogates(api_messages) + # Build a persistent-MoA request before measuring compression pressure. + # MoA reference output is injected into the aggregator prompt, but it + # is deliberately ephemeral and therefore absent from ``messages``. + # Preparing here makes the pre-API guard measure the exact prompt the + # aggregator will receive; ``create()`` consumes this private prepared + # request later without running the advisors a second time. + _moa_prepared_request = None + if agent.provider == "moa": + _moa_completions = getattr(getattr(agent.client, "chat", None), "completions", None) + if pending_moa_prepared_request is not None: + _rebase_moa_request = getattr(_moa_completions, "rebase_prepared_request", None) + if callable(_rebase_moa_request): + _moa_prepared_request = _rebase_moa_request( + pending_moa_prepared_request, api_messages + ) + pending_moa_prepared_request = None + if _moa_prepared_request is None: + _prepare_moa_request = getattr(_moa_completions, "prepare", None) + if callable(_prepare_moa_request): + _moa_prepared_request = _prepare_moa_request(api_messages) + if _moa_prepared_request is not None: + api_messages = _moa_prepared_request["messages"] + # One image-stripped message estimate feeds both figures. Was: a # str(msg) char walk (re-serialized base64 every call) + a second # messages walk inside estimate_request_tokens_rough. Tools added @@ -1181,6 +1208,8 @@ def run_conversation( and not _compression_cooldown and _compressor.should_compress(request_pressure_tokens) ): + if _moa_prepared_request is not None: + pending_moa_prepared_request = _moa_prepared_request compression_attempts += 1 logger.info( "Pre-API compression: ~%s request tokens >= %s threshold " @@ -1428,6 +1457,12 @@ def run_conversation( if env_var_enabled("HERMES_DUMP_REQUESTS"): agent._dump_api_request_debug(api_kwargs, reason="preflight") + # This object is private to the in-process MoA facade. Add it + # only after middleware, hooks, and debug dumps so none of them + # attempts to serialize it as part of the provider payload. + if _moa_prepared_request is not None and agent.provider == "moa": + api_kwargs["_moa_prepared_request"] = _moa_prepared_request + # Always prefer the streaming path — even without stream # consumers. Streaming gives us fine-grained health # checking (90s stale-stream detection, 60s read timeout) diff --git a/agent/moa_loop.py b/agent/moa_loop.py index d6a0bebf101d..6d1aad9c3750 100644 --- a/agent/moa_loop.py +++ b/agent/moa_loop.py @@ -905,7 +905,114 @@ class MoAChatCompletions: except Exception as exc: # pragma: no cover - display must never break the turn logger.debug("MoA reference_callback failed for %s: %s", event, exc) + def prepare(self, messages: list[dict[str, Any]]) -> dict[str, Any]: + """Run the advisor fan-out and return the exact aggregator request. + + The normal agent loop needs to measure this augmented prompt before its + compression gate. ``create()`` also uses this method for direct callers; + when the loop supplies the returned private object back to ``create()``, + the advisor fan-out is not repeated. + """ + return self.create(messages=messages, _moa_prepare_only=True) + + def rebase_prepared_request( + self, prepared: dict[str, Any], messages: list[dict[str, Any]] + ) -> dict[str, Any]: + """Apply already-generated advisor guidance to a rebuilt API transcript. + + Context compression changes the persisted transcript but not the + ephemeral advisor result. Reusing the guidance avoids a second costly + fan-out while keeping the aggregator request aligned with the compacted + history. + """ + guidance = prepared.get("guidance") + agg_messages = [dict(message) for message in messages] + if guidance: + _attach_reference_guidance(agg_messages, str(guidance)) + return {**prepared, "messages": agg_messages} + + def _call_prepared_aggregator( + self, prepared: dict[str, Any], api_kwargs: dict[str, Any] + ) -> Any: + """Send an already prepared MoA aggregator request exactly once.""" + agg_messages = prepared["messages"] + aggregator = prepared["aggregator"] + aggregator_temperature = prepared["aggregator_temperature"] + if aggregator.get("provider") == "moa": + raise RuntimeError("MoA aggregator cannot be another MoA preset") + agg_kwargs = dict(api_kwargs) + max_tokens: Any = agg_kwargs.get("max_tokens") + tools: Any = agg_kwargs.get("tools") + extra_body: Any = agg_kwargs.get("extra_body") + # Record the exact aggregator INPUT (incl. the injected reference + # context) into the pending trace so a trace captures what the + # aggregator actually saw, not a reconstruction. + if self._pending_trace is not None: + self._pending_trace["aggregator_input_messages"] = agg_messages + self._pending_trace["aggregator_label"] = _slot_label(aggregator) + # The aggregator is the acting model. Resolve its slot to the provider's + # real runtime (base_url/api_key/api_mode) and call it through the same + # request-building path any model uses — so per-model wire-format + # handling (anthropic_messages, max_completion_tokens, fixed/forbidden + # temperature) applies identically to it. MoA imposes no output cap: + # max_tokens is passed through from the caller (normally None → omitted + # → the model's real maximum). The preset's old hardcoded 4096 default + # is gone — it truncated long syntheses. + # When the agent's streaming consumer calls us with stream=True, run the + # references first (above) and then return the aggregator's RAW token + # stream so the acting model's output reaches the user live. The consumer + # reassembles chunks + tool_calls, runs stale-stream detection, and falls + # back to a non-streaming retry on error. The non-streaming path + # (stream=False) is unchanged — no stream/stream_options/timeout are + # forwarded, so its behavior is byte-for-byte identical to before. + stream = bool(api_kwargs.get("stream")) + stream_kwargs: dict[str, Any] = {} + if stream: + stream_kwargs["stream"] = True + stream_kwargs["stream_options"] = ( + api_kwargs.get("stream_options") or {"include_usage": True} + ) + # Forward the consumer's per-request (stream read) timeout so it + # actually governs the aggregator stream, not just call_llm's default. + if api_kwargs.get("timeout") is not None: + stream_kwargs["timeout"] = api_kwargs["timeout"] + _agg_response = call_llm( + task="moa_aggregator", + messages=agg_messages, + temperature=aggregator_temperature, + max_tokens=max_tokens, + tools=tools, + extra_body=extra_body, + # Prepared requests must retain the acting aggregator's reasoning + # policy exactly as the direct create() path does (#64187). + reasoning_config=_aggregator_reasoning_config(aggregator), + **stream_kwargs, + **_slot_runtime(aggregator), + ) + # Non-streaming path (quiet mode / eval / subagents): the aggregator + # output is available inline, so capture it into the pending trace now. + # Streaming path: the aggregator's raw token stream is returned to the + # consumer live and its acting output lands as the turn's assistant + # message; the trace marks it streamed and points there. + if self._pending_trace is not None: + if stream: + self._pending_trace["aggregator_streamed"] = True + self._pending_trace["aggregator_output"] = None + else: + self._pending_trace["aggregator_streamed"] = False + try: + self._pending_trace["aggregator_output"] = _extract_text(_agg_response) + except Exception: # pragma: no cover - defensive + self._pending_trace["aggregator_output"] = None + return _agg_response + def create(self, **api_kwargs: Any) -> Any: + prepared_request = api_kwargs.pop("_moa_prepared_request", None) + if prepared_request is not None: + if not isinstance(prepared_request, dict): + raise TypeError("_moa_prepared_request must be a dict") + return self._call_prepared_aggregator(prepared_request, api_kwargs) + from hermes_cli.config import load_config from hermes_cli.moa_config import resolve_moa_preset @@ -1065,6 +1172,7 @@ class MoAChatCompletions: ref_count=_ref_count, ) + guidance: str | None = None agg_messages = [dict(m) for m in messages] if reference_outputs: joined = "\n\n".join( @@ -1082,69 +1190,15 @@ class MoAChatCompletions: ) _attach_reference_guidance(agg_messages, guidance) - if aggregator.get("provider") == "moa": - raise RuntimeError("MoA aggregator cannot be another MoA preset") - agg_kwargs = dict(api_kwargs) - agg_kwargs["messages"] = agg_messages - # Record the exact aggregator INPUT (incl. the injected reference - # context) into the pending trace so a trace captures what the - # aggregator actually saw, not a reconstruction. - if self._pending_trace is not None: - self._pending_trace["aggregator_input_messages"] = agg_messages - self._pending_trace["aggregator_label"] = _slot_label(aggregator) - # The aggregator is the acting model. Resolve its slot to the provider's - # real runtime (base_url/api_key/api_mode) and call it through the same - # request-building path any model uses — so per-model wire-format - # handling (anthropic_messages, max_completion_tokens, fixed/forbidden - # temperature) applies identically to it. MoA imposes no output cap: - # max_tokens is passed through from the caller (normally None → omitted - # → the model's real maximum). The preset's old hardcoded 4096 default - # is gone — it truncated long syntheses. - # When the agent's streaming consumer calls us with stream=True, run the - # references first (above) and then return the aggregator's RAW token - # stream so the acting model's output reaches the user live. The consumer - # reassembles chunks + tool_calls, runs stale-stream detection, and falls - # back to a non-streaming retry on error. The non-streaming path - # (stream=False) is unchanged — no stream/stream_options/timeout are - # forwarded, so its behavior is byte-for-byte identical to before. - stream = bool(api_kwargs.get("stream")) - stream_kwargs: dict[str, Any] = {} - if stream: - stream_kwargs["stream"] = True - stream_kwargs["stream_options"] = ( - api_kwargs.get("stream_options") or {"include_usage": True} - ) - # Forward the consumer's per-request (stream read) timeout so it - # actually governs the aggregator stream, not just call_llm's default. - if api_kwargs.get("timeout") is not None: - stream_kwargs["timeout"] = api_kwargs["timeout"] - _agg_response = call_llm( - task="moa_aggregator", - messages=agg_messages, - temperature=aggregator_temperature, - max_tokens=agg_kwargs.get("max_tokens"), - tools=agg_kwargs.get("tools"), - extra_body=agg_kwargs.get("extra_body"), - reasoning_config=_aggregator_reasoning_config(aggregator), - **stream_kwargs, - **_slot_runtime(aggregator), - ) - # Non-streaming path (quiet mode / eval / subagents): the aggregator - # output is available inline, so capture it into the pending trace now. - # Streaming path: the aggregator's raw token stream is returned to the - # consumer live and its acting output lands as the turn's assistant - # message; the trace marks it streamed and points there. - if self._pending_trace is not None: - if stream: - self._pending_trace["aggregator_streamed"] = True - self._pending_trace["aggregator_output"] = None - else: - self._pending_trace["aggregator_streamed"] = False - try: - self._pending_trace["aggregator_output"] = _extract_text(_agg_response) - except Exception: # pragma: no cover - defensive - self._pending_trace["aggregator_output"] = None - return _agg_response + prepared_request = { + "messages": agg_messages, + "guidance": guidance, + "aggregator": aggregator, + "aggregator_temperature": aggregator_temperature, + } + if api_kwargs.pop("_moa_prepare_only", False): + return prepared_request + return self._call_prepared_aggregator(prepared_request, api_kwargs) class MoAClient: diff --git a/tests/run_agent/test_moa_loop_mode.py b/tests/run_agent/test_moa_loop_mode.py index 1c1ab8f71eb7..50e3a0087244 100644 --- a/tests/run_agent/test_moa_loop_mode.py +++ b/tests/run_agent/test_moa_loop_mode.py @@ -1265,3 +1265,117 @@ def test_reference_messages_drops_whitespace_only_string_user_turn(): assert view[0] == {"role": "assistant", "content": "a"} assert view[-1] == {"role": "user", "content": "real"} assert all(str(m["content"]).strip() for m in view) + +def test_moa_pre_api_compression_includes_reference_guidance(monkeypatch, tmp_path): + """The aggregator must not receive guidance that pushes it past compression. + + The normal pre-API check sees only the persisted conversation. MoA adds + reference guidance later, inside ``MoAChatCompletions.create()``, so this + regression drives a raw request just below the threshold and makes the + injected guidance cross it. Compression must occur before the aggregator + request and leave the rebuilt request below the threshold. + """ + home = tmp_path / ".hermes" + home.mkdir() + (home / "config.yaml").write_text( + """ +moa: + default_preset: review + presets: + review: + reference_models: + - provider: openrouter + model: advisor + aggregator: + provider: openrouter + model: aggregator +""".strip(), + encoding="utf-8", + ) + monkeypatch.setenv("HERMES_HOME", str(home)) + + events = [] + compression_inputs = [] + aggregator_request_tokens = [] + + def fake_estimate(messages, *args, **kwargs): + rendered = str(messages) + raw_tokens = 80 if "PRE_COMPACTION_HISTORY" in rendered else 20 + guidance_tokens = 40 if "Mixture of Agents reference context" in rendered else 0 + return raw_tokens + guidance_tokens + + def fake_call_llm(**kwargs): + if kwargs["task"] == "moa_reference": + events.append("reference") + return _response("advisor guidance") + events.append("aggregator") + aggregator_request_tokens.append(fake_estimate(kwargs["messages"])) + return _response("aggregator acted") + + monkeypatch.setattr("agent.moa_loop.call_llm", fake_call_llm) + monkeypatch.setattr("agent.turn_context.estimate_request_tokens_rough", fake_estimate) + monkeypatch.setattr("agent.conversation_loop.estimate_request_tokens_rough", fake_estimate) + + agent = AIAgent( + api_key="moa-virtual-provider", + base_url="moa://local", + model="review", + provider="moa", + quiet_mode=True, + skip_context_files=True, + skip_memory=True, + enabled_toolsets=["file"], + max_iterations=3, + ) + compressor = getattr(agent, "context_compressor") + compressor.threshold_tokens = 100 + + def fake_compress(messages, *_args, **_kwargs): + events.append("compress") + compression_inputs.append(messages) + return ([{"role": "user", "content": "SUMMARY"}], "system") + + monkeypatch.setattr(agent, "_compress_context", fake_compress) + + result = agent.run_conversation( + "PRE_COMPACTION_HISTORY", + conversation_history=[{"role": "assistant", "content": "prior response"}], + ) + + assert result["final_response"] == "aggregator acted" + assert events.index("compress") < events.index("aggregator") + assert events.count("reference") == 1 + assert all("Mixture of Agents reference context" not in str(item) for item in compression_inputs) + assert aggregator_request_tokens == [60] + + +def test_prepared_aggregator_preserves_reasoning_config(monkeypatch): + """Prepared MoA requests retain the acting aggregator reasoning policy.""" + from agent import moa_loop + + captured = {} + expected_reasoning = {"enabled": True, "effort": "high"} + + def fake_call_llm(**kwargs): + captured.update(kwargs) + return _response("aggregator acted") + + monkeypatch.setattr(moa_loop, "call_llm", fake_call_llm) + monkeypatch.setattr(moa_loop, "_aggregator_reasoning_config", lambda _slot: expected_reasoning) + monkeypatch.setattr( + moa_loop, + "_slot_runtime", + lambda slot: {"provider": slot["provider"], "model": slot["model"]}, + ) + + facade = moa_loop.MoAChatCompletions("review") + facade._call_prepared_aggregator( + { + "messages": [{"role": "user", "content": "question"}], + "aggregator": {"provider": "openrouter", "model": "aggregator"}, + "aggregator_temperature": None, + }, + {}, + ) + + assert captured["reasoning_config"] == expected_reasoning