From 3e496e220b16d0eb0111aededb0f6c5b2074c8ee Mon Sep 17 00:00:00 2001 From: Serge Panev Date: Fri, 4 Sep 2026 20:25:17 -0700 Subject: [PATCH 1/2] [GDPVal] Fix long-context budgeting and compaction Signed-off-by: Serge Panev --- benchmarks/gdpval/config.yaml | 8 + responses_api_agents/stirrup_agent/README.md | 22 +- responses_api_agents/stirrup_agent/app.py | 35 +- .../stirrup_agent/nemo_agent.py | 164 ++++++- .../stirrup_agent/nemo_client.py | 269 ++++++++++-- .../stirrup_agent/tests/test_app.py | 17 + .../tests/test_nemo_agent_compaction.py | 303 +++++++++++++ .../stirrup_agent/tests/test_nemo_client.py | 401 +++++++++++++++++- 8 files changed, 1179 insertions(+), 40 deletions(-) create mode 100644 responses_api_agents/stirrup_agent/tests/test_nemo_agent_compaction.py diff --git a/benchmarks/gdpval/config.yaml b/benchmarks/gdpval/config.yaml index 04030acb9a9..85275b49bf6 100644 --- a/benchmarks/gdpval/config.yaml +++ b/benchmarks/gdpval/config.yaml @@ -103,6 +103,14 @@ gdpval_stirrup_agent: agent_max_turns: 250 concurrency: 32 temperature: 1.0 + # GDPVal tasks routinely need multi-thousand-token scripts. A 1,024-token + # completion cannot form a useful tool call and can create a sticky + # code_exec({}) loop after long reasoning turns. + min_completion_tokens: 8192 + # Reject the observed near-empty compaction outputs that erase progress + # on long GDPVal sessions. Generic Stirrup agents accept any non-empty + # summary for backward compatibility. + min_compaction_summary_words: 50 system_prompt_template: ${oc.env:SYSTEM_PROMPT_TEMPLATE,null} user_prompt_template: ${oc.env:USER_PROMPT_TEMPLATE,null} gdpval_container_path: ${oc.env:GDPVAL_CONTAINER_PATH,null} diff --git a/responses_api_agents/stirrup_agent/README.md b/responses_api_agents/stirrup_agent/README.md index 68b71fe0603..ce8dcd2c22c 100644 --- a/responses_api_agents/stirrup_agent/README.md +++ b/responses_api_agents/stirrup_agent/README.md @@ -131,6 +131,9 @@ The agent reads its Hydra config at `configs/stirrup_agent.yaml`. Notable keys: | `rerun_incomplete` | `false` | If true, skip the rollout for tasks that already **finished** (a finish marker is cached) and only re-run the ones that didn't (see [Task Re-run Mode](#task-re-run-mode)). Composes with `execute_only` and `judge_only`. | | `model_id` | `null` | HF model id or local path used to load a tokenizer for dynamic output sizing. | | `completion_token_buffer` | `1000` | Safety margin (in tokens) reserved when sizing `max_completion_tokens` per call. | +| `min_completion_tokens` | `1024` | Target useful per-call completion floor. The hard cap always applies; estimated remaining context is also a strict bound when the tokenizer successfully renders the complete prompt. Approximate fallbacks preserve this floor even when their estimate exceeds context. GDPVal raises it to `8192`. | +| `prompt_estimator_truncate_history_thinking` | `null` | Optional estimator-only switch for templates that omit completed historical reasoning. It is never sent to the provider. | +| `min_compaction_summary_words` | `1` | Minimum summary length accepted during context compaction. GDPVal uses `50`; invalid summaries retry up to three times and never replace the live history. | Env vars honored: `TAVILY_API_KEY`, `HF_TOKEN`, `OPENAI_API_KEY`. @@ -155,13 +158,28 @@ The wrapper ships a `DynamicMaxTokensChatCompletionsClient` Set `model_id` to the same checkpoint (or HF id) you are serving via vLLM and the tokeniser match is exact. Leave it unset and the client -falls back to a conservative character-count estimate — slower to -allocate completion budget but always safe. `completion_token_buffer` +falls back to a character-count estimate. When +`prompt_estimator_truncate_history_thinking=true`, that fallback mirrors +Nemotron's prior-assistant transformation without changing the retained +trajectory or provider request. `completion_token_buffer` absorbs the residual gap between our estimate and the exact prompt the server renders (chat-template wrappers, tool-schema injection). The default `1000` works in practice; raise it (e.g. 2000–5000) if you see sporadic HTTP 400 responses at the vLLM proxy. +When the loaded tokenizer successfully renders the complete prompt, estimated +remaining context is a strict output bound and a prompt with no remaining room +fails before dispatch. If full rendering is unsupported, JSON and character +estimates can substantially overcount retained reasoning, so the client still +dispatches the configured completion floor; the hard completion cap applies in +both modes. + +GDPVal uses an `8192`-token minimum because a 1024-token truncation cannot +reliably produce a complete script or tool call. Generic Stirrup agents keep +the historical 1024-token default. Launchers serving a local checkpoint should +still pass that checkpoint as `model_id`; the fallback is a resilience path, +not a replacement for template-exact tokenization. + ## Advanced Features NeMo Gym splits a GDPVal evaluation into two halves — *executing* the task diff --git a/responses_api_agents/stirrup_agent/app.py b/responses_api_agents/stirrup_agent/app.py index 9937b103206..bc1f3ced3e7 100644 --- a/responses_api_agents/stirrup_agent/app.py +++ b/responses_api_agents/stirrup_agent/app.py @@ -489,7 +489,10 @@ async def _run_stirrup_agent( top_p: float = 0.95, enable_thinking: bool = True, max_completion_tokens_cap: int = 64000, + min_completion_tokens: int = 1024, + prompt_estimator_truncate_history_thinking: Optional[bool] = None, truncation_recovery: bool = True, + min_compaction_summary_words: int = 1, tavily_api_key: Optional[Union[str, List[str]]] = None, tavily_max_sweeps: int = 1, ) -> Dict[str, Any]: @@ -557,6 +560,8 @@ async def run_command(self, cmd: str, *, timeout: int = SHELL_TIMEOUT) -> Comman top_p=top_p, enable_thinking=enable_thinking, max_completion_tokens_cap=max_completion_tokens_cap, + min_completion_tokens=min_completion_tokens, + prompt_estimator_truncate_history_thinking=prompt_estimator_truncate_history_thinking, truncation_recovery=truncation_recovery, ) @@ -609,6 +614,7 @@ async def run_command(self, cmd: str, *, timeout: int = SHELL_TIMEOUT) -> Comman "tools": tools, "tool_response_as_user": True, "skip_input_file_listing": is_gdpval, + "min_compaction_summary_words": min_compaction_summary_words, } if system_prompt: agent_kwargs["system_prompt"] = system_prompt @@ -969,8 +975,8 @@ class StirrupAgentWrapperConfig(BaseResponsesAPIAgentConfig): default=None, description="HuggingFace model ID (or local checkpoint path) used to load a tokenizer " "for dynamic max_completion_tokens sizing. E.g. 'Qwen/Qwen3-Coder-30B-A3B-Instruct'. " - "When None, a character-count fallback is used (conservative, but slightly over-allocates " - "input tokens). See ``nemo_client.DynamicMaxTokensChatCompletionsClient``.", + "When None, an approximate character-count fallback is used; it can substantially " + "overcount retained reasoning. See ``nemo_client.DynamicMaxTokensChatCompletionsClient``.", ) completion_token_buffer: int = Field( default=1000, @@ -990,10 +996,25 @@ class StirrupAgentWrapperConfig(BaseResponsesAPIAgentConfig): ) max_completion_tokens_cap: int = Field( default=64000, + ge=1, description="Hard ceiling on per-call ``max_completion_tokens``. Dynamic sizing computes " "context_window - input_tokens - completion_token_buffer, then caps to this value. " "Set to match the training-side response-length budget for RL.", ) + min_completion_tokens: int = Field( + default=1024, + ge=1, + description="Target minimum per-call ``max_completion_tokens``. The hard cap always applies; " + "estimated remaining context is also a strict bound when the loaded tokenizer successfully " + "renders the complete prompt. The " + "approximate fallback preserves this floor even when its estimate exceeds context. " + "Long-horizon benchmarks can opt into a larger usable floor.", + ) + prompt_estimator_truncate_history_thinking: Optional[bool] = Field( + default=None, + description="Optional prompt-estimator setting for checkpoints whose template omits historical " + "reasoning before the last user turn. This is never forwarded to the model request.", + ) truncation_recovery: bool = Field( default=True, description="When a model call spends its entire completion budget without emitting a " @@ -1004,6 +1025,13 @@ class StirrupAgentWrapperConfig(BaseResponsesAPIAgentConfig): "deliverable writes (median 34.6k tokens). The instruction is sent to the server but not " "recorded in the trajectory; set False for RL rollouts that require an unsteered policy.", ) + min_compaction_summary_words: int = Field( + default=1, + ge=1, + description="Minimum whitespace-delimited word count accepted from context compaction. " + "The generic default rejects only empty output; long-horizon benchmarks can require a " + "more complete summary.", + ) tavily_api_key: Optional[Union[str, List[str]]] = Field( default=None, description="Tavily API key(s) for the ``web_search`` / ``fetch_web_page`` tools. " @@ -1197,7 +1225,10 @@ async def responses( "top_p": getattr(body, "top_p", None) or self.config.top_p, "enable_thinking": self.config.enable_thinking, "max_completion_tokens_cap": self.config.max_completion_tokens_cap, + "min_completion_tokens": self.config.min_completion_tokens, + "prompt_estimator_truncate_history_thinking": (self.config.prompt_estimator_truncate_history_thinking), "truncation_recovery": self.config.truncation_recovery, + "min_compaction_summary_words": self.config.min_compaction_summary_words, "tavily_api_key": self.config.tavily_api_key, "tavily_max_sweeps": self.config.tavily_max_sweeps, } diff --git a/responses_api_agents/stirrup_agent/nemo_agent.py b/responses_api_agents/stirrup_agent/nemo_agent.py index 20d83562a8e..e432ce54204 100644 --- a/responses_api_agents/stirrup_agent/nemo_agent.py +++ b/responses_api_agents/stirrup_agent/nemo_agent.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""NeMoAgent — Stirrup ``Agent`` subclass with two behavioural overrides. +"""NeMoAgent — Stirrup ``Agent`` subclass with Gym-specific behaviour. 1. **tool_response_as_user** — Return a ``UserMessage`` (not ``ToolMessage``) from ``run_tool()`` so the conversation history presents tool results @@ -23,6 +23,10 @@ injects into the system prompt; useful when a task prompt already lists its own reference files (e.g. GDPVal). +3. **safe context compaction** — Removes inline reasoning from generated + summaries and retries empty or degenerate summaries before replacing the + live history. + To preserve tool-call metadata that ``Agent.step()`` reads immediately after ``run_tool()`` returns (``.name``, ``.success``, ``.tool_call_id``), the conversion uses :class:`NeMoUserMessage` — a ``UserMessage`` subclass with @@ -32,12 +36,69 @@ from __future__ import annotations +import logging +import re +from itertools import takewhile from typing import Any, Optional from pydantic import ConfigDict from stirrup import Agent from stirrup.core.agent import SessionAgent -from stirrup.core.models import ToolCall, ToolMessage, UserMessage +from stirrup.core.exceptions import ContextOverflowError +from stirrup.core.models import AssistantMessage, ChatMessage, SummaryMessage, ToolCall, ToolMessage, UserMessage +from stirrup.prompts import MESSAGE_SUMMARIZER, MESSAGE_SUMMARIZER_BRIDGE_TEMPLATE + + +LOGGER = logging.getLogger(__name__) + +_MAX_COMPACTION_SUMMARY_ATTEMPTS = 3 +_LEADING_THINK_OPEN_RE = re.compile(r"\A\s*", re.IGNORECASE) +_THINK_CLOSE_RE = re.compile(r"", re.IGNORECASE) + + +def _sanitize_compaction_summary(content: Any) -> str: + """Remove leading inline-thinking blocks from a generated summary. + + Only a leading block is special: a later literal ```` may be part + of task content and must survive. An unclosed leading block has no safe + answer portion, so it is treated as empty and retried. + """ + if not isinstance(content, str): + return "" + + sanitized = content.strip() + while match := _LEADING_THINK_OPEN_RE.match(sanitized): + closing = _THINK_CLOSE_RE.search(sanitized, match.end()) + if closing is None: + return "" + sanitized = sanitized[closing.end() :].lstrip() + return sanitized.strip() + + +def _is_meaningful_compaction_summary(summary: str, *, min_words: int = 1) -> bool: + """Return whether a summary is large enough to replace a full context.""" + return len(summary.split()) >= min_words + + +def _summary_text_only_instruction(min_words: int) -> str: + minimum = ( + f"The summary must contain at least {min_words} words." + if min_words > 1 + else "The summary must contain useful, non-empty text." + ) + return ( + "Return a concise but comprehensive summary as text only. " + "Do not call tools. Do not include analysis, reasoning, or tags. " + f"{minimum}" + ) + + +def _summary_rejection_reason(content: Any, sanitized: str, min_words: int) -> str: + if not isinstance(content, str): + return f"non-text content ({type(content).__name__})" + if not sanitized: + return "empty after removing leading reasoning" + return f"{len(sanitized.split())} words, below configured minimum {min_words}" class NeMoUserMessage(UserMessage): @@ -76,11 +137,15 @@ def __init__( *, tool_response_as_user: bool = False, skip_input_file_listing: bool = False, + min_compaction_summary_words: int = 1, **kwargs: Any, ) -> None: + if min_compaction_summary_words < 1: + raise ValueError("min_compaction_summary_words must be at least 1") super().__init__(**kwargs) self._tool_response_as_user = tool_response_as_user self._skip_input_file_listing = skip_input_file_listing + self._min_compaction_summary_words = min_compaction_summary_words def _build_system_prompt(self) -> str: """Override to optionally skip the input file listing.""" @@ -103,6 +168,101 @@ def _build_system_prompt(self) -> str: return result + async def summarize_messages( + self, + messages: list[ChatMessage], + run_metadata_by_turn: dict[str, dict[str, list[Any]]], + ) -> tuple[list[ChatMessage], list[ChatMessage]]: + """Compact history without replaying reasoning or erasing progress. + + Stirrup 0.1 inserts ``summary.content`` into the next context verbatim. + Some reasoning-model endpoints return a leading ```` block in + that field, including think-only responses. Sanitize at the adapter + boundary and retry bad summaries before replacing any history. + """ + current_messages = messages + while True: + try: + task_context: list[ChatMessage] = list( + takewhile( + lambda message: not isinstance(message, (AssistantMessage, SummaryMessage)), + current_messages, + ) + ) + + tool_docs = "\n".join(f"- {tool.name}: {tool.description}" for tool in self._active_tools.values()) + strict_prompt = ( + f"{MESSAGE_SUMMARIZER}\n\n{_summary_text_only_instruction(self._min_compaction_summary_words)}" + ) + no_tools_prompt = f"{strict_prompt}\n\nTools are disabled for this response." + ( + f" Tools used earlier were:\n{tool_docs}" if tool_docs else "" + ) + attempts = ( + (MESSAGE_SUMMARIZER, self._active_tools), + (strict_prompt, self._active_tools), + (no_tools_prompt, {}), + ) + + summary_content = "" + for attempt, (prompt, tools) in enumerate(attempts, start=1): + # Compaction uses the policy client internally, but its + # one-shot truncation recovery belongs to policy turns. + # Do not let a summary consume an armed recovery or arm a + # recovery that leaks into the next policy turn. + sentinel = object() + saved_recovery = getattr(self._client, "_recover_from_truncation", sentinel) + saved_overruns = getattr(self._client, "_truncation_overruns", sentinel) + if saved_recovery is not sentinel: + self._client._recover_from_truncation = False + try: + summary = await self._client.generate( + [*current_messages, UserMessage(content=prompt)], + tools, + ) + finally: + if saved_recovery is not sentinel: + self._client._recover_from_truncation = saved_recovery + if saved_overruns is not sentinel: + self._client._truncation_overruns = saved_overruns + summary_content = _sanitize_compaction_summary(summary.content) + if _is_meaningful_compaction_summary( + summary_content, + min_words=self._min_compaction_summary_words, + ): + break + usage = getattr(summary, "token_usage", None) + LOGGER.warning( + "Compaction summary attempt %d/%d rejected: %s (input=%s answer=%s reasoning=%s); %s", + attempt, + _MAX_COMPACTION_SUMMARY_ATTEMPTS, + _summary_rejection_reason( + summary.content, + summary_content, + self._min_compaction_summary_words, + ), + getattr(usage, "input", None), + getattr(usage, "answer", None), + getattr(usage, "reasoning", None), + "retrying" if attempt < _MAX_COMPACTION_SUMMARY_ATTEMPTS else "no attempts remain", + ) + else: + raise RuntimeError( + "Compaction summarizer produced no usable text after " + f"{_MAX_COMPACTION_SUMMARY_ATTEMPTS} attempts; refusing to erase context" + ) + + summary_bridge_prompt = MESSAGE_SUMMARIZER_BRIDGE_TEMPLATE.format(summary=summary_content) + summary_bridge = SummaryMessage(content=summary_bridge_prompt) + acknowledgement_msg = UserMessage(content="Got it, thanks!") + + self._logger.context_summarization_complete(summary_content, summary_bridge_prompt) + return current_messages, [*task_context, summary_bridge, acknowledgement_msg] + except ContextOverflowError: + # Preserve Stirrup's recovery contract: discard the latest + # completed turn, including its per-turn metadata, and retry. + current_messages, dropped_turn_id = self._unwind_context_overflow(current_messages) + run_metadata_by_turn.pop(dropped_turn_id, None) + async def run_tool(self, tool_call: ToolCall, run_metadata: dict[str, list[Any]]) -> ToolMessage: """Run a tool and optionally return a ``NeMoUserMessage`` instead of ``ToolMessage``. diff --git a/responses_api_agents/stirrup_agent/nemo_client.py b/responses_api_agents/stirrup_agent/nemo_client.py index 790ec2df482..1749044d2ba 100644 --- a/responses_api_agents/stirrup_agent/nemo_client.py +++ b/responses_api_agents/stirrup_agent/nemo_client.py @@ -32,18 +32,22 @@ context_window − tokenized(messages) − completion_token_buffer -clamped to a minimum of ``_MIN_COMPLETION_TOKENS``. On the response +raised toward a configurable minimum (historically 1,024 tokens), subject to +the hard cap. Estimated remaining context is also a strict bound when the +tokenizer successfully renders the complete prompt; approximate fallbacks keep +the configured floor because they can substantially overcount retained reasoning. On the response side, we replicate Stirrup parsing but do *not* raise on ``finish_reason=length`` — the agent loop will either terminate when the model invokes the ``finish`` tool or exhaust ``max_turns``, yielding a clean timeout instead of a crash. ``model_id`` selects the HuggingFace tokenizer (or local checkpoint path). -When unset, a conservative character-count fallback is used. +When unset, an approximate character-count fallback is used. """ from __future__ import annotations +import json import logging from time import perf_counter from typing import Any, Optional @@ -52,6 +56,7 @@ from pydantic import ValidationError as _PydanticValidationError from stirrup.clients.chat_completions_client import ChatCompletionsClient from stirrup.clients.utils import to_openai_tools +from stirrup.core.exceptions import ContextOverflowError from stirrup.core.models import ( AssistantMessage, ChatMessage, @@ -153,8 +158,9 @@ def _install_coercing_finish_tool() -> None: _install_coercing_finish_tool() -# Floor for per-call max_completion_tokens. Below this the model basically -# cannot produce a useful answer — treat as a hard minimum. +# Target floor for per-call max_completion_tokens. Below this the model usually +# cannot produce a useful answer. The hard cap always takes precedence; exact +# tokenizer estimates also enforce remaining context as a strict bound. _MIN_COMPLETION_TOKENS = 1024 # Hard cap on per-call max_completion_tokens. Oversized completion budgets @@ -231,15 +237,23 @@ def __init__( top_p: float = 0.95, enable_thinking: bool = True, max_completion_tokens_cap: int = _DEFAULT_MAX_COMPLETION_TOKENS_CAP, + min_completion_tokens: int = _MIN_COMPLETION_TOKENS, + prompt_estimator_truncate_history_thinking: Optional[bool] = None, truncation_recovery: bool = True, **kwargs: Any, ) -> None: super().__init__(*args, **kwargs) + if max_completion_tokens_cap < 1: + raise ValueError("max_completion_tokens_cap must be at least 1") + if min_completion_tokens < 1: + raise ValueError("min_completion_tokens must be at least 1") self._completion_token_buffer = completion_token_buffer self._temperature = temperature self._top_p = top_p self._enable_thinking = enable_thinking self._max_completion_tokens_cap = max_completion_tokens_cap + self._min_completion_tokens = min_completion_tokens + self._prompt_estimator_truncate_history_thinking = prompt_estimator_truncate_history_thinking self._truncation_recovery = truncation_recovery # Set when the previous call exhausted its completion budget without # emitting a tool call; consumed by the very next generate(). @@ -256,11 +270,141 @@ def __init__( # Token counting # ------------------------------------------------------------------ + @staticmethod + def _truncate_prior_assistant_text_for_estimate(content: str, *, has_tool_calls: bool) -> str: + """Mirror Nemotron's historical-assistant rendering for one text value.""" + # The template first inserts an empty reasoning block when neither tag + # is present. This matters for both its normal and tool-call branches. + if "" not in content and "" not in content: + content = f"{content}" + + if has_tool_calls: + if not content.strip(): + return "" + if "" in content: + content = content.rsplit("", 1)[-1] + elif "" in content: + content = content.split("", 1)[0] + return f"{content}" + + if "" in content and "" in content: + content = f"{content.rsplit('', 1)[-1]}" + return content.strip() + + def _messages_for_estimator_count( + self, + messages: list[NeMoGymChatCompletionMessageParam], + ) -> list[NeMoGymChatCompletionMessageParam]: + """Return the history shape rendered by Nemotron's truncation mode. + + This is estimator-only: it never changes the retained trajectory or + request. Tool calls and every other serialized field stay intact. + """ + if self._prompt_estimator_truncate_history_thinking is not True: + return messages + + last_user_index = max( + (index for index, message in enumerate(messages) if message.get("role") == "user"), + default=-1, + ) + if last_user_index < 0: + return messages + + counted_messages: list[NeMoGymChatCompletionMessageParam] = [] + for index, message in enumerate(messages): + if index >= last_user_index or message.get("role") != "assistant": + counted_messages.append(message) + continue + + content = message.get("content") + estimated_content: Any = content + has_tool_calls = bool(message.get("tool_calls")) + if isinstance(content, str): + estimated_content = self._truncate_prior_assistant_text_for_estimate( + content, + has_tool_calls=has_tool_calls, + ) + elif isinstance(content, list): + estimated_parts = [] + changed = False + for part in content: + if isinstance(part, dict) and isinstance(part.get("text"), str): + estimated_text = self._truncate_prior_assistant_text_for_estimate( + part["text"], + has_tool_calls=has_tool_calls, + ) + if estimated_text != part["text"]: + part = {**part, "text": estimated_text} + changed = True + estimated_parts.append(part) + if changed: + estimated_content = estimated_parts + + if estimated_content is content or estimated_content == content: + counted_messages.append(message) + else: + counted_messages.append({**message, "content": estimated_content}) + return counted_messages + + def _messages_for_template_count( + self, + messages: list[NeMoGymChatCompletionMessageParam], + ) -> list[NeMoGymChatCompletionMessageParam]: + """Mirror vLLM's tool-argument decoding before template rendering. + + OpenAI history carries ``function.arguments`` as a JSON string, while + Nemotron's Jinja template iterates it as a mapping. vLLM decodes that + field before rendering; HuggingFace ``apply_chat_template`` does not. + Keep this as a fallback template-input variant so templates that + natively accept OpenAI's string representation retain their existing + behavior. The provider payload is never mutated. + """ + normalized_messages: list[NeMoGymChatCompletionMessageParam] = [] + messages_changed = False + for message in messages: + tool_calls = message.get("tool_calls") + if not isinstance(tool_calls, list): + normalized_messages.append(message) + continue + + normalized_calls = [] + changed = False + for tool_call in tool_calls: + if not isinstance(tool_call, dict): + normalized_calls.append(tool_call) + continue + function = tool_call.get("function") + if not isinstance(function, dict) or not isinstance(function.get("arguments"), str): + normalized_calls.append(tool_call) + continue + try: + arguments = json.loads(function["arguments"]) + except json.JSONDecodeError: + normalized_calls.append(tool_call) + continue + if not isinstance(arguments, dict): + normalized_calls.append(tool_call) + continue + normalized_calls.append({**tool_call, "function": {**function, "arguments": arguments}}) + changed = True + + normalized_messages.append({**message, "tool_calls": normalized_calls} if changed else message) + messages_changed = messages_changed or changed + return normalized_messages if messages_changed else messages + def _count_input_tokens( self, messages: list[NeMoGymChatCompletionMessageParam], tools: Optional[dict[str, Tool]] = None, ) -> int: + """Return the best available prompt-token estimate.""" + return self._count_input_tokens_with_confidence(messages, tools)[0] + + def _count_input_tokens_with_confidence( + self, + messages: list[NeMoGymChatCompletionMessageParam], + tools: Optional[dict[str, Tool]] = None, + ) -> tuple[int, bool]: """Estimate the full prompt token count the server will see. ``messages`` must already be serialized for the provider. This keeps @@ -279,19 +423,26 @@ def _count_input_tokens( rough but serialises everything. 4. Character-count fallback when no tokenizer is present. + Returns ``(count, exact_template_render)``. Only a successful render of + the complete prompt is exact enough to impose a hard context bound. + JSON and character fallbacks remain useful for budget sizing, but must + not turn an approximate over-count into a false context-overflow error. + Any residual gap is absorbed by ``completion_token_buffer``. """ import json as _json if self._tokenizer is None: - # Pure character-count fallback. - total = sum(len(str(m.get("content") or "")) for m in messages) // 3 + # Pure character-count fallback. Count the complete serialized + # payload, including assistant tool-call names and arguments. + counted_messages = self._messages_for_estimator_count(messages) + total = len(_json.dumps(counted_messages, ensure_ascii=False)) // 3 if tools: try: total += len(_json.dumps(to_openai_tools(tools))) // 3 except Exception: pass - return total + return total, False oai_tools = None if tools: @@ -301,38 +452,68 @@ def _count_input_tokens( LOGGER.warning(f"to_openai_tools failed ({exc}).") # Strategy 1: apply_chat_template with tools= + template_kwargs: dict[str, Any] = {} + if self._prompt_estimator_truncate_history_thinking is not None: + template_kwargs["truncate_history_thinking"] = self._prompt_estimator_truncate_history_thinking + normalized_template_messages = self._messages_for_template_count(messages) + template_message_variants = [messages] + if normalized_template_messages is not messages: + template_message_variants.append(normalized_template_messages) if oai_tools is not None: + last_template_error = None + for template_messages in template_message_variants: + try: + text = self._tokenizer.apply_chat_template( + template_messages, + tools=oai_tools, + tokenize=False, + add_generation_prompt=True, + **template_kwargs, + ) + return len(self._tokenizer(text, add_special_tokens=False)["input_ids"]), True + except Exception as exc: + last_template_error = exc + LOGGER.debug( + "apply_chat_template(tools=) unsupported (%s); trying separate tool count.", + last_template_error, + ) + + # Strategy 2: apply_chat_template on messages only + separate tool JSON count + last_template_error = None + for template_messages in template_message_variants: try: text = self._tokenizer.apply_chat_template( - messages, tools=oai_tools, tokenize=False, add_generation_prompt=True + template_messages, + tokenize=False, + add_generation_prompt=True, + **template_kwargs, ) - return len(self._tokenizer(text, add_special_tokens=False)["input_ids"]) + total = len(self._tokenizer(text, add_special_tokens=False)["input_ids"]) + if oai_tools is not None: + total += len(self._tokenizer(_json.dumps(oai_tools), add_special_tokens=False)["input_ids"]) + # With no tools this is the complete rendered prompt. When tools + # exist, tokenizing their JSON separately does not reproduce the + # template's wrappers or token boundaries and remains approximate. + return total, oai_tools is None except Exception as exc: - LOGGER.debug(f"apply_chat_template(tools=) unsupported ({exc}); trying separate tool count.") - - # Strategy 2: apply_chat_template on messages only + separate tool JSON count - try: - text = self._tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) - total = len(self._tokenizer(text, add_special_tokens=False)["input_ids"]) - if oai_tools is not None: - total += len(self._tokenizer(_json.dumps(oai_tools), add_special_tokens=False)["input_ids"]) - return total - except Exception as exc: - LOGGER.warning(f"apply_chat_template(messages) failed ({exc}); falling back to JSON count.") + last_template_error = exc + LOGGER.warning("apply_chat_template(messages) failed (%s); falling back to JSON count.", last_template_error) # Strategy 3: tokenise the full JSON payload try: - blob = _json.dumps(messages) + counted_messages = self._messages_for_estimator_count(messages) + blob = _json.dumps(counted_messages) total = len(self._tokenizer(blob, add_special_tokens=False)["input_ids"]) if oai_tools is not None: total += len(self._tokenizer(_json.dumps(oai_tools), add_special_tokens=False)["input_ids"]) - return total + return total, False except Exception as exc: LOGGER.warning(f"JSON tokenisation failed ({exc}); falling back to character count.") # Strategy 4: character count - total = sum(len(str(m.get("content") or "")) for m in messages) // 3 - return total + counted_messages = self._messages_for_estimator_count(messages) + total = len(_json.dumps(counted_messages, ensure_ascii=False)) // 3 + return total, False async def generate( self, @@ -352,13 +533,20 @@ async def generate( if recovering: provider_messages = [*provider_messages, {"role": "user", "content": _TRUNCATION_RECOVERY_NUDGE}] - input_tokens = self._count_input_tokens(provider_messages, tools) + input_tokens, exact_template_render = self._count_input_tokens_with_confidence(provider_messages, tools) context_window = self._max_tokens + estimated_remaining_context = context_window - input_tokens + if exact_template_render and estimated_remaining_context <= 0: + raise ContextOverflowError( + f"Estimated prompt ({input_tokens} tokens) leaves no room in the {context_window}-token context window" + ) dynamic_max = max( - context_window - input_tokens - self._completion_token_buffer, - _MIN_COMPLETION_TOKENS, + estimated_remaining_context - self._completion_token_buffer, + self._min_completion_tokens, ) capped_max = min(dynamic_max, self._max_completion_tokens_cap) + if exact_template_render: + capped_max = min(capped_max, estimated_remaining_context) # ``self._kwargs`` is spread last so explicit per-request kwargs override # the agent-level defaults. @@ -447,15 +635,30 @@ async def generate( for tc in (msg.tool_calls or []) ] - # A call that exhausted its completion budget *and* produced no tool call - # advanced the task by nothing. Arm the recovery turn (see - # _TRUNCATION_RECOVERY_NUDGE). A truncated call that still carried a tool - # call made progress and is left alone. - if choice.finish_reason in ("length", "max_tokens") and not tool_calls: + # A call that exhausted its completion budget without a schema-valid + # tool call advanced the task by nothing. Parser-truncated calls often + # surface as ``code_exec({})``; treating their mere presence as progress + # creates a sticky invalid-call loop. + usable_tool_call = False + for tool_call in tool_calls: + tool = tools.get(tool_call.name) + if tool is None: + # We cannot safely validate an unknown provider-specific tool, + # so preserve the historical assumption that it made progress. + usable_tool_call = True + break + try: + tool.parameters.model_validate_json(tool_call.arguments or "{}") + except _PydanticValidationError: + continue + usable_tool_call = True + break + + if choice.finish_reason in ("length", "max_tokens") and not usable_tool_call: self._truncation_overruns += 1 self._recover_from_truncation = True LOGGER.warning( - "completion budget (%d tokens) exhausted with no tool call [overrun #%d]; %s", + "completion budget (%d tokens) exhausted with no usable tool call [overrun #%d]; %s", capped_max, self._truncation_overruns, "next turn will run with thinking disabled" diff --git a/responses_api_agents/stirrup_agent/tests/test_app.py b/responses_api_agents/stirrup_agent/tests/test_app.py index 0a5e99ec166..983b41917d4 100644 --- a/responses_api_agents/stirrup_agent/tests/test_app.py +++ b/responses_api_agents/stirrup_agent/tests/test_app.py @@ -17,6 +17,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from omegaconf import OmegaConf from stirrup.core.models import AssistantMessage, TokenUsage, ToolCall from nemo_gym.config_types import ModelServerRef, ResourcesServerRef @@ -129,6 +130,22 @@ def test_sanity(self) -> None: ) StirrupAgentWrapper(config=config, server_client=MagicMock(spec=ServerClient)) + # Generic Stirrup users retain the historical behavior. GDPVal opts in + # to its larger floor and matching template semantics in its benchmark + # config below. + assert config.min_completion_tokens == 1024 + assert config.prompt_estimator_truncate_history_thinking is None + assert config.min_compaction_summary_words == 1 + + def test_gdpval_config_opts_into_safe_context_budget(self) -> None: + repo_root = STIRRUP_AGENT_DIR.parents[1] + config = OmegaConf.load(repo_root / "benchmarks" / "gdpval" / "config.yaml") + agent = config.gdpval_stirrup_agent.responses_api_agents.stirrup_agent + + assert agent.min_completion_tokens == 8192 + assert "prompt_estimator_truncate_history_thinking" not in agent + assert agent.min_compaction_summary_words == 50 + def test_output_history_preserves_nemo_user_tool_results(self) -> None: """Run-history export should keep NeMo user-role tool results as tool outputs.""" history = [ diff --git a/responses_api_agents/stirrup_agent/tests/test_nemo_agent_compaction.py b/responses_api_agents/stirrup_agent/tests/test_nemo_agent_compaction.py new file mode 100644 index 00000000000..af860a976d3 --- /dev/null +++ b/responses_api_agents/stirrup_agent/tests/test_nemo_agent_compaction.py @@ -0,0 +1,303 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Compaction hardening for the Gym Stirrup adapter.""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest +from stirrup.core.exceptions import ContextOverflowError +from stirrup.core.models import ( + AssistantMessage, + Reasoning, + SummaryMessage, + SystemMessage, + TokenUsage, + ToolMessage, + UserMessage, +) +from stirrup.prompts import MESSAGE_SUMMARIZER + +from responses_api_agents.stirrup_agent.nemo_agent import ( + NeMoAgent, + NeMoSessionAgent, + _is_meaningful_compaction_summary, + _sanitize_compaction_summary, +) + + +def _words(count: int) -> str: + return " ".join(f"word{index}" for index in range(count)) + + +def _assistant(content: str, *, message_id: str = "assistant-1") -> AssistantMessage: + return AssistantMessage( + id=message_id, + content=content, + token_usage=TokenUsage(input=1, answer=1, reasoning=0), + ) + + +class _FakeClient: + model_slug = "fake-model" + max_tokens = 262_144 + + def __init__(self, *results: AssistantMessage | Exception) -> None: + self.results = list(results) + self.calls: list[tuple[list, dict]] = [] + + async def generate(self, messages, tools): + self.calls.append((messages, tools)) + result = self.results.pop(0) + if isinstance(result, Exception): + raise result + return result + + +class _RecoveryAwareFakeClient(_FakeClient): + def __init__( + self, + *results: AssistantMessage | Exception, + recovery: bool, + overruns: int, + mutate_during_generate: bool = False, + ) -> None: + super().__init__(*results) + self._recover_from_truncation = recovery + self._truncation_overruns = overruns + self.mutate_during_generate = mutate_during_generate + self.recovery_seen: list[bool] = [] + + async def generate(self, messages, tools): + self.recovery_seen.append(self._recover_from_truncation) + if self.mutate_during_generate: + self._recover_from_truncation = True + self._truncation_overruns += 1 + return await super().generate(messages, tools) + + +def _agent( + client: _FakeClient, + logger: MagicMock | None = None, + *, + min_words: int = 1, +) -> NeMoAgent: + return NeMoAgent( + client=client, + name="test_agent", + tools=[], + logger=logger or MagicMock(), + min_compaction_summary_words=min_words, + ) + + +def _history() -> list: + return [ + SystemMessage(content="system"), + UserMessage(content="original task"), + _assistant("worked", message_id="assistant-1"), + ToolMessage(content="tool result", tool_call_id="call-1", name="code_exec", success=True), + ] + + +def test_sanitize_compaction_summary_removes_only_leading_think_blocks() -> None: + summary = " private reasoning\n\n" + _words(50) + " literal tag" + + sanitized = _sanitize_compaction_summary(summary) + + assert sanitized.startswith("word0 word1") + assert "private reasoning" not in sanitized + assert sanitized.endswith("literal tag") + + +@pytest.mark.parametrize( + "summary", + [ + "", + " ", + "reasoning without a closing tag", + "reasoning\n\n", + "first second but unclosed", + ], +) +def test_sanitize_compaction_summary_rejects_empty_or_unclosed_think(summary: str) -> None: + assert _sanitize_compaction_summary(summary) == "" + + +def test_meaningful_compaction_summary_uses_configured_minimum() -> None: + assert not _is_meaningful_compaction_summary(_words(49), min_words=50) + assert _is_meaningful_compaction_summary(_words(50), min_words=50) + + +@pytest.mark.parametrize("summary", ["Done.", "已完成并保留所有关键进展。"]) +def test_generic_compaction_accepts_concise_and_cjk_summaries(summary: str) -> None: + assert _is_meaningful_compaction_summary(summary) + + +def test_compaction_minimum_must_be_positive() -> None: + with pytest.raises(ValueError, match="at least 1"): + _agent(_FakeClient(), min_words=0) + + +def test_session_agent_uses_the_hardened_compaction_override() -> None: + assert NeMoSessionAgent.summarize_messages is NeMoAgent.summarize_messages + + +@pytest.mark.asyncio +async def test_summarization_retries_degenerate_outputs_and_persists_only_sanitized_text() -> None: + good_summary = _words(60) + client = _FakeClient( + _assistant("reasoning only"), + _assistant("too short to preserve a long conversation"), + AssistantMessage( + content=f"private chain of thought\n{good_summary}", + reasoning=Reasoning(content="separate private reasoning"), + ), + ) + logger = MagicMock() + agent = _agent(client, logger, min_words=50) + history = _history() + + archived, compacted = await agent.summarize_messages(history, {"assistant-1": {}}) + + assert archived == history + assert len(client.calls) == 3 + assert client.calls[0][0][-1].content == MESSAGE_SUMMARIZER + assert client.calls[0][1] + assert "text only" in client.calls[1][0][-1].content.lower() + assert client.calls[1][1] + assert client.calls[2][1] == {} + + summaries = [message for message in compacted if isinstance(message, SummaryMessage)] + assert len(summaries) == 1 + persisted = summaries[0].content + assert good_summary in persisted + assert "" not in persisted.lower() + assert "private" not in persisted.lower() + assert compacted[-1] == UserMessage(content="Got it, thanks!") + + logger.context_summarization_complete.assert_called_once() + logged_summary, logged_bridge = logger.context_summarization_complete.call_args.args + assert logged_summary == good_summary + assert "" not in logged_bridge.lower() + + +@pytest.mark.asyncio +async def test_summarization_fails_closed_after_three_degenerate_outputs() -> None: + client = _FakeClient( + _assistant("one"), + _assistant(_words(10)), + _assistant("unclosed"), + ) + logger = MagicMock() + agent = _agent(client, logger, min_words=50) + metadata = {"assistant-1": {"code_exec": []}} + + with pytest.raises(RuntimeError, match="usable text after 3 attempts"): + await agent.summarize_messages(_history(), metadata) + + assert len(client.calls) == 3 + assert metadata == {"assistant-1": {"code_exec": []}} + logger.context_summarization_complete.assert_not_called() + + +@pytest.mark.asyncio +async def test_summarization_preserves_context_overflow_unwind_semantics() -> None: + history = [ + SystemMessage(content="system"), + UserMessage(content="original task"), + _assistant("first turn", message_id="assistant-1"), + ToolMessage(content="first result", tool_call_id="call-1", name="code_exec", success=True), + _assistant("second turn", message_id="assistant-2"), + ToolMessage(content="second result", tool_call_id="call-2", name="code_exec", success=True), + ] + client = _FakeClient(ContextOverflowError("too large"), _assistant(_words(60))) + agent = _agent(client, min_words=50) + metadata = {"assistant-1": {"first": []}, "assistant-2": {"second": []}} + + archived, compacted = await agent.summarize_messages(history, metadata) + + assert archived == history[:4] + assert "assistant-1" in metadata + assert "assistant-2" not in metadata + assert any(isinstance(message, SummaryMessage) for message in compacted) + assert len(client.calls) == 2 + + +@pytest.mark.asyncio +async def test_summary_calls_do_not_consume_an_armed_policy_recovery() -> None: + client = _RecoveryAwareFakeClient(_assistant(_words(60)), recovery=True, overruns=7) + agent = _agent(client, min_words=50) + + await agent.summarize_messages(_history(), {}) + + assert client.recovery_seen == [False] + assert client._recover_from_truncation is True + assert client._truncation_overruns == 7 + + +@pytest.mark.asyncio +async def test_summary_calls_do_not_arm_policy_recovery_or_increment_overruns() -> None: + client = _RecoveryAwareFakeClient( + _assistant("too short"), + _assistant("still too short"), + _assistant(_words(60)), + recovery=False, + overruns=2, + mutate_during_generate=True, + ) + agent = _agent(client, min_words=50) + + await agent.summarize_messages(_history(), {}) + + assert client.recovery_seen == [False, False, False] + assert client._recover_from_truncation is False + assert client._truncation_overruns == 2 + + +@pytest.mark.asyncio +async def test_session_run_replays_only_the_sanitized_summary_after_compaction() -> None: + good_summary = _words(60) + client = _FakeClient( + AssistantMessage( + content="first policy turn", + token_usage=TokenUsage(input=190_000, answer=10, reasoning=0), + ), + _assistant(f"private summary reasoning\n{good_summary}"), + _assistant("second policy turn"), + ) + agent = NeMoAgent( + client=client, + name="test_agent", + tools=[], + logger=MagicMock(), + max_turns=2, + context_summarization_cutoff=0.7, + min_compaction_summary_words=50, + ) + + async with agent.session(cache_on_interrupt=False) as session: + assert isinstance(session, NeMoSessionAgent) + _, history, _ = await session.run("original task") + + assert len(history) == 2 + assert any(isinstance(message, SummaryMessage) for message in history[1]) + second_policy_messages = client.calls[2][0] + serialized_context = "\n".join( + message.content for message in second_policy_messages if isinstance(message.content, str) + ) + assert good_summary in serialized_context + assert "private summary reasoning" not in serialized_context + assert "" not in serialized_context.lower() diff --git a/responses_api_agents/stirrup_agent/tests/test_nemo_client.py b/responses_api_agents/stirrup_agent/tests/test_nemo_client.py index 41b8833fe14..24287d77b28 100644 --- a/responses_api_agents/stirrup_agent/tests/test_nemo_client.py +++ b/responses_api_agents/stirrup_agent/tests/test_nemo_client.py @@ -20,10 +20,14 @@ from __future__ import annotations +import json +from copy import deepcopy from unittest.mock import AsyncMock, MagicMock import pytest -from stirrup.core.models import AssistantMessage, SystemMessage, TokenUsage, ToolCall, ToolMessage, UserMessage +from pydantic import BaseModel +from stirrup.core.exceptions import ContextOverflowError +from stirrup.core.models import AssistantMessage, SystemMessage, TokenUsage, Tool, ToolCall, ToolMessage, UserMessage from responses_api_agents.stirrup_agent.nemo_agent import NeMoUserMessage from responses_api_agents.stirrup_agent.nemo_client import ( @@ -80,6 +84,7 @@ async def test_generate_forwards_configured_sampling_kwargs() -> None: assert sent["temperature"] == pytest.approx(0.42) assert sent["top_p"] == pytest.approx(0.7) assert sent["extra_body"]["chat_template_kwargs"]["enable_thinking"] is False + assert "truncate_history_thinking" not in sent["extra_body"]["chat_template_kwargs"] assert sent["max_completion_tokens"] <= 2048 @@ -209,9 +214,375 @@ def test_defaults_match_pre_lift_behaviour() -> None: assert client._top_p == 0.95 assert client._enable_thinking is True assert client._max_completion_tokens_cap == 64000 + assert client._min_completion_tokens == 1024 + assert client._prompt_estimator_truncate_history_thinking is None assert client._truncation_recovery is True +def _long_thinking_history(*, closed: bool = True) -> list: + messages = [UserMessage(content="do the task")] + for _ in range(4): + reasoning = "r" * 195_000 + content = f"{reasoning}visible" if closed else f"{reasoning}" + messages.extend( + [ + AssistantMessage(content=content, token_usage=TokenUsage(input=1, answer=1, reasoning=0)), + UserMessage(content="Please continue the task"), + ] + ) + return messages + + +@pytest.mark.asyncio +async def test_fallback_ignores_closed_historical_thinking_when_template_truncates_it() -> None: + """Regression: raw 64k thinking turns must not collapse later calls to the floor. + + Nemotron's chat template removes each completed reasoning block once a later + user message exists. The tokenizer-free fallback must estimate that rendered + prompt, not the much larger forensic history retained by Stirrup. + """ + client = DynamicMaxTokensChatCompletionsClient( + model="m", + max_tokens=262_144, + base_url="http://test", + api_key="k", + min_completion_tokens=8192, + max_completion_tokens_cap=64_000, + prompt_estimator_truncate_history_thinking=True, + ) + fake_create = AsyncMock(return_value=_make_response()) + client._client = MagicMock() + client._client.chat.completions.create = fake_create + + history = _long_thinking_history() + original_content = history[1].content + await client.generate(history, tools={}) + + sent = fake_create.await_args.kwargs + assert sent["max_completion_tokens"] == 64_000 + assert "truncate_history_thinking" not in sent["extra_body"]["chat_template_kwargs"] + assert sent["messages"][1]["content"][0]["text"] == original_content + assert history[1].content == original_content + + +def test_tokenizer_receives_estimator_history_truncation_setting() -> None: + class RecordingTokenizer: + def __init__(self) -> None: + self.template_kwargs = None + + def apply_chat_template(self, messages, **kwargs): + self.template_kwargs = kwargs + return "rendered" + + def __call__(self, text, **kwargs): + return {"input_ids": [1, 2, 3]} + + client = DynamicMaxTokensChatCompletionsClient( + model="m", + max_tokens=262_144, + base_url="http://test", + api_key="k", + prompt_estimator_truncate_history_thinking=True, + ) + tokenizer = RecordingTokenizer() + client._tokenizer = tokenizer + + assert client._count_input_tokens(to_provider_openai_messages([UserMessage(content="hi")]), tools={}) == 3 + assert tokenizer.template_kwargs["truncate_history_thinking"] is True + + +def test_token_count_is_exact_only_when_the_complete_prompt_renders() -> None: + class MessagesOnlyTokenizer: + def apply_chat_template(self, messages, **kwargs): + if "tools" in kwargs: + raise TypeError("template does not accept tools") + return "rendered messages" + + def __call__(self, text, **kwargs): + return {"input_ids": [1, 2, 3]} + + client = DynamicMaxTokensChatCompletionsClient( + model="m", + max_tokens=262_144, + base_url="http://test", + api_key="k", + ) + client._tokenizer = MessagesOnlyTokenizer() + messages = to_provider_openai_messages([UserMessage(content="hi")]) + + count_without_tools, exact_without_tools = client._count_input_tokens_with_confidence(messages, tools={}) + count_with_tools, exact_with_tools = client._count_input_tokens_with_confidence( + messages, + tools={"code_exec": _code_exec_tool()}, + ) + + assert count_without_tools == 3 + assert exact_without_tools is True + assert count_with_tools > 3 + assert exact_with_tools is False + + +def test_nemotron_estimator_decodes_tool_arguments_for_exact_template_render() -> None: + class MappingArgumentsTokenizer: + def __init__(self) -> None: + self.rendered_messages = None + + def apply_chat_template(self, messages, **kwargs): + arguments = messages[1]["tool_calls"][0]["function"]["arguments"] + if not isinstance(arguments, dict): + raise TypeError("tool arguments must be a mapping") + self.rendered_messages = messages + return "complete rendered prompt" + + def __call__(self, text, **kwargs): + return {"input_ids": [1, 2, 3, 4]} + + client = DynamicMaxTokensChatCompletionsClient( + model="m", + max_tokens=262_144, + base_url="http://test", + api_key="k", + prompt_estimator_truncate_history_thinking=True, + ) + tokenizer = MappingArgumentsTokenizer() + client._tokenizer = tokenizer + messages = _serialized_history_with_prior_assistant("oldvisible", with_tool_call=True) + original = deepcopy(messages) + + count, exact = client._count_input_tokens_with_confidence( + messages, + tools={"code_exec": _code_exec_tool()}, + ) + + assert (count, exact) == (4, True) + assert tokenizer.rendered_messages[1]["tool_calls"][0]["function"]["arguments"] == {"cmd": "echo hello"} + assert messages == original + + +def test_nemotron_estimator_does_not_claim_exact_render_for_invalid_tool_arguments() -> None: + class MappingArgumentsTokenizer: + def apply_chat_template(self, messages, **kwargs): + arguments = messages[1]["tool_calls"][0]["function"]["arguments"] + if not isinstance(arguments, dict): + raise TypeError("tool arguments must be a mapping") + return "complete rendered prompt" + + def __call__(self, text, **kwargs): + return {"input_ids": [1, 2, 3]} + + client = DynamicMaxTokensChatCompletionsClient( + model="m", + max_tokens=262_144, + base_url="http://test", + api_key="k", + prompt_estimator_truncate_history_thinking=True, + ) + client._tokenizer = MappingArgumentsTokenizer() + messages = _serialized_history_with_prior_assistant("visible", with_tool_call=True) + messages[1]["tool_calls"][0]["function"]["arguments"] = "not-json" + + _, exact = client._count_input_tokens_with_confidence( + messages, + tools={"code_exec": _code_exec_tool()}, + ) + + assert exact is False + + +def _serialized_history_with_prior_assistant(content: str, *, with_tool_call: bool) -> list: + tool_calls = ( + [ToolCall(tool_call_id="call_1", name="code_exec", arguments='{"cmd":"echo hello"}')] if with_tool_call else [] + ) + return to_provider_openai_messages( + [ + UserMessage(content="start"), + AssistantMessage( + content=content, + tool_calls=tool_calls, + token_usage=TokenUsage(input=1, answer=1, reasoning=0), + ), + UserMessage(content="continue"), + ] + ) + + +@pytest.mark.parametrize( + ("content", "with_tool_call", "expected"), + [ + ("oldvisible", False, "visible"), + ("unfinished", False, "unfinished"), + ("prefixoldvisible", True, "visible"), + ("prefixunfinished", True, "prefix"), + ], + ids=("closed-no-tool", "unclosed-no-tool", "closed-tool", "unclosed-tool"), +) +def test_estimator_mirrors_nemotron_prior_assistant_semantics( + content: str, + with_tool_call: bool, + expected: str, +) -> None: + client = DynamicMaxTokensChatCompletionsClient( + model="m", + max_tokens=262_144, + base_url="http://test", + api_key="k", + prompt_estimator_truncate_history_thinking=True, + ) + serialized = _serialized_history_with_prior_assistant(content, with_tool_call=with_tool_call) + original = deepcopy(serialized) + + estimated = client._messages_for_estimator_count(serialized) + + assert estimated[1]["content"][0]["text"] == expected + assert estimated[1].get("tool_calls") == serialized[1].get("tool_calls") + assert serialized == original + + +def test_estimator_uses_content_after_last_closed_think_block() -> None: + client = DynamicMaxTokensChatCompletionsClient( + model="m", + max_tokens=262_144, + base_url="http://test", + api_key="k", + prompt_estimator_truncate_history_thinking=True, + ) + serialized = _serialized_history_with_prior_assistant( + "firstmiddlesecondtail", + with_tool_call=False, + ) + + estimated = client._messages_for_estimator_count(serialized) + + assert estimated[1]["content"][0]["text"] == "tail" + + +def test_character_fallback_counts_complete_serialized_messages() -> None: + client = DynamicMaxTokensChatCompletionsClient( + model="m", + max_tokens=262_144, + base_url="http://test", + api_key="k", + ) + serialized = _serialized_history_with_prior_assistant("same content", with_tool_call=True) + expected = len(json.dumps(serialized, ensure_ascii=False)) // 3 + + count = client._count_input_tokens(serialized, tools={}) + + assert count == expected + long_arguments = deepcopy(serialized) + long_arguments[1]["tool_calls"][0]["function"]["arguments"] = "x" * 3000 + assert client._count_input_tokens(long_arguments, tools={}) >= count + 990 + + +@pytest.mark.asyncio +async def test_completion_budget_is_bounded_by_positive_estimated_remaining_context() -> None: + client = DynamicMaxTokensChatCompletionsClient( + model="m", + max_tokens=10_000, + base_url="http://test", + api_key="k", + completion_token_buffer=1000, + min_completion_tokens=8192, + max_completion_tokens_cap=64_000, + ) + client._count_input_tokens_with_confidence = MagicMock(return_value=(9500, True)) + client._tokenizer = object() + fake_create = AsyncMock(return_value=_make_response()) + client._client = MagicMock() + client._client.chat.completions.create = fake_create + + await client.generate([UserMessage(content="hi")], tools={}) + + assert fake_create.await_args.kwargs["max_completion_tokens"] == 500 + + +@pytest.mark.asyncio +async def test_no_estimated_context_remaining_raises_before_dispatch() -> None: + client = DynamicMaxTokensChatCompletionsClient( + model="m", + max_tokens=10_000, + base_url="http://test", + api_key="k", + min_completion_tokens=8192, + ) + client._count_input_tokens_with_confidence = MagicMock(return_value=(10_000, True)) + client._tokenizer = object() + fake_create = AsyncMock(return_value=_make_response()) + client._client = MagicMock() + client._client.chat.completions.create = fake_create + + with pytest.raises(ContextOverflowError, match="leaves no room"): + await client.generate([UserMessage(content="hi")], tools={}) + + fake_create.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_fallback_overestimate_still_dispatches_configured_floor() -> None: + client = DynamicMaxTokensChatCompletionsClient( + model="m", + max_tokens=10_000, + base_url="http://test", + api_key="k", + min_completion_tokens=8192, + max_completion_tokens_cap=64_000, + ) + client._count_input_tokens_with_confidence = MagicMock(return_value=(20_000, False)) + fake_create = AsyncMock(return_value=_make_response()) + client._client = MagicMock() + client._client.chat.completions.create = fake_create + + await client.generate([UserMessage(content="hi")], tools={}) + + assert client._tokenizer is None + assert fake_create.await_args.kwargs["max_completion_tokens"] == 8192 + + +@pytest.mark.asyncio +async def test_loaded_tokenizer_approximation_does_not_impose_a_false_context_bound() -> None: + client = DynamicMaxTokensChatCompletionsClient( + model="m", + max_tokens=10_000, + base_url="http://test", + api_key="k", + min_completion_tokens=8192, + max_completion_tokens_cap=64_000, + ) + client._tokenizer = object() + client._count_input_tokens_with_confidence = MagicMock(return_value=(20_000, False)) + fake_create = AsyncMock(return_value=_make_response()) + client._client = MagicMock() + client._client.chat.completions.create = fake_create + + await client.generate([UserMessage(content="hi")], tools={}) + + assert fake_create.await_args.kwargs["max_completion_tokens"] == 8192 + + +@pytest.mark.asyncio +async def test_minimum_completion_floor_never_consumes_more_than_remaining_context() -> None: + client = DynamicMaxTokensChatCompletionsClient( + model="m", + max_tokens=10_000, + base_url="http://test", + api_key="k", + completion_token_buffer=1000, + min_completion_tokens=8192, + max_completion_tokens_cap=64_000, + ) + client._count_input_tokens_with_confidence = MagicMock(return_value=(1500, True)) + client._tokenizer = object() + fake_create = AsyncMock(return_value=_make_response()) + client._client = MagicMock() + client._client.chat.completions.create = fake_create + + await client.generate([UserMessage(content="hi")], tools={}) + + # 8,500 raw tokens remain. The floor may reclaim part of the 1,000-token + # safety buffer, but it may not exceed the actual estimated remainder. + assert fake_create.await_args.kwargs["max_completion_tokens"] == 8192 + + # --------------------------------------------------------------------------- # Thinking-overrun recovery # --------------------------------------------------------------------------- @@ -249,6 +620,19 @@ def _tool_call_response(finish_reason: str = "stop"): return response +class _CodeExecParams(BaseModel): + cmd: str + + +def _code_exec_tool() -> Tool: + return Tool( + name="code_exec", + description="Run a command.", + parameters=_CodeExecParams, + executor=lambda _: None, + ) + + @pytest.mark.asyncio async def test_budget_overrun_without_tool_call_arms_recovery() -> None: """The next call drops thinking and carries an explicit instruction.""" @@ -325,6 +709,21 @@ async def test_truncated_call_that_still_made_a_tool_call_is_left_alone() -> Non assert len(second["messages"]) == 1 +@pytest.mark.asyncio +@pytest.mark.parametrize("arguments", ["", "{}", '{"wrong":"field"}']) +async def test_truncated_schema_invalid_tool_call_arms_recovery(arguments: str) -> None: + """A parser-emitted call with unusable arguments made no task progress.""" + client = _make_client() + response = _tool_call_response(finish_reason="length") + response.choices[0].message.tool_calls[0].function.arguments = arguments + client._client.chat.completions.create = AsyncMock(return_value=response) + + await client.generate([UserMessage(content="hi")], tools={"code_exec": _code_exec_tool()}) + + assert client._recover_from_truncation is True + assert client._truncation_overruns == 1 + + @pytest.mark.asyncio async def test_text_only_response_that_finished_normally_is_left_alone() -> None: """A short text-only turn is a different problem; don't touch it.""" From 470a5e324db2a1b9341a2fd2da438fd70e4c43a5 Mon Sep 17 00:00:00 2001 From: Serge Panev Date: Fri, 4 Sep 2026 22:07:21 -0700 Subject: [PATCH 2/2] [GDPVal] Separate output cap from context window Signed-off-by: Serge Panev --- benchmarks/gdpval/config.yaml | 1 + responses_api_agents/stirrup_agent/README.md | 1 + responses_api_agents/stirrup_agent/app.py | 23 ++++++++--- .../stirrup_agent/tests/test_app.py | 41 +++++++++++++++++++ .../stirrup_agent/tests/test_nemo_client.py | 21 ++++++++++ 5 files changed, 82 insertions(+), 5 deletions(-) diff --git a/benchmarks/gdpval/config.yaml b/benchmarks/gdpval/config.yaml index 85275b49bf6..c22c63fc69c 100644 --- a/benchmarks/gdpval/config.yaml +++ b/benchmarks/gdpval/config.yaml @@ -106,6 +106,7 @@ gdpval_stirrup_agent: # GDPVal tasks routinely need multi-thousand-token scripts. A 1,024-token # completion cannot form a useful tool call and can create a sticky # code_exec({}) loop after long reasoning turns. + context_window_tokens: 262144 min_completion_tokens: 8192 # Reject the observed near-empty compaction outputs that erase progress # on long GDPVal sessions. Generic Stirrup agents accept any non-empty diff --git a/responses_api_agents/stirrup_agent/README.md b/responses_api_agents/stirrup_agent/README.md index ce8dcd2c22c..87b34d88800 100644 --- a/responses_api_agents/stirrup_agent/README.md +++ b/responses_api_agents/stirrup_agent/README.md @@ -131,6 +131,7 @@ The agent reads its Hydra config at `configs/stirrup_agent.yaml`. Notable keys: | `rerun_incomplete` | `false` | If true, skip the rollout for tasks that already **finished** (a finish marker is cached) and only re-run the ones that didn't (see [Task Re-run Mode](#task-re-run-mode)). Composes with `execute_only` and `judge_only`. | | `model_id` | `null` | HF model id or local path used to load a tokenizer for dynamic output sizing. | | `completion_token_buffer` | `1000` | Safety margin (in tokens) reserved when sizing `max_completion_tokens` per call. | +| `context_window_tokens` | `262144` | Model context-window size used for dynamic completion budgeting. Independent of a request's `max_output_tokens`. | | `min_completion_tokens` | `1024` | Target useful per-call completion floor. The hard cap always applies; estimated remaining context is also a strict bound when the tokenizer successfully renders the complete prompt. Approximate fallbacks preserve this floor even when their estimate exceeds context. GDPVal raises it to `8192`. | | `prompt_estimator_truncate_history_thinking` | `null` | Optional estimator-only switch for templates that omit completed historical reasoning. It is never sent to the provider. | | `min_compaction_summary_words` | `1` | Minimum summary length accepted during context compaction. GDPVal uses `50`; invalid summaries retry up to three times and never replace the live history. | diff --git a/responses_api_agents/stirrup_agent/app.py b/responses_api_agents/stirrup_agent/app.py index bc1f3ced3e7..5d376e3b053 100644 --- a/responses_api_agents/stirrup_agent/app.py +++ b/responses_api_agents/stirrup_agent/app.py @@ -475,7 +475,7 @@ async def _run_stirrup_agent( api_key: str = "dummy", max_turns: int = 250, temperature: float = 0.6, - max_tokens: int = 262144, + context_window_tokens: int = 262144, reference_files: Optional[list] = None, reference_file_urls: Optional[list] = None, exec_provider_class: Optional[str] = None, @@ -505,6 +505,9 @@ async def _run_stirrup_agent( a tokenizer that sizes ``max_completion_tokens`` dynamically per call (see ``DynamicMaxTokensChatCompletionsClient``). When unset, a character-count fallback is used. + + *context_window_tokens* describes model capacity. It is deliberately + separate from the outer Responses API request's ``max_output_tokens``. """ from stirrup.tools import DEFAULT_TOOLS from stirrup.tools.code_backends.base import SHELL_TIMEOUT, CodeExecToolProvider, CommandResult @@ -553,7 +556,7 @@ async def run_command(self, cmd: str, *, timeout: int = SHELL_TIMEOUT) -> Comman model=model_name, base_url=model_base_url, api_key=api_key, - max_tokens=max_tokens, + max_tokens=context_window_tokens, model_id=model_id, completion_token_buffer=completion_token_buffer, temperature=temperature, @@ -985,6 +988,13 @@ class StirrupAgentWrapperConfig(BaseResponsesAPIAgentConfig): "(messages + tool-schema JSON) and the exact prompt the server sees after chat-template " "rendering. See ``nemo_client.DynamicMaxTokensChatCompletionsClient``.", ) + context_window_tokens: int = Field( + default=262144, + ge=1, + description="Model context-window size used for dynamic per-call completion budgeting. " + "This is independent of the Responses API request's max_output_tokens, which limits " + "output rather than describing model capacity.", + ) top_p: float = Field( default=0.95, description="Top-p sampling cutoff for the policy model. Forwarded to the LLM client.", @@ -1191,7 +1201,10 @@ async def responses( model_name = getattr(body, "model", None) or "default" temperature = getattr(body, "temperature", None) or self.config.temperature - max_tokens = getattr(body, "max_output_tokens", 262144) or 262144 + requested_output_tokens = getattr(body, "max_output_tokens", None) + max_completion_tokens_cap = self.config.max_completion_tokens_cap + if requested_output_tokens is not None: + max_completion_tokens_cap = min(max_completion_tokens_cap, requested_output_tokens) exec_provider = self.task_strategy.get_exec_provider(task_info, self.config) exec_provider_class = None @@ -1211,7 +1224,7 @@ async def responses( "api_key": "dummy", # pragma: allowlist secret "max_turns": self.config.agent_max_turns, "temperature": temperature, - "max_tokens": max_tokens, + "context_window_tokens": self.config.context_window_tokens, "reference_files": task_info.get("reference_files") if self.config.task == "gdpval" else None, "reference_file_urls": task_info.get("reference_file_urls") if self.config.task == "gdpval" else None, "exec_provider_class": exec_provider_class, @@ -1224,7 +1237,7 @@ async def responses( "completion_token_buffer": self.config.completion_token_buffer, "top_p": getattr(body, "top_p", None) or self.config.top_p, "enable_thinking": self.config.enable_thinking, - "max_completion_tokens_cap": self.config.max_completion_tokens_cap, + "max_completion_tokens_cap": max_completion_tokens_cap, "min_completion_tokens": self.config.min_completion_tokens, "prompt_estimator_truncate_history_thinking": (self.config.prompt_estimator_truncate_history_thinking), "truncation_recovery": self.config.truncation_recovery, diff --git a/responses_api_agents/stirrup_agent/tests/test_app.py b/responses_api_agents/stirrup_agent/tests/test_app.py index 983b41917d4..8e75b4514e5 100644 --- a/responses_api_agents/stirrup_agent/tests/test_app.py +++ b/responses_api_agents/stirrup_agent/tests/test_app.py @@ -134,6 +134,7 @@ def test_sanity(self) -> None: # to its larger floor and matching template semantics in its benchmark # config below. assert config.min_completion_tokens == 1024 + assert config.context_window_tokens == 262144 assert config.prompt_estimator_truncate_history_thinking is None assert config.min_compaction_summary_words == 1 @@ -143,9 +144,49 @@ def test_gdpval_config_opts_into_safe_context_budget(self) -> None: agent = config.gdpval_stirrup_agent.responses_api_agents.stirrup_agent assert agent.min_completion_tokens == 8192 + assert agent.context_window_tokens == 262144 assert "prompt_estimator_truncate_history_thinking" not in agent assert agent.min_compaction_summary_words == 50 + @pytest.mark.asyncio + async def test_request_output_cap_does_not_replace_model_context_window(self) -> None: + config = _make_config() + config.context_window_tokens = 262144 + config.max_completion_tokens_cap = 64000 + wrapper = StirrupAgentWrapper(config=config, server_client=MagicMock(spec=ServerClient)) + body = NeMoGymResponseCreateParamsNonStreaming( + input="ignored", + model="policy", + max_output_tokens=8192, + metadata={"task_id": "task-1", "prompt": "do the thing"}, + ) + captured = {} + + async def completed_rollout(): + return { + "input_items": [], + "output_items": [], + "deliverable_text": "", + "elapsed_seconds": 0, + } + + def launch(params): + captured.update(params) + return completed_rollout() + + with ( + patch.object(StirrupAgentWrapper, "resolve_model_base_url", return_value="http://policy.invalid/v1"), + patch.object(wrapper.task_strategy, "get_exec_provider", return_value=None), + patch( + "responses_api_agents.stirrup_agent.app.run_stirrup_agent_remote.remote", + side_effect=launch, + ), + ): + await wrapper.responses(body) + + assert captured["context_window_tokens"] == 262144 + assert captured["max_completion_tokens_cap"] == 8192 + def test_output_history_preserves_nemo_user_tool_results(self) -> None: """Run-history export should keep NeMo user-role tool results as tool outputs.""" history = [ diff --git a/responses_api_agents/stirrup_agent/tests/test_nemo_client.py b/responses_api_agents/stirrup_agent/tests/test_nemo_client.py index 24287d77b28..2b2690aea0d 100644 --- a/responses_api_agents/stirrup_agent/tests/test_nemo_client.py +++ b/responses_api_agents/stirrup_agent/tests/test_nemo_client.py @@ -517,6 +517,27 @@ async def test_no_estimated_context_remaining_raises_before_dispatch() -> None: fake_create.assert_not_awaited() +@pytest.mark.asyncio +async def test_output_cap_smaller_than_prompt_does_not_shrink_context_window() -> None: + client = DynamicMaxTokensChatCompletionsClient( + model="m", + max_tokens=262_144, + base_url="http://test", + api_key="k", + min_completion_tokens=8192, + max_completion_tokens_cap=8192, + ) + client._count_input_tokens_with_confidence = MagicMock(return_value=(10_000, True)) + client._tokenizer = object() + fake_create = AsyncMock(return_value=_make_response()) + client._client = MagicMock() + client._client.chat.completions.create = fake_create + + await client.generate([UserMessage(content="hi")], tools={}) + + assert fake_create.await_args.kwargs["max_completion_tokens"] == 8192 + + @pytest.mark.asyncio async def test_fallback_overestimate_still_dispatches_configured_floor() -> None: client = DynamicMaxTokensChatCompletionsClient(