diff --git a/areal/experimental/openai/client.py b/areal/experimental/openai/client.py
index 21aeaf77f0..f415ac1e01 100644
--- a/areal/experimental/openai/client.py
+++ b/areal/experimental/openai/client.py
@@ -57,6 +57,10 @@
from areal.api import ModelRequest, ModelResponse
from areal.api.cli_args import GenerationHyperparameters
from areal.experimental.openai.cache import InteractionCache
+from areal.experimental.openai.prompt_renderer import (
+ IncrementalPromptRenderer,
+ tools_signature,
+)
from areal.experimental.openai.tool_call_parser import process_tool_calls
from areal.experimental.openai.types import InteractionWithTokenLogpReward
from areal.utils import logging
@@ -672,6 +676,26 @@ def _concat_prompt_token_ids_with_parent(
all_message_list = _parse_tool_call_arguments(all_message_list)
if full_prompt_token_ids is None:
+ if (
+ parent is not None
+ and parent.output_message_list is not None
+ and IncrementalPromptRenderer.is_supported(
+ tokenizer,
+ tools=tools,
+ chat_template_kwargs=extra_body.get("chat_template_kwargs", {}),
+ )
+ ):
+ child_tokens = IncrementalPromptRenderer.render_concat_child_tokens(
+ tokenizer,
+ parent_output_messages=parent.output_message_list,
+ message_list=message_list,
+ tools=tools,
+ chat_template_kwargs=extra_body.get("chat_template_kwargs", {}),
+ )
+ if child_tokens is not None:
+ prompt_token_ids = parent_tokens + child_tokens
+ return prompt_token_ids, len(parent_tokens) - 1, len(parent_tokens)
+
all_tokens = apply_chat_template(
tokenizer,
all_message_list,
@@ -713,6 +737,7 @@ async def _prepare_prompt(
tools: Iterable[ChatCompletionToolParam] | None,
extra_body: Body,
require_multimodal_processor: bool = False,
+ interaction: InteractionWithTokenLogpReward | None = None,
) -> _PreparedPrompt:
"""Prepare text or multimodal prompt data for one agent interaction."""
chat_template_kwargs = extra_body.get("chat_template_kwargs", {})
@@ -734,18 +759,97 @@ async def _prepare_prompt(
)
if chat_template_type == "hf":
- input_ids = (
- processed_prompt.input_ids
- if processed_prompt is not None
- else apply_chat_template(
+ if processed_prompt is not None:
+ input_ids = processed_prompt.input_ids
+ elif (
+ processor is None
+ and parent is not None
+ and parent.messages
+ and len(tokenizer_messages) > len(parent.messages)
+ and IncrementalPromptRenderer.is_supported(
+ tokenizer, tools=tools, chat_template_kwargs=chat_template_kwargs
+ )
+ ):
+ delta_messages = tokenizer_messages[len(parent.messages) :]
+ current_tools_signature = tools_signature(tools)
+ parent_base = parent.prompt_base_token_ids
+ parent_tools_signature = parent.prompt_tools_signature
+ if parent_base is None:
+ parent_base = apply_chat_template(
+ tokenizer,
+ parent.messages,
+ tools=tools,
+ add_generation_prompt=False,
+ tokenize=True,
+ **chat_template_kwargs,
+ )
+ parent.prompt_base_token_ids = parent_base
+ # Re-rendered here with the current tool set, so it is up to date.
+ parent_tools_signature = current_tools_signature
+ parent.prompt_tools_signature = parent_tools_signature
+ rendered = IncrementalPromptRenderer.render_incremental(
tokenizer,
- tokenizer_messages,
+ parent_base,
+ delta_messages,
tools=tools,
- add_generation_prompt=True,
- tokenize=True,
- **chat_template_kwargs,
+ chat_template_kwargs=chat_template_kwargs,
+ parent_tools_signature=parent_tools_signature,
+ parent_messages=parent.messages,
)
- )
+ if rendered is not None:
+ input_ids, new_base = rendered
+ if interaction is not None:
+ interaction.prompt_base_token_ids = new_base
+ interaction.prompt_tools_signature = current_tools_signature
+ interaction.prompt_token_ids = list(input_ids)
+ else:
+ input_ids = apply_chat_template(
+ tokenizer,
+ tokenizer_messages,
+ tools=tools,
+ add_generation_prompt=True,
+ tokenize=True,
+ **chat_template_kwargs,
+ )
+ if interaction is not None:
+ interaction.prompt_token_ids = list(input_ids)
+ else:
+ # Rendering the base prefix costs an extra chat-template pass, so only
+ # pay for it when a later turn could actually reuse it.
+ if (
+ processor is None
+ and IncrementalPromptRenderer.is_supported(
+ tokenizer, tools=tools, chat_template_kwargs=chat_template_kwargs
+ )
+ and IncrementalPromptRenderer.is_history_safe(
+ tokenizer,
+ tokenizer_messages,
+ tools=tools,
+ chat_template_kwargs=chat_template_kwargs,
+ )
+ ):
+ input_ids, base_ids = IncrementalPromptRenderer.render_initial(
+ tokenizer,
+ tokenizer_messages,
+ tools=tools,
+ chat_template_kwargs=chat_template_kwargs,
+ )
+ if interaction is not None:
+ interaction.prompt_base_token_ids = base_ids
+ interaction.prompt_tools_signature = tools_signature(tools)
+ interaction.prompt_token_ids = list(input_ids)
+ else:
+ input_ids = apply_chat_template(
+ tokenizer,
+ tokenizer_messages,
+ tools=tools,
+ add_generation_prompt=True,
+ tokenize=True,
+ **chat_template_kwargs,
+ )
+ if interaction is not None:
+ interaction.prompt_token_ids = list(input_ids)
+
if processor is None:
return _PreparedPrompt(input_ids=input_ids)
return _PreparedPrompt(
@@ -1051,6 +1155,7 @@ async def create(
tools=tools_list,
extra_body=extra_body,
require_multimodal_processor=self.require_multimodal_processor,
+ interaction=interaction,
)
prompt_token_ids = prepared_prompt.input_ids
if interaction is not None and self.processor is not None:
@@ -1508,6 +1613,7 @@ async def create(
tools=tools_list,
extra_body=extra_body,
require_multimodal_processor=self.require_multimodal_processor,
+ interaction=interaction,
)
prompt_token_ids = prepared_prompt.input_ids
if self.processor is not None:
diff --git a/areal/experimental/openai/prompt_renderer.py b/areal/experimental/openai/prompt_renderer.py
new file mode 100644
index 0000000000..2ad3286f69
--- /dev/null
+++ b/areal/experimental/openai/prompt_renderer.py
@@ -0,0 +1,468 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Incremental prompt rendering for multi-turn agent rollout.
+
+In multi-turn tool-using rollouts (e.g., coding/search agents), re-running
+Jinja2 chat templates and tokenization over the full message history on every
+turn produces O(N^2) cumulative message processing over N turns.
+
+This module provides incremental prompt rendering:
+1. Turn 1 renders the full prompt and caches the base token prefix (without the
+ final generation prompt).
+2. Turns 2..N render only the newly appended delta messages against a bounded
+ synthetic context, appending the resulting token slice to the parent's base
+ prefix.
+3. Automatically probes tokenizer template capability on first use to ensure
+ 100% token-for-token mathematical identity with canonical full-history
+ rendering, safely falling back to full-history rendering for dynamic or
+ unsupported templates.
+"""
+
+from __future__ import annotations
+
+import json
+import threading
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Any
+
+from openai.types.chat import ChatCompletionToolParam
+
+from areal.utils import logging
+from areal.utils.hf_utils import apply_chat_template
+
+if TYPE_CHECKING:
+ from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
+
+logger = logging.getLogger("PromptRenderer")
+
+
+def _find_kth(lst: list[int], val: int, k: int) -> int:
+ """Find the index of the k-th (1-indexed) occurrence of val in lst."""
+ count = 0
+ for idx, item in enumerate(lst):
+ if item == val:
+ count += 1
+ if count == k:
+ return idx
+ return -1
+
+
+_THINK_START = ""
+_THINK_END = ""
+
+
+def tools_signature(tools: Iterable[ChatCompletionToolParam] | None) -> str:
+ """Return a stable signature for the tool set rendered into a prompt prefix."""
+ if not tools:
+ return ""
+ try:
+ return json.dumps(list(tools), sort_keys=True, default=str)
+ except (TypeError, ValueError):
+ return repr(list(tools))
+
+
+def contains_reasoning(message: dict[str, Any]) -> bool:
+ """Check whether a message carries an inline reasoning block."""
+ content = message.get("content")
+ return isinstance(content, str) and _THINK_END in content
+
+
+def has_superseded_reasoning(messages: Iterable[dict[str, Any]] | None) -> bool:
+ """Check whether reasoning appears before the final user turn of a history.
+
+ Templates such as Qwen3's keep the reasoning of the current turn but drop it
+ from every turn preceding the last user message. Such a history cannot be
+ served from an append-only token cache, because appending the new user turn
+ is supposed to remove tokens that are already in the cached prefix.
+ """
+ if not messages:
+ return False
+ message_list = list(messages)
+ last_user_idx = -1
+ for idx, message in enumerate(message_list):
+ if message.get("role") == "user":
+ last_user_idx = idx
+ if last_user_idx <= 0:
+ return False
+ return any(contains_reasoning(m) for m in message_list[:last_user_idx])
+
+
+class IncrementalPromptRenderer:
+ """Renders multi-turn agent prompts incrementally with token parity guarantees."""
+
+ _capability_cache: dict[tuple[Any, ...], tuple[bool, bool]] = {}
+ _dummy_d0_cache: dict[tuple[Any, ...], int] = {}
+ _lock = threading.Lock()
+
+ @classmethod
+ def _get_cache_key(
+ cls,
+ tokenizer: PreTrainedTokenizerFast,
+ chat_template_kwargs: dict[str, Any] | None,
+ ) -> tuple[Any, ...]:
+ kw_items = (
+ tuple(sorted((k, str(v)) for k, v in chat_template_kwargs.items()))
+ if chat_template_kwargs
+ else ()
+ )
+ return (
+ id(tokenizer),
+ getattr(tokenizer, "name_or_path", None),
+ kw_items,
+ )
+
+ @classmethod
+ def is_supported(
+ cls,
+ tokenizer: PreTrainedTokenizerFast,
+ tools: Iterable[ChatCompletionToolParam] | None = None,
+ chat_template_kwargs: dict[str, Any] | None = None,
+ ) -> bool:
+ """Check whether the tokenizer chat template supports incremental delta rendering."""
+ return cls._get_capability(tokenizer, tools, chat_template_kwargs)[0]
+
+ @classmethod
+ def is_history_safe(
+ cls,
+ tokenizer: PreTrainedTokenizerFast,
+ messages: Iterable[dict[str, Any]] | None,
+ tools: Iterable[ChatCompletionToolParam] | None = None,
+ chat_template_kwargs: dict[str, Any] | None = None,
+ ) -> bool:
+ """Check whether ``messages`` can be served from an append-only token cache."""
+ if not has_superseded_reasoning(messages):
+ return True
+ return cls._get_capability(tokenizer, tools, chat_template_kwargs)[1]
+
+ @classmethod
+ def _get_capability(
+ cls,
+ tokenizer: PreTrainedTokenizerFast,
+ tools: Iterable[ChatCompletionToolParam] | None,
+ chat_template_kwargs: dict[str, Any] | None,
+ ) -> tuple[bool, bool]:
+ """Return (delta rendering supported, reasoning history safe) for a tokenizer."""
+ if not hasattr(tokenizer, "chat_template") or not tokenizer.chat_template:
+ return False, False
+
+ key = cls._get_cache_key(tokenizer, chat_template_kwargs)
+ with cls._lock:
+ if key in cls._capability_cache:
+ return cls._capability_cache[key]
+
+ # Probe capability with synthetic sequences
+ capability = cls._probe_capability(tokenizer, tools, chat_template_kwargs)
+ with cls._lock:
+ cls._capability_cache[key] = capability
+
+ supported, reasoning_safe = capability
+ if supported:
+ logger.debug(
+ "Incremental prompt rendering verified and enabled for tokenizer: %s "
+ "(reasoning history safe: %s)",
+ getattr(tokenizer, "name_or_path", type(tokenizer).__name__),
+ reasoning_safe,
+ )
+ else:
+ logger.debug(
+ "Incremental prompt rendering not supported for tokenizer: %s; using full fallback.",
+ getattr(tokenizer, "name_or_path", type(tokenizer).__name__),
+ )
+ return capability
+
+ @classmethod
+ def _probe_capability(
+ cls,
+ tokenizer: PreTrainedTokenizerFast,
+ tools: Iterable[ChatCompletionToolParam] | None,
+ chat_template_kwargs: dict[str, Any] | None,
+ ) -> tuple[bool, bool]:
+ """Run probes to verify token-for-token equality between incremental and full rendering.
+
+ The first probe covers an append-only tool-calling delta. The second probe
+ covers a cached prefix that already contains a reasoning block, which some
+ templates rewrite once a later turn is appended.
+ """
+ kwargs = chat_template_kwargs or {}
+ try:
+ m1 = [{"role": "user", "content": "probe user query"}]
+ delta = [
+ {
+ "role": "assistant",
+ "content": "probe response",
+ "tool_calls": [
+ {
+ "id": "call_probe_1",
+ "type": "function",
+ "function": {
+ "name": "probe_tool",
+ "arguments": '{"param": "val"}',
+ },
+ }
+ ],
+ },
+ {
+ "role": "tool",
+ "tool_call_id": "call_probe_1",
+ "name": "probe_tool",
+ "content": "probe result",
+ },
+ ]
+ full_2 = apply_chat_template(
+ tokenizer,
+ m1 + delta,
+ tools=tools,
+ add_generation_prompt=True,
+ tokenize=True,
+ **kwargs,
+ )
+ base_1 = apply_chat_template(
+ tokenizer,
+ m1,
+ tools=tools,
+ add_generation_prompt=False,
+ tokenize=True,
+ **kwargs,
+ )
+ dummy = [{"role": "user", "content": "x"}]
+ d0 = apply_chat_template(
+ tokenizer,
+ dummy,
+ add_generation_prompt=False,
+ tokenize=True,
+ **kwargs,
+ )
+ d_gen = apply_chat_template(
+ tokenizer,
+ dummy + delta,
+ add_generation_prompt=True,
+ tokenize=True,
+ **kwargs,
+ )
+ if not isinstance(full_2, list) or not isinstance(base_1, list):
+ return False, False
+ incr_2 = base_1 + d_gen[len(d0) :]
+ supported = full_2 == incr_2
+ except Exception as e:
+ logger.debug("PromptRenderer probe failed with error: %s", e)
+ return False, False
+
+ if not supported:
+ return False, False
+ return True, cls._probe_reasoning_history(
+ tokenizer, tools, chat_template_kwargs
+ )
+
+ @classmethod
+ def _probe_reasoning_history(
+ cls,
+ tokenizer: PreTrainedTokenizerFast,
+ tools: Iterable[ChatCompletionToolParam] | None,
+ chat_template_kwargs: dict[str, Any] | None,
+ ) -> bool:
+ """Verify that a cached prefix containing reasoning survives appending a turn."""
+ kwargs = chat_template_kwargs or {}
+ try:
+ history = [
+ {"role": "user", "content": "probe user query"},
+ {
+ "role": "assistant",
+ "content": f"{_THINK_START}probe reasoning{_THINK_END}probe answer",
+ },
+ ]
+ delta = [{"role": "user", "content": "probe follow-up"}]
+ base = apply_chat_template(
+ tokenizer,
+ history,
+ tools=tools,
+ add_generation_prompt=False,
+ tokenize=True,
+ **kwargs,
+ )
+ full = apply_chat_template(
+ tokenizer,
+ history + delta,
+ tools=tools,
+ add_generation_prompt=True,
+ tokenize=True,
+ **kwargs,
+ )
+ dummy = [{"role": "user", "content": "x"}]
+ d0 = apply_chat_template(
+ tokenizer,
+ dummy,
+ add_generation_prompt=False,
+ tokenize=True,
+ **kwargs,
+ )
+ d_gen = apply_chat_template(
+ tokenizer,
+ dummy + delta,
+ add_generation_prompt=True,
+ tokenize=True,
+ **kwargs,
+ )
+ if not isinstance(full, list) or not isinstance(base, list):
+ return False
+ return full == base + d_gen[len(d0) :]
+ except Exception as e:
+ logger.debug("PromptRenderer reasoning probe failed with error: %s", e)
+ return False
+
+ @classmethod
+ def _get_dummy_d0_len(
+ cls,
+ tokenizer: PreTrainedTokenizerFast,
+ chat_template_kwargs: dict[str, Any] | None,
+ ) -> int:
+ key = cls._get_cache_key(tokenizer, chat_template_kwargs)
+ with cls._lock:
+ if key in cls._dummy_d0_cache:
+ return cls._dummy_d0_cache[key]
+
+ dummy = [{"role": "user", "content": "x"}]
+ d0 = apply_chat_template(
+ tokenizer,
+ dummy,
+ add_generation_prompt=False,
+ tokenize=True,
+ **(chat_template_kwargs or {}),
+ )
+ d0_len = len(d0)
+ with cls._lock:
+ cls._dummy_d0_cache[key] = d0_len
+ return d0_len
+
+ @classmethod
+ def render_initial(
+ cls,
+ tokenizer: PreTrainedTokenizerFast,
+ messages: list[dict[str, Any]],
+ tools: Iterable[ChatCompletionToolParam] | None = None,
+ chat_template_kwargs: dict[str, Any] | None = None,
+ ) -> tuple[list[int], list[int]]:
+ """Render the initial turn's prompt tokens and base prefix tokens."""
+ kwargs = chat_template_kwargs or {}
+ prompt_token_ids = apply_chat_template(
+ tokenizer,
+ messages,
+ tools=tools,
+ add_generation_prompt=True,
+ tokenize=True,
+ **kwargs,
+ )
+ prompt_base_token_ids = apply_chat_template(
+ tokenizer,
+ messages,
+ tools=tools,
+ add_generation_prompt=False,
+ tokenize=True,
+ **kwargs,
+ )
+ return prompt_token_ids, prompt_base_token_ids
+
+ @classmethod
+ def render_incremental(
+ cls,
+ tokenizer: PreTrainedTokenizerFast,
+ parent_base_token_ids: list[int],
+ delta_messages: list[dict[str, Any]],
+ tools: Iterable[ChatCompletionToolParam] | None = None,
+ chat_template_kwargs: dict[str, Any] | None = None,
+ parent_tools_signature: str | None = "",
+ parent_messages: list[dict[str, Any]] | None = None,
+ ) -> tuple[list[int], list[int]] | None:
+ """Render prompt tokens for delta messages appended to parent base tokens.
+
+ Only the delta is rendered, so the tool definitions baked into the parent
+ prefix are reused as-is. ``parent_tools_signature`` records the tools that
+ prefix was built with; when the current turn declares a different tool set
+ the prefix is stale and rendering is refused.
+
+ ``parent_messages`` are the messages the prefix was built from. Together
+ with ``delta_messages`` they are checked against
+ :meth:`is_history_safe`, so a prefix whose reasoning blocks a later turn
+ would strip is never produced or consumed.
+
+ Returns (prompt_token_ids, new_base_token_ids) or None on failure.
+ """
+ if not delta_messages:
+ return None
+
+ if (parent_tools_signature or "") != tools_signature(tools):
+ logger.debug(
+ "Tool set changed between turns; skipping incremental prompt rendering."
+ )
+ return None
+
+ if not cls.is_history_safe(
+ tokenizer,
+ list(parent_messages or []) + delta_messages,
+ tools=tools,
+ chat_template_kwargs=chat_template_kwargs,
+ ):
+ logger.debug(
+ "Chat template rewrites reasoning history; skipping incremental "
+ "prompt rendering."
+ )
+ return None
+
+ kwargs = chat_template_kwargs or {}
+ try:
+ dummy = [{"role": "user", "content": "x"}]
+ d0_len = cls._get_dummy_d0_len(tokenizer, kwargs)
+ d_gen = apply_chat_template(
+ tokenizer,
+ dummy + delta_messages,
+ add_generation_prompt=True,
+ tokenize=True,
+ **kwargs,
+ )
+ d_no_gen = apply_chat_template(
+ tokenizer,
+ dummy + delta_messages,
+ add_generation_prompt=False,
+ tokenize=True,
+ **kwargs,
+ )
+ if not isinstance(d_gen, list) or not isinstance(d_no_gen, list):
+ return None
+ prompt_token_ids = parent_base_token_ids + d_gen[d0_len:]
+ new_base_token_ids = parent_base_token_ids + d_no_gen[d0_len:]
+ return prompt_token_ids, new_base_token_ids
+ except Exception as e:
+ logger.debug("render_incremental failed: %s; falling back", e)
+ return None
+
+ @classmethod
+ def render_concat_child_tokens(
+ cls,
+ tokenizer: PreTrainedTokenizerFast,
+ parent_output_messages: list[dict[str, Any]],
+ message_list: list[dict[str, Any]],
+ tools: Iterable[ChatCompletionToolParam] | None = None,
+ chat_template_kwargs: dict[str, Any] | None = None,
+ ) -> list[int] | None:
+ """Render child tokens for concat mode from a bounded synthetic context."""
+ kwargs = chat_template_kwargs or {}
+ try:
+ dummy = [{"role": "user", "content": "x"}]
+ d_delta = dummy + parent_output_messages + message_list
+ d_gen = apply_chat_template(
+ tokenizer,
+ d_delta,
+ add_generation_prompt=True,
+ tokenize=True,
+ **kwargs,
+ )
+ if not isinstance(d_gen, list):
+ return None
+ eos_token_id = tokenizer.eos_token_id
+ dummy_parent_eos_count = len(dummy) + len(parent_output_messages)
+ child_truncate_idx = _find_kth(d_gen, eos_token_id, dummy_parent_eos_count)
+ if child_truncate_idx == -1 or child_truncate_idx + 1 >= len(d_gen):
+ return None
+ return d_gen[child_truncate_idx + 1 :]
+ except Exception as e:
+ logger.debug("render_concat_child_tokens failed: %s; falling back", e)
+ return None
diff --git a/areal/experimental/openai/types.py b/areal/experimental/openai/types.py
index 2ca1127021..30564e5c3b 100644
--- a/areal/experimental/openai/types.py
+++ b/areal/experimental/openai/types.py
@@ -45,8 +45,12 @@ class InteractionWithTokenLogpReward:
chat_template_type: str = "hf"
_cache: dict[str, Any] | None = None
- # Multimodal training data prepared from the complete prompt for this turn.
+ # Prompt token cache
prompt_token_ids: list[int] | None = None
+ prompt_base_token_ids: list[int] | None = None
+ # Tool set the cached base prefix was rendered with; see
+ # IncrementalPromptRenderer.render_incremental.
+ prompt_tools_signature: str | None = None
mm_token_type_ids: list[int] | None = None
multi_modal_input: dict[str, torch.Tensor] | None = None
diff --git a/tests/experimental/openai/test_prompt_renderer.py b/tests/experimental/openai/test_prompt_renderer.py
new file mode 100644
index 0000000000..4933f80c9d
--- /dev/null
+++ b/tests/experimental/openai/test_prompt_renderer.py
@@ -0,0 +1,791 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+from openai.types.chat import ChatCompletionToolParam
+
+from tests.utils import get_model_path
+
+from areal.api import ModelResponse
+from areal.experimental.openai.client import (
+ _concat_prompt_token_ids_with_parent,
+ _prepare_prompt,
+)
+from areal.experimental.openai.prompt_renderer import (
+ IncrementalPromptRenderer,
+ _find_kth,
+ has_superseded_reasoning,
+ tools_signature,
+)
+from areal.experimental.openai.types import InteractionWithTokenLogpReward
+from areal.utils.hf_utils import apply_chat_template, load_hf_tokenizer
+
+QWEN3_MODEL_PATH = "Qwen/Qwen3-0.6B"
+LOCAL_QWEN3_PATH = "/storage/openpsi/models/Qwen__Qwen3-0.6B"
+
+QWEN25_MODEL_PATH = "Qwen/Qwen2.5-0.5B-Instruct"
+LOCAL_QWEN25_PATH = "/storage/openpsi/models/Qwen__Qwen2.5-0.5B-Instruct"
+
+WEATHER_TOOL: ChatCompletionToolParam = {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "Get current weather in a given location",
+ "parameters": {
+ "type": "object",
+ "properties": {"location": {"type": "string", "description": "City name"}},
+ "required": ["location"],
+ },
+ },
+}
+
+CALCULATOR_TOOL: ChatCompletionToolParam = {
+ "type": "function",
+ "function": {
+ "name": "calculate",
+ "description": "Evaluate math expression",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "expr": {"type": "string", "description": "Math expression"}
+ },
+ "required": ["expr"],
+ },
+ },
+}
+
+
+@pytest.fixture(scope="module")
+def qwen3_tokenizer():
+ return load_hf_tokenizer(get_model_path(LOCAL_QWEN3_PATH, QWEN3_MODEL_PATH))
+
+
+@pytest.fixture(scope="module")
+def qwen25_tokenizer():
+ return load_hf_tokenizer(get_model_path(LOCAL_QWEN25_PATH, QWEN25_MODEL_PATH))
+
+
+class TestPromptRendererCapability:
+ def test_find_kth_helper(self):
+ lst = [1, 2, 3, 2, 4, 2, 5]
+ assert _find_kth(lst, 2, 1) == 1
+ assert _find_kth(lst, 2, 2) == 3
+ assert _find_kth(lst, 2, 3) == 5
+ assert _find_kth(lst, 2, 4) == -1
+ assert _find_kth(lst, 99, 1) == -1
+
+ def test_supported_tokenizers(self, qwen3_tokenizer, qwen25_tokenizer):
+ assert IncrementalPromptRenderer.is_supported(
+ qwen3_tokenizer, tools=[WEATHER_TOOL]
+ )
+ assert IncrementalPromptRenderer.is_supported(
+ qwen25_tokenizer, tools=[WEATHER_TOOL]
+ )
+
+ def test_tokenizer_without_chat_template(self):
+ class DummyTokenizer:
+ chat_template = None
+ name_or_path = "dummy"
+
+ dummy = DummyTokenizer()
+ assert not IncrementalPromptRenderer.is_supported(dummy) # type: ignore[arg-type]
+
+
+class TestIncrementalPromptRenderingParity:
+ def test_render_initial(self, qwen3_tokenizer):
+ messages = [{"role": "user", "content": "What is the weather in Paris?"}]
+ prompt_ids, base_ids = IncrementalPromptRenderer.render_initial(
+ qwen3_tokenizer, messages, tools=[WEATHER_TOOL]
+ )
+ canonical_full = apply_chat_template(
+ qwen3_tokenizer,
+ messages,
+ tools=[WEATHER_TOOL],
+ add_generation_prompt=True,
+ tokenize=True,
+ )
+ canonical_base = apply_chat_template(
+ qwen3_tokenizer,
+ messages,
+ tools=[WEATHER_TOOL],
+ add_generation_prompt=False,
+ tokenize=True,
+ )
+ assert prompt_ids == canonical_full
+ assert base_ids == canonical_base
+
+ def test_single_tool_call_round(self, qwen3_tokenizer):
+ messages = [{"role": "user", "content": "What is the weather in Paris?"}]
+ _, base_ids = IncrementalPromptRenderer.render_initial(
+ qwen3_tokenizer, messages, tools=[WEATHER_TOOL]
+ )
+
+ asst_msg = {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "c1",
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"location": "Paris"}',
+ },
+ }
+ ],
+ }
+ tool_msg = {
+ "role": "tool",
+ "tool_call_id": "c1",
+ "name": "get_weather",
+ "content": '{"temp": "20C"}',
+ }
+ delta = [asst_msg, tool_msg]
+ messages.extend(delta)
+
+ # Canonical full render
+ canonical_prompt = apply_chat_template(
+ qwen3_tokenizer,
+ messages,
+ tools=[WEATHER_TOOL],
+ add_generation_prompt=True,
+ tokenize=True,
+ )
+
+ # Incremental render
+ rendered = IncrementalPromptRenderer.render_incremental(
+ qwen3_tokenizer,
+ base_ids,
+ delta,
+ tools=[WEATHER_TOOL],
+ parent_tools_signature=tools_signature([WEATHER_TOOL]),
+ parent_messages=messages[: -len(delta)],
+ )
+ assert rendered is not None
+ incr_prompt, _ = rendered
+
+ assert incr_prompt == canonical_prompt
+
+ def test_multi_turn_tool_sequence_20_turns(self, qwen3_tokenizer):
+ messages = [{"role": "user", "content": "Start multi-turn episode"}]
+ _, current_base = IncrementalPromptRenderer.render_initial(
+ qwen3_tokenizer, messages, tools=[WEATHER_TOOL, CALCULATOR_TOOL]
+ )
+
+ for turn in range(1, 20):
+ asst_msg = {
+ "role": "assistant",
+ "content": f"Step {turn} analysis",
+ "tool_calls": [
+ {
+ "id": f"call_{turn}",
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": f'{{"location": "City_{turn}"}}',
+ },
+ }
+ ],
+ }
+ tool_msg = {
+ "role": "tool",
+ "tool_call_id": f"call_{turn}",
+ "name": "get_weather",
+ "content": f'{{"temp": "{20 + turn}C"}}',
+ }
+ delta = [asst_msg, tool_msg]
+ messages.extend(delta)
+
+ # Canonical full history
+ canonical_prompt = apply_chat_template(
+ qwen3_tokenizer,
+ messages,
+ tools=[WEATHER_TOOL, CALCULATOR_TOOL],
+ add_generation_prompt=True,
+ tokenize=True,
+ )
+
+ # Incremental
+ rendered = IncrementalPromptRenderer.render_incremental(
+ qwen3_tokenizer,
+ current_base,
+ delta,
+ tools=[WEATHER_TOOL, CALCULATOR_TOOL],
+ parent_tools_signature=tools_signature([WEATHER_TOOL, CALCULATOR_TOOL]),
+ parent_messages=messages[: -len(delta)],
+ )
+ assert rendered is not None
+ incr_prompt, current_base = rendered
+
+ assert incr_prompt == canonical_prompt, f"Mismatch at turn {turn}"
+
+ def test_parallel_tool_calls(self, qwen3_tokenizer):
+ messages = [{"role": "user", "content": "Compare Paris and London"}]
+ _, base_ids = IncrementalPromptRenderer.render_initial(
+ qwen3_tokenizer, messages, tools=[WEATHER_TOOL]
+ )
+
+ asst_msg = {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "c1",
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"location": "Paris"}',
+ },
+ },
+ {
+ "id": "c2",
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"location": "London"}',
+ },
+ },
+ ],
+ }
+ tool_1 = {
+ "role": "tool",
+ "tool_call_id": "c1",
+ "name": "get_weather",
+ "content": "20C",
+ }
+ tool_2 = {
+ "role": "tool",
+ "tool_call_id": "c2",
+ "name": "get_weather",
+ "content": "15C",
+ }
+ delta = [asst_msg, tool_1, tool_2]
+ messages.extend(delta)
+
+ canonical_prompt = apply_chat_template(
+ qwen3_tokenizer,
+ messages,
+ tools=[WEATHER_TOOL],
+ add_generation_prompt=True,
+ tokenize=True,
+ )
+
+ rendered = IncrementalPromptRenderer.render_incremental(
+ qwen3_tokenizer,
+ base_ids,
+ delta,
+ tools=[WEATHER_TOOL],
+ parent_tools_signature=tools_signature([WEATHER_TOOL]),
+ parent_messages=messages[: -len(delta)],
+ )
+ assert rendered is not None
+ incr_prompt, _ = rendered
+
+ assert incr_prompt == canonical_prompt
+
+ def test_conversational_and_custom_system_prompt(self, qwen25_tokenizer):
+ messages = [
+ {"role": "system", "content": "You are a specialized math assistant."},
+ {"role": "user", "content": "Solve 2+2"},
+ ]
+ _, base_ids = IncrementalPromptRenderer.render_initial(
+ qwen25_tokenizer, messages, tools=[CALCULATOR_TOOL]
+ )
+
+ delta1 = [
+ {"role": "assistant", "content": "2+2 equals 4."},
+ {"role": "user", "content": "Now compute 10 * 10"},
+ ]
+ messages.extend(delta1)
+
+ canonical_prompt = apply_chat_template(
+ qwen25_tokenizer,
+ messages,
+ tools=[CALCULATOR_TOOL],
+ add_generation_prompt=True,
+ tokenize=True,
+ )
+
+ rendered = IncrementalPromptRenderer.render_incremental(
+ qwen25_tokenizer,
+ base_ids,
+ delta1,
+ tools=[CALCULATOR_TOOL],
+ parent_tools_signature=tools_signature([CALCULATOR_TOOL]),
+ parent_messages=messages[: -len(delta1)],
+ )
+ assert rendered is not None
+ incr_prompt, _ = rendered
+
+ assert incr_prompt == canonical_prompt
+
+
+class TestIncrementalRenderingGuards:
+ """Regression tests for prefixes that an append-only cache cannot represent."""
+
+ def test_has_superseded_reasoning(self):
+ reasoning_turn = {"role": "assistant", "content": "ra"}
+ # No later user turn: nothing supersedes the reasoning.
+ assert not has_superseded_reasoning(
+ [{"role": "user", "content": "q"}, reasoning_turn]
+ )
+ # A later user turn supersedes it.
+ assert has_superseded_reasoning(
+ [
+ {"role": "user", "content": "q1"},
+ reasoning_turn,
+ {"role": "user", "content": "q2"},
+ ]
+ )
+ assert not has_superseded_reasoning([])
+
+ def test_reasoning_history_falls_back_for_qwen3(self, qwen3_tokenizer):
+ """Qwen3 drops reasoning before the last user turn, so the cache is stale."""
+ messages = [
+ {"role": "user", "content": "What is 2+2?"},
+ {"role": "assistant", "content": "simple sumIt is 4."},
+ ]
+ _, base_ids = IncrementalPromptRenderer.render_initial(
+ qwen3_tokenizer, messages
+ )
+ delta = [{"role": "user", "content": "And 3+3?"}]
+
+ assert not IncrementalPromptRenderer.is_history_safe(
+ qwen3_tokenizer, messages + delta
+ )
+ assert (
+ IncrementalPromptRenderer.render_incremental(
+ qwen3_tokenizer,
+ base_ids,
+ delta,
+ parent_messages=messages,
+ )
+ is None
+ )
+
+ # The cached prefix would have kept the reasoning the template removes.
+ canonical = apply_chat_template(
+ qwen3_tokenizer,
+ messages + delta,
+ add_generation_prompt=True,
+ tokenize=True,
+ )
+ stale = (
+ base_ids
+ + apply_chat_template(
+ qwen3_tokenizer,
+ [{"role": "user", "content": "x"}] + delta,
+ add_generation_prompt=True,
+ tokenize=True,
+ )[
+ len(
+ apply_chat_template(
+ qwen3_tokenizer,
+ [{"role": "user", "content": "x"}],
+ add_generation_prompt=False,
+ tokenize=True,
+ )
+ ) :
+ ]
+ )
+ assert stale != canonical
+
+ def test_reasoning_tool_loop_stays_incremental(self, qwen3_tokenizer):
+ """Tool-only deltas keep the current turn's reasoning, so the cache holds."""
+ messages = [{"role": "user", "content": "Weather in Paris?"}]
+ _, base_ids = IncrementalPromptRenderer.render_initial(
+ qwen3_tokenizer, messages, tools=[WEATHER_TOOL]
+ )
+ delta = [
+ {
+ "role": "assistant",
+ "content": "call the tool",
+ "tool_calls": [
+ {
+ "id": "c1",
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"location": "Paris"}',
+ },
+ }
+ ],
+ },
+ {
+ "role": "tool",
+ "tool_call_id": "c1",
+ "name": "get_weather",
+ "content": "20C",
+ },
+ ]
+ rendered = IncrementalPromptRenderer.render_incremental(
+ qwen3_tokenizer,
+ base_ids,
+ delta,
+ tools=[WEATHER_TOOL],
+ parent_tools_signature=tools_signature([WEATHER_TOOL]),
+ parent_messages=messages,
+ )
+ assert rendered is not None
+ assert rendered[0] == apply_chat_template(
+ qwen3_tokenizer,
+ messages + delta,
+ tools=[WEATHER_TOOL],
+ add_generation_prompt=True,
+ tokenize=True,
+ )
+
+ def test_reasoning_history_allowed_for_qwen25(self, qwen25_tokenizer):
+ """Qwen2.5 never rewrites history, so reasoning is safe to cache."""
+ messages = [
+ {"role": "user", "content": "What is 2+2?"},
+ {"role": "assistant", "content": "simple sumIt is 4."},
+ ]
+ _, base_ids = IncrementalPromptRenderer.render_initial(
+ qwen25_tokenizer, messages
+ )
+ delta = [{"role": "user", "content": "And 3+3?"}]
+
+ assert IncrementalPromptRenderer.is_history_safe(
+ qwen25_tokenizer, messages + delta
+ )
+ rendered = IncrementalPromptRenderer.render_incremental(
+ qwen25_tokenizer,
+ base_ids,
+ delta,
+ parent_messages=messages,
+ )
+ assert rendered is not None
+ assert rendered[0] == apply_chat_template(
+ qwen25_tokenizer,
+ messages + delta,
+ add_generation_prompt=True,
+ tokenize=True,
+ )
+
+ def test_tools_signature_is_order_stable(self):
+ assert tools_signature(None) == tools_signature([]) == ""
+ assert tools_signature([WEATHER_TOOL]) == tools_signature([WEATHER_TOOL])
+ assert tools_signature([WEATHER_TOOL]) != tools_signature(
+ [WEATHER_TOOL, CALCULATOR_TOOL]
+ )
+
+ @pytest.mark.parametrize(
+ "parent_tools,current_tools",
+ [
+ ([WEATHER_TOOL], [WEATHER_TOOL, CALCULATOR_TOOL]),
+ ([WEATHER_TOOL, CALCULATOR_TOOL], [WEATHER_TOOL]),
+ ([WEATHER_TOOL], None),
+ (None, [WEATHER_TOOL]),
+ ],
+ )
+ def test_changed_tools_fall_back(
+ self, qwen3_tokenizer, parent_tools, current_tools
+ ):
+ """A prefix built with one tool set cannot serve a turn declaring another."""
+ messages = [{"role": "user", "content": "Weather in Paris?"}]
+ _, base_ids = IncrementalPromptRenderer.render_initial(
+ qwen3_tokenizer, messages, tools=parent_tools
+ )
+ delta = [
+ {"role": "assistant", "content": "Let me check."},
+ {"role": "user", "content": "Also compute 2+2"},
+ ]
+ assert (
+ IncrementalPromptRenderer.render_incremental(
+ qwen3_tokenizer,
+ base_ids,
+ delta,
+ tools=current_tools,
+ parent_tools_signature=tools_signature(parent_tools),
+ parent_messages=messages,
+ )
+ is None
+ )
+
+ def test_unchanged_tools_stay_incremental(self, qwen3_tokenizer):
+ messages = [{"role": "user", "content": "Weather in Paris?"}]
+ _, base_ids = IncrementalPromptRenderer.render_initial(
+ qwen3_tokenizer, messages, tools=[WEATHER_TOOL]
+ )
+ delta = [
+ {"role": "assistant", "content": "Let me check."},
+ {"role": "user", "content": "Also, thanks"},
+ ]
+ rendered = IncrementalPromptRenderer.render_incremental(
+ qwen3_tokenizer,
+ base_ids,
+ delta,
+ tools=[WEATHER_TOOL],
+ parent_tools_signature=tools_signature([WEATHER_TOOL]),
+ parent_messages=messages,
+ )
+ assert rendered is not None
+ assert rendered[0] == apply_chat_template(
+ qwen3_tokenizer,
+ messages + delta,
+ tools=[WEATHER_TOOL],
+ add_generation_prompt=True,
+ tokenize=True,
+ )
+
+
+class TestIncrementalConcatPromptRendering:
+ def test_render_concat_child_tokens_identity(self, qwen3_tokenizer):
+ msg1 = [{"role": "user", "content": "Weather in Paris?"}]
+ t1 = apply_chat_template(
+ qwen3_tokenizer,
+ msg1,
+ tools=[WEATHER_TOOL],
+ add_generation_prompt=True,
+ tokenize=True,
+ )
+ eos_id = qwen3_tokenizer.eos_token_id
+ out_tokens = qwen3_tokenizer.encode(
+ '\n{"name": "get_weather", "arguments": {"location": "Paris"}}\n',
+ add_special_tokens=False,
+ ) + [eos_id]
+
+ resp = ModelResponse(
+ input_tokens=t1,
+ output_tokens=out_tokens,
+ output_logprobs=[0.0] * len(out_tokens),
+ output_versions=[0] * len(out_tokens),
+ stop_reason="stop",
+ tokenizer=qwen3_tokenizer,
+ )
+ asst = {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "c1",
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"location": "Paris"}',
+ },
+ }
+ ],
+ }
+ parent = InteractionWithTokenLogpReward(
+ messages=msg1,
+ output_message_list=[asst],
+ model_response=resp,
+ chat_template_type="concat",
+ )
+ tool_msg = [
+ {
+ "role": "tool",
+ "tool_call_id": "c1",
+ "name": "get_weather",
+ "content": "20C",
+ }
+ ]
+
+ # Full history concat
+ concat_full, _, _ = _concat_prompt_token_ids_with_parent(
+ message_list=tool_msg,
+ parent=parent,
+ tokenizer=qwen3_tokenizer,
+ tools=[WEATHER_TOOL],
+ )
+
+ # Incremental child tokens
+ child_tokens = IncrementalPromptRenderer.render_concat_child_tokens(
+ qwen3_tokenizer,
+ parent_output_messages=parent.output_message_list,
+ message_list=tool_msg,
+ tools=[WEATHER_TOOL],
+ )
+ assert child_tokens is not None
+ parent_tokens = (
+ parent.model_response.input_tokens
+ + parent.model_response.output_tokens_without_stop
+ + [eos_id]
+ )
+ concat_incr = parent_tokens + child_tokens
+
+ assert concat_incr == concat_full
+
+
+@pytest.mark.asyncio
+class TestPreparePromptIntegration:
+ async def test_prepare_prompt_hf_multi_turn_parity(self, qwen3_tokenizer):
+ messages = [{"role": "user", "content": "Weather query"}]
+ inter1 = InteractionWithTokenLogpReward(
+ messages=list(messages), chat_template_type="hf"
+ )
+ p1 = await _prepare_prompt(
+ tokenizer=qwen3_tokenizer,
+ processor=None,
+ tokenizer_messages=messages,
+ concat_messages=messages,
+ image_data=[],
+ parent=None,
+ chat_template_type="hf",
+ tools=[WEATHER_TOOL],
+ extra_body={},
+ interaction=inter1,
+ )
+ assert inter1.prompt_token_ids == p1.input_ids
+ assert inter1.prompt_base_token_ids is not None
+
+ asst = {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "c1",
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"location": "Paris"}',
+ },
+ }
+ ],
+ }
+ tool = {
+ "role": "tool",
+ "tool_call_id": "c1",
+ "name": "get_weather",
+ "content": "20C",
+ }
+ messages.extend([asst, tool])
+ inter2 = InteractionWithTokenLogpReward(
+ messages=list(messages), chat_template_type="hf", parent=inter1
+ )
+
+ p2 = await _prepare_prompt(
+ tokenizer=qwen3_tokenizer,
+ processor=None,
+ tokenizer_messages=messages,
+ concat_messages=[tool],
+ image_data=[],
+ parent=inter1,
+ chat_template_type="hf",
+ tools=[WEATHER_TOOL],
+ extra_body={},
+ interaction=inter2,
+ )
+
+ canonical = apply_chat_template(
+ qwen3_tokenizer,
+ messages,
+ tools=[WEATHER_TOOL],
+ add_generation_prompt=True,
+ tokenize=True,
+ )
+ assert p2.input_ids == canonical
+ assert inter2.prompt_token_ids == canonical
+
+ async def test_prepare_prompt_hf_fallback_without_parent_base_ids(
+ self, qwen3_tokenizer
+ ):
+ messages = [{"role": "user", "content": "Weather query"}]
+ inter1 = InteractionWithTokenLogpReward(
+ messages=list(messages), chat_template_type="hf"
+ )
+ # Simulate legacy interaction with prompt_base_token_ids as None
+ inter1.prompt_base_token_ids = None
+
+ asst = {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "c1",
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"location": "Paris"}',
+ },
+ }
+ ],
+ }
+ tool = {
+ "role": "tool",
+ "tool_call_id": "c1",
+ "name": "get_weather",
+ "content": "20C",
+ }
+ messages.extend([asst, tool])
+ inter2 = InteractionWithTokenLogpReward(
+ messages=list(messages), chat_template_type="hf", parent=inter1
+ )
+
+ p2 = await _prepare_prompt(
+ tokenizer=qwen3_tokenizer,
+ processor=None,
+ tokenizer_messages=messages,
+ concat_messages=[tool],
+ image_data=[],
+ parent=inter1,
+ chat_template_type="hf",
+ tools=[WEATHER_TOOL],
+ extra_body={},
+ interaction=inter2,
+ )
+
+ canonical = apply_chat_template(
+ qwen3_tokenizer,
+ messages,
+ tools=[WEATHER_TOOL],
+ add_generation_prompt=True,
+ tokenize=True,
+ )
+ assert p2.input_ids == canonical
+ assert inter2.prompt_token_ids == canonical
+ assert inter1.prompt_base_token_ids is not None
+
+ async def test_prepare_prompt_fallback_on_unsupported_template(
+ self, qwen3_tokenizer, monkeypatch
+ ):
+ monkeypatch.setattr(
+ IncrementalPromptRenderer, "is_supported", lambda *args, **kwargs: False
+ )
+
+ messages = [{"role": "user", "content": "Query 1"}]
+ inter1 = InteractionWithTokenLogpReward(
+ messages=list(messages), chat_template_type="hf"
+ )
+ _ = await _prepare_prompt(
+ tokenizer=qwen3_tokenizer,
+ processor=None,
+ tokenizer_messages=messages,
+ concat_messages=messages,
+ image_data=[],
+ parent=None,
+ chat_template_type="hf",
+ tools=[WEATHER_TOOL],
+ extra_body={},
+ interaction=inter1,
+ )
+
+ asst = {"role": "assistant", "content": "Response 1"}
+ user2 = {"role": "user", "content": "Query 2"}
+ messages.extend([asst, user2])
+ inter2 = InteractionWithTokenLogpReward(
+ messages=list(messages), chat_template_type="hf", parent=inter1
+ )
+
+ p2 = await _prepare_prompt(
+ tokenizer=qwen3_tokenizer,
+ processor=None,
+ tokenizer_messages=messages,
+ concat_messages=[user2],
+ image_data=[],
+ parent=inter1,
+ chat_template_type="hf",
+ tools=[WEATHER_TOOL],
+ extra_body={},
+ interaction=inter2,
+ )
+
+ canonical = apply_chat_template(
+ qwen3_tokenizer,
+ messages,
+ tools=[WEATHER_TOOL],
+ add_generation_prompt=True,
+ tokenize=True,
+ )
+ assert p2.input_ids == canonical