diff --git a/areal/experimental/openai/client.py b/areal/experimental/openai/client.py
index 21aeaf77f0..221621232e 100644
--- a/areal/experimental/openai/client.py
+++ b/areal/experimental/openai/client.py
@@ -57,6 +57,7 @@
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
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 +673,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 +734,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 +756,77 @@ 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) :]
+ parent_base = parent.prompt_base_token_ids
+ 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
+ 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,
)
- )
+ if rendered is not None:
+ input_ids, new_base = rendered
+ if interaction is not None:
+ interaction.prompt_base_token_ids = new_base
+ 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:
+ if processor is None and IncrementalPromptRenderer.is_supported(
+ tokenizer, 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_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 +1132,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 +1590,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..97882a0750
--- /dev/null
+++ b/areal/experimental/openai/prompt_renderer.py
@@ -0,0 +1,304 @@
+# 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 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
+
+
+class IncrementalPromptRenderer:
+ """Renders multi-turn agent prompts incrementally with token parity guarantees."""
+
+ _capability_cache: dict[tuple[Any, ...], 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."""
+ if not hasattr(tokenizer, "chat_template") or not tokenizer.chat_template:
+ return 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 a synthetic 2-turn sequence
+ supported = cls._probe_capability(tokenizer, tools, chat_template_kwargs)
+ with cls._lock:
+ cls._capability_cache[key] = supported
+
+ if supported:
+ logger.debug(
+ "Incremental prompt rendering verified and enabled for tokenizer: %s",
+ getattr(tokenizer, "name_or_path", type(tokenizer).__name__),
+ )
+ else:
+ logger.debug(
+ "Incremental prompt rendering not supported for tokenizer: %s; using full fallback.",
+ getattr(tokenizer, "name_or_path", type(tokenizer).__name__),
+ )
+ return supported
+
+ @classmethod
+ def _probe_capability(
+ cls,
+ tokenizer: PreTrainedTokenizerFast,
+ tools: Iterable[ChatCompletionToolParam] | None,
+ chat_template_kwargs: dict[str, Any] | None,
+ ) -> bool:
+ """Run a probe to verify token-for-token equality between incremental and full rendering."""
+ 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
+ incr_2 = base_1 + d_gen[len(d0) :]
+ return full_2 == incr_2
+ except Exception as e:
+ logger.debug("PromptRenderer 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,
+ ) -> tuple[list[int], list[int]] | None:
+ """Render prompt tokens for delta messages appended to parent base tokens.
+
+ Returns (prompt_token_ids, new_base_token_ids) or None on failure.
+ """
+ if not delta_messages:
+ 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..ed01e1dbee 100644
--- a/areal/experimental/openai/types.py
+++ b/areal/experimental/openai/types.py
@@ -45,8 +45,9 @@ 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
mm_token_type_ids: list[int] | None = None
multi_modal_input: dict[str, torch.Tensor] | None = None
diff --git a/areal/utils/pkg_version.py b/areal/utils/pkg_version.py
index 75e18f4050..54c155947d 100644
--- a/areal/utils/pkg_version.py
+++ b/areal/utils/pkg_version.py
@@ -39,7 +39,10 @@ def is_version_greater_or_equal(package_name: str, target_version: str) -> bool:
:param target_version: Target version to compare against.
:return: True if the installed version is greater than or equal to the target version, False otherwise.
"""
- installed_version = get_version(package_name)
+ try:
+ installed_version = get_version(package_name)
+ except PackageNotFoundError:
+ return False
return compare_versions(installed_version, target_version) >= 0
@@ -51,7 +54,10 @@ def is_version_less(package_name: str, target_version: str) -> bool:
:param target_version: Target version to compare against.
:return: True if the installed version is less than the target version, False otherwise.
"""
- installed_version = get_version(package_name)
+ try:
+ installed_version = get_version(package_name)
+ except PackageNotFoundError:
+ return False
return compare_versions(installed_version, target_version) < 0
@@ -63,5 +69,8 @@ def is_version_equal(package_name: str, target_version: str) -> bool:
:param target_version: Target version to compare against.
:return: True if the installed version is equal to the target version, False otherwise.
"""
- installed_version = get_version(package_name)
+ try:
+ installed_version = get_version(package_name)
+ except PackageNotFoundError:
+ return False
return compare_versions(installed_version, target_version) == 0
diff --git a/areal/utils/testing_utils.py b/areal/utils/testing_utils.py
index 6c03b472e5..df155e24d3 100644
--- a/areal/utils/testing_utils.py
+++ b/areal/utils/testing_utils.py
@@ -15,7 +15,6 @@
from transformers import AutoConfig
from areal.api import InferenceEngine, RolloutWorkflow
-from areal.experimental.models.archon import get_model_spec, is_supported_model
from areal.experimental.openai.types import InteractionWithTokenLogpReward
from areal.utils import logging
from areal.utils.save_load import get_state_dict_from_repo_id_or_path
@@ -84,49 +83,85 @@ def get_dataset_path(local_path: str, hf_id: str) -> str:
raise
-# Model paths for testing (keyed by HF model_type)
-# Dense models (fast to instantiate even on meta device)
-DENSE_MODEL_PATHS = {
- "qwen2": get_model_path(
+class _LazyModelDict(dict):
+ """Dictionary that lazily resolves model paths upon access."""
+
+ def __init__(self, specs: dict[str, tuple[str, str]]):
+ super().__init__()
+ self._specs = specs
+
+ def __getitem__(self, key: str) -> str:
+ if key not in self:
+ if key in self._specs:
+ local_path, hf_id = self._specs[key]
+ self[key] = get_model_path(local_path, hf_id)
+ else:
+ raise KeyError(key)
+ return super().__getitem__(key)
+
+ def get(self, key: str, default: Any = None) -> Any:
+ try:
+ return self[key]
+ except KeyError:
+ return default
+
+ def __contains__(self, key: object) -> bool:
+ return key in self._specs or super().__contains__(key)
+
+ def items(self):
+ for k in self._specs:
+ yield k, self[k]
+
+ def values(self):
+ for k in self._specs:
+ yield self[k]
+
+ def keys(self):
+ return self._specs.keys()
+
+
+_DENSE_SPECS = {
+ "qwen2": (
"/storage/openpsi/models/Qwen__Qwen2.5-0.5B-Instruct/",
"Qwen/Qwen2.5-0.5B-Instruct",
),
- "qwen3": get_model_path(
+ "qwen3": (
"/storage/openpsi/models/Qwen__Qwen3-0.6B/",
"Qwen/Qwen3-0.6B",
),
- "qwen3_5": get_model_path(
+ "qwen3_5": (
"/storage/openpsi/models/Qwen__Qwen3.5-0.8B/",
"Qwen/Qwen3.5-0.8B",
),
- "qwen2_5_vl": get_model_path(
+ "qwen2_5_vl": (
"/storage/openpsi/models/Qwen__Qwen2.5-VL-3B-Instruct/",
"Qwen/Qwen2.5-VL-3B-Instruct",
),
- "qwen3_vl": get_model_path(
+ "qwen3_vl": (
"/storage/openpsi/models/Qwen__Qwen3-VL-2B-Instruct/",
"Qwen/Qwen3-VL-2B-Instruct",
),
}
-# MoE models (slow to instantiate due to large number of experts)
-MOE_MODEL_PATHS = {
- "qwen3_moe": get_model_path(
+_MOE_SPECS = {
+ "qwen3_moe": (
"/storage/openpsi/models/Qwen__Qwen3-30B-A3B/",
"Qwen/Qwen3-30B-A3B",
),
- "qwen3_5_moe": get_model_path(
+ "qwen3_5_moe": (
"/storage/openpsi/models/Qwen__Qwen3.5-35B-A3B",
"Qwen/Qwen3.5-35B-A3B",
),
- "qwen3_vl_moe": get_model_path(
+ "qwen3_vl_moe": (
"/storage/openpsi/models/Qwen__Qwen3-VL-30B-A3B-Instruct/",
"Qwen/Qwen3-VL-30B-A3B-Instruct",
),
}
-# Combined for backward compatibility
-MODEL_PATHS = {**DENSE_MODEL_PATHS, **MOE_MODEL_PATHS}
+# Model paths for testing (keyed by HF model_type)
+DENSE_MODEL_PATHS = _LazyModelDict(_DENSE_SPECS)
+MOE_MODEL_PATHS = _LazyModelDict(_MOE_SPECS)
+MODEL_PATHS = _LazyModelDict({**_DENSE_SPECS, **_MOE_SPECS})
def load_archon_model(
@@ -142,6 +177,7 @@ def load_archon_model(
Returns:
Tuple of (model, adapter) or (None, None) if skip_unsupported and model not supported.
"""
+ from areal.experimental.models.archon import get_model_spec, is_supported_model
from areal.infra.platforms import current_platform
config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
diff --git a/tests/experimental/openai/test_prompt_renderer.py b/tests/experimental/openai/test_prompt_renderer.py
new file mode 100644
index 0000000000..ed6fe2f42d
--- /dev/null
+++ b/tests/experimental/openai/test_prompt_renderer.py
@@ -0,0 +1,571 @@
+# 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,
+)
+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],
+ )
+ 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],
+ )
+ 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],
+ )
+ 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],
+ )
+ assert rendered is not None
+ incr_prompt, _ = rendered
+
+ assert incr_prompt == canonical_prompt
+
+
+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