diff --git a/renderers/__init__.py b/renderers/__init__.py index 5baf242..f422e76 100644 --- a/renderers/__init__.py +++ b/renderers/__init__.py @@ -7,10 +7,10 @@ __version__ = "0+unknown" from renderers.base import ( + MULTIMODAL_MODELS, Content, ContentPart, ImagePart, - MULTIMODAL_MODELS, Message, MultiModalData, MultimodalRenderer, @@ -43,17 +43,16 @@ from renderers.configs import ( AutoRendererConfig, BaseRendererConfig, - config_from_name, - DefaultRendererConfig, DeepSeekR1RendererConfig, DeepSeekV3RendererConfig, + DefaultRendererConfig, + GLM5RendererConfig, GLM45RendererConfig, GLM51RendererConfig, - GLM5RendererConfig, GptOssRendererConfig, Hy3RendererConfig, - KimiK25RendererConfig, KimiK2RendererConfig, + KimiK25RendererConfig, LagunaXS2RendererConfig, LagunaXS21RendererConfig, Llama3RendererConfig, @@ -61,11 +60,13 @@ Nemotron3RendererConfig, Nemotron3UltraRendererConfig, PrimeQwen3RendererConfig, - Qwen35RendererConfig, - Qwen36RendererConfig, Qwen3RendererConfig, Qwen3VLRendererConfig, + Qwen35RendererConfig, + Qwen36RendererConfig, RendererConfig, + RlmRendererConfig, + config_from_name, ) # Concrete renderer classes are lazy-loaded so that consumers needing @@ -101,6 +102,7 @@ "Qwen36Renderer": "renderers.qwen36", "Qwen3Renderer": "renderers.qwen3", "Qwen3VLRenderer": "renderers.qwen3_vl", + "RlmRenderer": "renderers.rlm", } @@ -175,6 +177,8 @@ def __dir__() -> list[str]: "Qwen3RendererConfig", "Qwen3VLRenderer", "Qwen3VLRendererConfig", + "RlmRenderer", + "RlmRendererConfig", "RenderedConversation", "RenderedTokens", "RenderedTrainingSample", diff --git a/renderers/base.py b/renderers/base.py index 304d6cc..d866823 100644 --- a/renderers/base.py +++ b/renderers/base.py @@ -1066,6 +1066,8 @@ def bridge_to_next_turn(self, *args: Any, **kwargs: Any) -> "RenderedTokens | No "meta-llama/Llama-3.2-3B-Instruct": "llama-3", # Poolside Laguna. The two checkpoints ship different chat templates, # each mirrored by its own renderer class. + # RLM chat format (minimal role tags, rlm harness). + "PrimeIntellect/RLM-Chat-Template": "rlm", "poolside/Laguna-XS.2": "laguna-xs.2", "poolside/Laguna-XS-2.1": "laguna-xs-2.1", # GPT-OSS. @@ -1324,6 +1326,7 @@ def _populate_registry(): from renderers.qwen3_vl import Qwen3VLRenderer from renderers.qwen35 import Qwen35Renderer from renderers.qwen36 import Qwen36Renderer + from renderers.rlm import RlmRenderer RENDERER_REGISTRY.update( { @@ -1347,6 +1350,7 @@ def _populate_registry(): "llama-3": Llama3Renderer, "nemotron-3": Nemotron3Renderer, "nemotron-3-ultra": Nemotron3UltraRenderer, + "rlm": RlmRenderer, "gpt-oss": GptOssRenderer, } ) diff --git a/renderers/configs.py b/renderers/configs.py index ad32044..d199841 100644 --- a/renderers/configs.py +++ b/renderers/configs.py @@ -663,6 +663,29 @@ class DeepSeekR1RendererConfig(BaseRendererConfig): name: Literal["deepseek-r1"] = "deepseek-r1" +class RlmRendererConfig(BaseRendererConfig): + """RLM chat format (``PrimeIntellect/RLM-Chat-Template``) — minimal + role-tag-only format for the rlm harness. + + One message, one tag block: `` `` + plus the single inline tool call ``...``. Thinking is + structurally never dropped, so ``thinking_retention`` may only be left + unset or ``"all"``. Tool schemas are never rendered; ``tools=`` is + validation-only (exactly ``[ipython]`` is accepted). + """ + + name: Literal["rlm"] = "rlm" + + @model_validator(mode="after") + def _check_thinking_retention(self): + if self.thinking_retention not in (None, "all"): + raise ValueError( + "rlm never drops thinking; thinking_retention must be unset " + f"or 'all', got {self.thinking_retention!r}" + ) + return self + + RendererConfig = Annotated[ Union[ AutoRendererConfig, @@ -687,6 +710,7 @@ class DeepSeekR1RendererConfig(BaseRendererConfig): Nemotron3UltraRendererConfig, DeepSeekV3RendererConfig, DeepSeekR1RendererConfig, + RlmRendererConfig, ], Field(discriminator="name"), ] @@ -727,6 +751,7 @@ class DeepSeekR1RendererConfig(BaseRendererConfig): "nemotron-3-ultra": Nemotron3UltraRendererConfig, "deepseek-v3": DeepSeekV3RendererConfig, "deepseek-r1": DeepSeekR1RendererConfig, + "rlm": RlmRendererConfig, } diff --git a/renderers/rlm.py b/renderers/rlm.py new file mode 100644 index 0000000..829887d --- /dev/null +++ b/renderers/rlm.py @@ -0,0 +1,357 @@ +"""Renderer for the RLM chat format (``PrimeIntellect/RLM-Chat-Template``). + +The format is deliberately minimal — one message, one tag block, nothing else: + + ......{content}{code}{result} + +All ten delimiters are single added tokens (reserved ````..```` +slots of the Nemotron-3 tokenizer, renamed), so every tag is one token and +```` (id 23 on the reference tokenizer) is the eos/stop token. + +Semantics: + +- **Thinking is never dropped.** Assistant content renders verbatim on every + turn; thinking retention is structurally ``"all"`` and the bridge never + re-renders for retention. A separate ``reasoning_content`` field (as parse + APIs and some datasets carry) is glued back as ``{reasoning}`` + before the content — plain text, since the format has no native thinking + syntax. The HF chat template ignores ``reasoning_content`` (jinja cannot + guarantee the field exists); the renderer restores it instead of dropping + it, which is the only divergence from byte-parity with the template. +- **ipython is the only tool.** An assistant turn may end with exactly one + ``{code}`` call; the tool result comes back as an + ``...`` turn (OpenAI ``role: "tool"``). Tool *schemas* are + never rendered: the ``tools=`` kwarg is validation-only — the rlm harness + sends ``tools=[ipython]`` on every chat-completion call, which is accepted + and ignored; anything else raises. (The HF template raises on any + ``tools=`` because ``apply_chat_template`` is never in the serving path — + the renderer is; see the PR description for the end-to-end trace.) +""" + +from __future__ import annotations + +import json +from typing import Any + +from transformers.tokenization_utils import PreTrainedTokenizer + +from renderers.base import ( + Message, + ParsedResponse, + ParsedToolCall, + RenderedTokens, + ToolCallParseStatus, + ToolSpec, + extract_message_tool_names, + reject_assistant_in_extension, + resolve_thinking_retention, + trim_to_turn_close, +) +from renderers.configs import RlmRendererConfig + + +class RlmRenderer: + """Deterministic message ↔ token renderer for the RLM chat format.""" + + def __init__( + self, + tokenizer: PreTrainedTokenizer, + config: RlmRendererConfig | None = None, + ): + self._tokenizer = tokenizer + self.config = config or RlmRendererConfig() + # Retention is structurally "all": the format has no truncate rule and + # the config validator rejects anything else. + self.effective_thinking_retention = resolve_thinking_retention(self.config, "all") + + self._system = self._token_id("") + self._system_end = self._token_id("") + self._user = self._token_id("") + self._user_end = self._token_id("") + self._assistant = self._token_id("") + self._assistant_end = self._token_id("") + self._ipython = self._token_id("") + self._ipython_end = self._token_id("") + self._output = self._token_id("") + self._output_end = self._token_id("") + + def _token_id(self, token: str) -> int: + tid = self._tokenizer.convert_tokens_to_ids(token) + if not isinstance(tid, int) or tid == self._tokenizer.unk_token_id: + raise AssertionError( + f"Special token {token!r} not found in tokenizer vocabulary — " + "the rlm renderer requires the RLM-Chat-Template tokenizer " + "(single-token role tags)." + ) + return tid + + def _encode(self, text: str) -> list[int]: + if not text: + return [] + return self._tokenizer.encode(text, add_special_tokens=False) + + @staticmethod + def _render_content(content: Any) -> str: + if content is None: + return "" + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for item in content: + if isinstance(item, str): + parts.append(item) + elif isinstance(item, dict) and "text" in item: + parts.append(item["text"]) + else: + raise ValueError(f"Unexpected content item: {item}") + return "".join(parts) + raise TypeError(f"Unexpected content type: {type(content)}") + + @staticmethod + def _validate_tools(tools: list[ToolSpec] | None) -> None: + """tools= is validation-only: exactly [ipython] is accepted, nothing is rendered.""" + if not tools: + return + if len(tools) != 1: + raise ValueError(f"rlm accepts exactly one tool (ipython), got {len(tools)}") + fn = tools[0].get("function") or {} + if fn.get("name") != "ipython": + raise ValueError(f"rlm's only tool is ipython, got {fn.get('name')!r}") + + def _assistant_body(self, msg: Message, content: str) -> str: + """Assemble the assistant body: verbatim content (+ restored + reasoning), then the single optional ipython call.""" + reasoning = msg.get("reasoning_content") + if isinstance(reasoning, str) and reasoning: + content = "" + reasoning + "" + content + + tool_calls = msg.get("tool_calls") or [] + if not tool_calls: + return content + if len(tool_calls) != 1: + raise ValueError("rlm allows exactly one ipython call per assistant turn") + name, arguments = self._normalize_tool_call(tool_calls[0]) + if name != "ipython": + raise ValueError(f"unknown tool: {name!r} (only ipython exists)") + if not isinstance(arguments, dict) or "code" not in arguments: + raise ValueError("ipython arguments must carry a 'code' key") + code = arguments["code"] + if not isinstance(code, str): + raise ValueError("ipython 'code' must be a string") + return content + "" + code + "" + + @staticmethod + def _normalize_tool_call(tc: Any) -> tuple[Any, Any]: + """Return ``(name, arguments)`` from any of the shapes rlm data carries. + + Accepted: the OpenAI shape (``{"function": {"name", "arguments"}}``), + the verifiers trace shape (flat ``{"name", "arguments"}``), and either + as a JSON string (HF datasets store trace tool_calls as strings). + ``arguments`` may itself be a JSON string; it is parsed here. + """ + if isinstance(tc, str): + tc = json.loads(tc) + if not isinstance(tc, dict): + raise ValueError(f"unparseable tool call: {tc!r}") + fn = tc.get("function") if isinstance(tc.get("function"), dict) else tc + arguments = fn.get("arguments") + if isinstance(arguments, str): + arguments = json.loads(arguments) + return fn.get("name"), arguments + + # The /<\ipython> markers are single added tokens, so encoding the + # assembled body yields the same ids as the chat template (which encodes a + # rendered string). Mirrors Nemotron3Renderer._render_assistant. + + def render( + self, + messages: list[Message], + *, + tools: list[ToolSpec] | None = None, + add_generation_prompt: bool = False, + ) -> RenderedTokens: + if not messages: + raise ValueError("No messages provided.") + self._validate_tools(tools) + + tokens: list[int] = [] + indices: list[int] = [] + sampled: list[bool] = [] + content_mask: list[bool] = [] + + def emit(ids: list[int], msg_idx: int, *, is_sampled: bool, is_content: bool) -> None: + tokens.extend(ids) + indices.extend([msg_idx] * len(ids)) + sampled.extend([is_sampled] * len(ids)) + content_mask.extend([is_content] * len(ids)) + + for i, msg in enumerate(messages): + role = msg.get("role") + content = self._render_content(msg.get("content")) + + if role == "system": + # Any position: the chat template renders wherever it + # appears, and rlm compaction traces carry mid-list systems. + emit([self._system], i, is_sampled=False, is_content=False) + emit(self._encode(content), i, is_sampled=False, is_content=True) + emit([self._system_end], i, is_sampled=False, is_content=False) + elif role == "user": + emit([self._user], i, is_sampled=False, is_content=False) + emit(self._encode(content), i, is_sampled=False, is_content=True) + emit([self._user_end], i, is_sampled=False, is_content=False) + elif role == "assistant": + # The opening tag is the generation prompt (never sampled); + # body and closing are what the model produces — + # on assistant the invariant is_content == sampled_mask holds + # (mirrors Nemotron3Renderer._render_assistant). + emit([self._assistant], i, is_sampled=False, is_content=False) + body = self._assistant_body(msg, content) + emit(self._encode(body), i, is_sampled=True, is_content=True) + emit([self._assistant_end], i, is_sampled=True, is_content=True) + elif role == "tool": + emit([self._output], i, is_sampled=False, is_content=False) + emit(self._encode(content), i, is_sampled=False, is_content=True) + emit([self._output_end], i, is_sampled=False, is_content=False) + else: + raise ValueError(f"Unexpected message role: {role}") + + if add_generation_prompt: + emit([self._assistant], -1, is_sampled=False, is_content=False) + + return RenderedTokens( + token_ids=tokens, + message_indices=indices, + sampled_mask=sampled, + is_content=content_mask, + message_roles=[m.get("role") or "" for m in messages], + message_tool_names=extract_message_tool_names(messages), + ) + + def render_ids( + self, + messages: list[Message], + *, + tools: list[ToolSpec] | None = None, + add_generation_prompt: bool = False, + ) -> list[int]: + return self.render(messages, tools=tools, add_generation_prompt=add_generation_prompt).token_ids + + def parse_response( + self, + token_ids: list[int], + *, + tools: list[ToolSpec] | None = None, + ) -> ParsedResponse: + ids = list(token_ids) + while ids and ids[-1] == self._assistant_end: + ids = ids[:-1] + + content_segments: list[str] = [] + tool_calls: list[ParsedToolCall] = [] + pos = 0 + while True: + try: + start = ids.index(self._ipython, pos) + except ValueError: + content_segments.append(self._decode(ids[pos:])) + break + content_segments.append(self._decode(ids[pos:start])) + try: + end = ids.index(self._ipython_end, start + 1) + except ValueError: + raw = self._decode(ids[start + 1 :]) + tool_calls.append( + ParsedToolCall( + raw=raw, + name="ipython", + token_span=(start, len(ids)), + status=ToolCallParseStatus.UNCLOSED_BLOCK, + ) + ) + break + code = self._decode(ids[start + 1 : end]) + tool_calls.append( + ParsedToolCall( + raw=code, + name="ipython", + arguments={"code": code}, + token_span=(start, end + 1), + status=ToolCallParseStatus.OK, + ) + ) + pos = end + 1 + + # Thinking stays inline in content, verbatim — never split out. + return ParsedResponse( + content="".join(content_segments), + reasoning_content=None, + tool_calls=tool_calls, + ) + + def _decode(self, ids: list[int]) -> str: + if not ids: + return "" + return self._tokenizer.decode(ids) + + def get_stop_token_ids(self) -> list[int]: + return [self._assistant_end] + + def bridge_to_next_turn( + self, + previous_prompt_ids: list[int], + previous_completion_ids: list[int], + new_messages: list[Message], + *, + tools: list[ToolSpec] | None = None, + ) -> RenderedTokens | None: + if not previous_prompt_ids or not new_messages or reject_assistant_in_extension(new_messages): + return None + self._validate_tools(tools) + # Retention is "all": appending never requires re-rendering history. + + previous_ids = trim_to_turn_close( + previous_prompt_ids, + previous_completion_ids, + {self._assistant_end}, + synthesize_close=self._assistant_end, + ) + if previous_ids is None: + return None + + ext: list[int] = [] + ext_indices: list[int] = [] + ext_content: list[bool] = [] + + def emit(ids: list[int], msg_idx: int, *, is_content: bool) -> None: + ext.extend(ids) + ext_indices.extend([msg_idx] * len(ids)) + ext_content.extend([is_content] * len(ids)) + + for i, msg in enumerate(new_messages): + role = msg.get("role") + content = self._render_content(msg.get("content")) + if role == "user": + emit([self._user], i, is_content=False) + emit(self._encode(content), i, is_content=True) + emit([self._user_end], i, is_content=False) + elif role == "tool": + emit([self._output], i, is_content=False) + emit(self._encode(content), i, is_content=True) + emit([self._output_end], i, is_content=False) + else: + # System (or anything else) mid-conversation: not bridgeable. + return None + + # Generation prompt. + emit([self._assistant], -1, is_content=False) + + total_len = len(previous_ids) + len(ext) + return RenderedTokens( + token_ids=previous_ids + ext, + message_indices=[-1] * len(previous_ids) + ext_indices, + sampled_mask=[False] * total_len, + is_content=[False] * len(previous_ids) + ext_content, + message_roles=[m.get("role") or "" for m in new_messages], + message_tool_names=extract_message_tool_names(new_messages), + ) diff --git a/tests/test_rlm.py b/tests/test_rlm.py new file mode 100644 index 0000000..1bdce8d --- /dev/null +++ b/tests/test_rlm.py @@ -0,0 +1,262 @@ +"""Tests for the rlm renderer (RLM chat format). + +The reference tokenizer is ``PrimeIntellect/RLM-Chat-Template`` (private): +the Nemotron-3 Super tokenizer with reserved ````..```` +slots renamed to the ten single-token role tags, plus the minimal chat +template. The fixture below reconstructs it from the public Nemotron-3 Super +tokenizer so CI needs no private-hub access; the byte-level rename recipe is +identical to the published repo's build. +""" + +import json +from pathlib import Path + +import pytest +from renderers import create_renderer +from renderers.base import ToolCallParseStatus, load_tokenizer +from renderers.configs import RlmRendererConfig, config_from_name + +BASE_MODEL = "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16" + +TAGS = { + 18: "", + 19: "", + 20: "", + 21: "", + 22: "", + 23: "", + 24: "", + 25: "", + 26: "", + 27: "", +} + +# Mirrors PrimeIntellect/RLM-Chat-Template chat_template.jinja exactly. +CHAT_TEMPLATE = ( + "{%- if tools -%}" + '{{- raise_exception("This template does not support tools=. The single built-in tool is ipython, invoked inline as ... in assistant turns.") -}}' + "{%- endif -%}" + "{%- for message in messages -%}" + "{%- if message['role'] == 'system' -%}" + "{{- '' + message['content'] + '' -}}" + "{%- elif message['role'] == 'user' -%}" + "{{- '' + message['content'] + '' -}}" + "{%- elif message['role'] == 'assistant' -%}" + "{{- '' -}}" + "{%- if message['content'] -%}{{- message['content'] -}}{%- endif -%}" + "{%- if message['tool_calls'] is defined and message['tool_calls'] -%}" + "{%- set call = message['tool_calls'][0]['function'] -%}" + "{{- '' + call['arguments']['code'] + '' -}}" + "{%- endif -%}" + "{{- '' -}}" + "{%- elif message['role'] == 'tool' -%}" + "{{- '' + message['content'] + '' -}}" + "{%- endif -%}" + "{%- endfor -%}" + "{%- if add_generation_prompt -%}{{- '' -}}{%- endif -%}" +) + + +@pytest.fixture(scope="module") +def rlm_tokenizer(tmp_path_factory): + """Rebuild the RLM-Chat-Template tokenizer from the public base model.""" + from huggingface_hub import hf_hub_download + + out = tmp_path_factory.mktemp("rlm-tokenizer") + for fname in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json"): + src = Path(hf_hub_download(BASE_MODEL, fname)) + data = json.loads(src.read_text()) + if fname == "tokenizer.json": + for tok in data["added_tokens"]: + if tok["id"] in TAGS: + tok["content"] = TAGS[tok["id"]] + vocab = data["model"]["vocab"] + for i, tag in TAGS.items(): + del vocab[f""] + vocab[tag] = i + elif fname == "tokenizer_config.json": + for i, tag in TAGS.items(): + data["added_tokens_decoder"][str(i)]["content"] = tag + data["eos_token"] = "" + data.pop("chat_template", None) + else: + data["eos_token"]["content"] = "" + (out / fname).write_text(json.dumps(data, ensure_ascii=False)) + (out / "chat_template.jinja").write_text(CHAT_TEMPLATE) + return load_tokenizer(str(out)) + + +@pytest.fixture(scope="module") +def renderer(rlm_tokenizer): + return create_renderer(rlm_tokenizer, RlmRendererConfig()) + + +MESSAGES = [ + {"role": "system", "content": "You are a coding agent."}, + {"role": "user", "content": "Fix the bug in foo.py"}, + { + "role": "assistant", + "content": "Look at the file first.I'll inspect foo.py.", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": { + "name": "ipython", + "arguments": {"code": "print(open('foo.py').read())"}, + }, + } + ], + }, + {"role": "tool", "tool_call_id": "c1", "content": "def foo():\n return 1/0\n"}, + {"role": "assistant", "content": "Division by zero.Fixed."}, +] + + +def test_tags_are_single_tokens(rlm_tokenizer): + for i, tag in TAGS.items(): + assert rlm_tokenizer.encode(tag, add_special_tokens=False) == [i] + assert rlm_tokenizer.eos_token_id == 23 + + +@pytest.mark.parametrize("add_generation_prompt", [False, True]) +def test_render_parity_with_chat_template(renderer, rlm_tokenizer, add_generation_prompt): + for upto in (2, 3, 4, 5): + msgs = MESSAGES[:upto] + if add_generation_prompt and msgs[-1]["role"] == "assistant": + continue + expected = rlm_tokenizer.apply_chat_template(msgs, tokenize=True, add_generation_prompt=add_generation_prompt) + if not isinstance(expected, list): + expected = expected["input_ids"] + got = renderer.render_ids(msgs, add_generation_prompt=add_generation_prompt) + assert got == expected, f"parity mismatch at upto={upto}" + + +def test_thinking_never_dropped(renderer, rlm_tokenizer): + ids = renderer.render_ids(MESSAGES) + text = rlm_tokenizer.decode(ids) + assert "Look at the file first." in text + assert "Division by zero." in text + + +def test_reasoning_content_is_restored_not_dropped(renderer, rlm_tokenizer): + msgs = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "answer", "reasoning_content": "hidden plan"}, + ] + text = rlm_tokenizer.decode(renderer.render_ids(msgs)) + assert "hidden plananswer" in text + + +def test_tools_validation(renderer): + ipython_tool = { + "type": "function", + "function": {"name": "ipython", "parameters": {"type": "object"}}, + } + msgs = MESSAGES[:2] + # exactly [ipython]: accepted, renders identically to no tools + assert renderer.render_ids(msgs, tools=[ipython_tool]) == renderer.render_ids(msgs) + with pytest.raises(ValueError, match="only tool is ipython"): + renderer.render_ids(msgs, tools=[{"type": "function", "function": {"name": "bash"}}]) + with pytest.raises(ValueError, match="exactly one tool"): + renderer.render_ids(msgs, tools=[ipython_tool, ipython_tool]) + + +def test_render_rejects_bad_tool_calls(renderer): + base = {"role": "assistant", "content": ""} + with pytest.raises(ValueError, match="unknown tool"): + renderer.render_ids( + [ + {"role": "user", "content": "u"}, + { + **base, + "tool_calls": [{"type": "function", "function": {"name": "bash", "arguments": {"code": "x"}}}], + }, + ] + ) + with pytest.raises(ValueError, match="exactly one ipython call"): + tc = {"type": "function", "function": {"name": "ipython", "arguments": {"code": "x"}}} + renderer.render_ids([{"role": "user", "content": "u"}, {**base, "tool_calls": [tc, tc]}]) + + +def test_render_accepts_trace_shape_tool_calls(renderer, rlm_tokenizer): + # HF datasets store verifiers-trace tool calls as flat JSON strings. + msgs = [ + {"role": "user", "content": "u"}, + { + "role": "assistant", + "content": "c", + "tool_calls": ['{"id": "t1", "name": "ipython", "arguments": "{\\"code\\": \\"print(1)\\"}"}'], + }, + ] + text = rlm_tokenizer.decode(renderer.render_ids(msgs)) + assert "print(1)" in text + + +def test_parse_response_roundtrip(renderer, rlm_tokenizer): + body = "plantextprint(1)" + ids = rlm_tokenizer.encode(body, add_special_tokens=False) + [23] + parsed = renderer.parse_response(ids) + assert parsed.content == "plantext" + assert parsed.reasoning_content is None + assert len(parsed.tool_calls) == 1 + call = parsed.tool_calls[0] + assert call.status == ToolCallParseStatus.OK + assert call.name == "ipython" + assert call.arguments == {"code": "print(1)"} + start, end = call.token_span + stripped = ids[:-1] + assert stripped[start] == 24 and stripped[end - 1] == 25 + + +def test_parse_response_content_only(renderer, rlm_tokenizer): + ids = rlm_tokenizer.encode("just an answer", add_special_tokens=False) + [23] + parsed = renderer.parse_response(ids) + assert parsed.content == "just an answer" + assert parsed.tool_calls == [] + + +def test_parse_response_unclosed_ipython(renderer, rlm_tokenizer): + ids = rlm_tokenizer.encode("xprint(", add_special_tokens=False) + parsed = renderer.parse_response(ids) + assert parsed.content == "x" + assert len(parsed.tool_calls) == 1 + assert parsed.tool_calls[0].status == ToolCallParseStatus.UNCLOSED_BLOCK + + +def test_stop_tokens(renderer): + assert renderer.get_stop_token_ids() == [23] + + +def test_bridge_extends_exactly(renderer): + prompt = renderer.render_ids(MESSAGES[:2], add_generation_prompt=True) + # completion the model would sample: body + + full_turn = renderer.render(MESSAGES[:3]) + completion = full_turn.token_ids[len(prompt) :] + bridged = renderer.bridge_to_next_turn(prompt, completion, [MESSAGES[3]]) + assert bridged is not None + expected = renderer.render_ids(MESSAGES[:4], add_generation_prompt=True) + assert bridged.token_ids == expected + assert bridged.token_ids[: len(prompt) + len(completion)] == prompt + completion + + +def test_bridge_rejects_assistant_messages(renderer): + prompt = renderer.render_ids(MESSAGES[:2], add_generation_prompt=True) + out = renderer.bridge_to_next_turn(prompt, [23], [{"role": "assistant", "content": "no"}]) + assert out is None + + +def test_config_rejects_non_all_retention(): + with pytest.raises(Exception, match="never drops thinking"): + RlmRendererConfig(thinking_retention="tool_cycle") + assert config_from_name("rlm").thinking_retention is None + + +def test_content_mask_roles(renderer): + rendered = renderer.render(MESSAGES) + spans = rendered.content_token_spans_by_role() + assert set(spans) >= {"assistant", "tool", "user", "system"} + # assistant invariant: is_content == sampled_mask + for is_c, is_s, idx in zip(rendered.is_content, rendered.sampled_mask, rendered.message_indices): + if idx >= 0 and MESSAGES[idx]["role"] == "assistant": + assert is_c == is_s