diff --git a/renderers/__init__.py b/renderers/__init__.py index fa6dd75e..c37d08a1 100644 --- a/renderers/__init__.py +++ b/renderers/__init__.py @@ -51,6 +51,7 @@ DeepSeekR1RendererConfig, DeepSeekV3RendererConfig, DeepSeekV4RendererConfig, + DeepSeekV41RendererConfig, DefaultRendererConfig, Gemma4RendererConfig, GLM5RendererConfig, @@ -99,6 +100,7 @@ "DeepSeekR1Renderer": "renderers.deepseek_r1", "DeepSeekV3Renderer": "renderers.deepseek_v3", "DeepSeekV4Renderer": "renderers.deepseek_v4", + "DeepSeekV41Renderer": "renderers.deepseek_v41", "DefaultRenderer": "renderers.default", "GLM45Renderer": "renderers.glm45", "GLM51Renderer": "renderers.glm5", @@ -160,6 +162,8 @@ def __dir__() -> list[str]: "DeepSeekV3RendererConfig", "DeepSeekV4Renderer", "DeepSeekV4RendererConfig", + "DeepSeekV41Renderer", + "DeepSeekV41RendererConfig", "DefaultRenderer", "DefaultRendererConfig", "GLM45Renderer", diff --git a/renderers/base.py b/renderers/base.py index f3418841..f86ea1d2 100644 --- a/renderers/base.py +++ b/renderers/base.py @@ -1355,6 +1355,7 @@ def _populate_registry(): from renderers.deepseek_r1 import DeepSeekR1Renderer from renderers.deepseek_v3 import DeepSeekV3Renderer from renderers.deepseek_v4 import DeepSeekV4Renderer + from renderers.deepseek_v41 import DeepSeekV41Renderer from renderers.default import DefaultRenderer from renderers.gemma4 import Gemma4Renderer from renderers.glm5 import ( @@ -1409,6 +1410,7 @@ def _populate_registry(): "deepseek-v3": DeepSeekV3Renderer, "deepseek-r1": DeepSeekR1Renderer, "deepseek-v4": DeepSeekV4Renderer, + "deepseek-v4.1": DeepSeekV41Renderer, "hy3": Hy3Renderer, "inkling": InklingRenderer, "kimi-k2": KimiK2Renderer, diff --git a/renderers/configs.py b/renderers/configs.py index d6013132..9a444d15 100644 --- a/renderers/configs.py +++ b/renderers/configs.py @@ -1028,6 +1028,34 @@ def _check_thinking_retention(self) -> Self: return self +class DeepSeekV41RendererConfig(DeepSeekV4RendererConfig): + """DeepSeek-V4.1-Flash reference-encoder configuration. + + V4.1's encoder is V4's with a numeric reasoning-effort budget, a + ``<｜System｜>`` token in front of the conversation's opening effort prefix + or system message, and thinking mode on by default. + """ + + name: Literal["deepseek-v4.1"] = "deepseek-v4.1" + + enable_thinking: bool = True + """Select thinking mode, the V4.1 encoder's default.""" + + reasoning_effort: Literal["low", "high", "max"] | int = "high" + """Thinking-only budget: an int in [1, 100], or ``low`` / ``high`` / ``max`` (50 / 75 / 100).""" + + @model_validator(mode="after") + def _check_reasoning_effort(self) -> Self: + if ( + isinstance(self.reasoning_effort, int) + and not 1 <= self.reasoning_effort <= 100 + ): + raise ValueError( + f"reasoning_effort must be in [1, 100], got {self.reasoning_effort}" + ) + return self + + class CustomRendererConfig(BaseRendererConfig): """Config for a renderer that lives outside this package. @@ -1096,6 +1124,7 @@ def _validate_renderer_fields(self) -> Self: DeepSeekV3RendererConfig, DeepSeekR1RendererConfig, DeepSeekV4RendererConfig, + DeepSeekV41RendererConfig, CustomRendererConfig, ], Field(discriminator="name"), @@ -1147,6 +1176,7 @@ def _validate_renderer_fields(self) -> Self: "deepseek-v3": DeepSeekV3RendererConfig, "deepseek-r1": DeepSeekR1RendererConfig, "deepseek-v4": DeepSeekV4RendererConfig, + "deepseek-v4.1": DeepSeekV41RendererConfig, } @@ -1180,6 +1210,7 @@ def config_from_name(name: str) -> BaseRendererConfig | None: "DeepSeekR1RendererConfig", "DeepSeekV3RendererConfig", "DeepSeekV4RendererConfig", + "DeepSeekV41RendererConfig", "GLM45RendererConfig", "GLM51RendererConfig", "GLM53FlashRendererConfig", diff --git a/renderers/deepseek_v4.py b/renderers/deepseek_v4.py index e3ba16ca..749ea84d 100644 --- a/renderers/deepseek_v4.py +++ b/renderers/deepseek_v4.py @@ -40,6 +40,7 @@ from renderers.configs import DeepSeekV4RendererConfig from renderers.parsing import parse_deepseek_v4 from renderers.reasoning import prompt_ends_in_reasoning +from renderers.tagged import DEEPSEEK_V4_GRAMMAR _BOS = "<｜begin▁of▁sentence｜>" _EOS = "<｜end▁of▁sentence｜>" @@ -87,17 +88,17 @@ _TOOLS_TEMPLATE = """## Tools -You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<{dsml}tool_calls>" block like the following: +You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<{calls}>" block like the following: -<{dsml}tool_calls> -<{dsml}invoke name="$TOOL_NAME"> -<{dsml}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE +<{calls}> +<{invoke} name="$TOOL_NAME"> +<{parameter} name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE ... - -<{dsml}invoke name="$TOOL_NAME2"> + +<{invoke} name="$TOOL_NAME2"> ... - - + + String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`. @@ -300,6 +301,12 @@ def is_prefix_stable(self) -> bool: _implied_thinking_retention = "tool_cycle" + # Full DSML tag names, as in ``<{_invoke_tag} name="...">``. + _calls_tag = f"{_DSML}tool_calls" + _invoke_tag = f"{_DSML}invoke" + _parameter_tag = f"{_DSML}parameter" + _grammar = DEEPSEEK_V4_GRAMMAR + def __init__( self, tokenizer: Tokenizer, @@ -335,46 +342,69 @@ def _special_id(self, token: str) -> int: raise ValueError(f"Expected one token for {token!r}, got {ids}") return ids[0] - def _reasoning_effort_prompt(self) -> str: - """Text emitted after BOS in thinking mode; subclasses can replace it.""" + def _preamble(self, *, opens_with_system: bool) -> str: # noqa: ARG002 + """Text emitted after BOS, before the first message. + + ``opens_with_system`` tells whether the rendered conversation starts + with a system message (including the synthetic one carrying tools). + """ + if not self.config.enable_thinking: + return "" return _REASONING_EFFORT_PROMPTS[self.config.reasoning_effort] + def _is_turn_start(self, message: _LogicalMessage) -> bool: + """Whether ``message`` opens a turn, bounding dropped historical reasoning.""" + return message.role in _QUERY_ROLES + + def _tool_target(self, messages: list[_LogicalMessage]) -> int | None: + """Index of the message carrying the tools block; ``None`` adds a leading one.""" + for role in ("developer", "system"): + for index, message in enumerate(messages): + if message.role == role: + return index + return None + @staticmethod - def _render_tools(tools: list[ToolSpec]) -> str: + def _tool_call_arguments(raw_arguments: Any) -> Any: + if isinstance(raw_arguments, str): + try: + return json.loads(raw_arguments) + except (json.JSONDecodeError, TypeError): + return {"arguments": raw_arguments} + if isinstance(raw_arguments, Mapping): + return dict(raw_arguments) + return {"arguments": raw_arguments} + + def _render_tools(self, tools: list[ToolSpec]) -> str: schemas = [_json(dict(_tool_function(tool))) for tool in tools] return _TOOLS_TEMPLATE.format( - dsml=_DSML, + calls=self._calls_tag, + invoke=self._invoke_tag, + parameter=self._parameter_tag, think_start=_THINK_START, think_end=_THINK_END, tool_schemas="\n".join(schemas), ) - @staticmethod - def _render_tool_call(tool_call: Mapping[str, Any]) -> str: + def _render_tool_call(self, tool_call: Mapping[str, Any]) -> str: function = _tool_call_function(tool_call) name = function.get("name") - raw_arguments = function.get("arguments", {}) - if isinstance(raw_arguments, str): - try: - arguments = json.loads(raw_arguments) - except (json.JSONDecodeError, TypeError): - arguments = {"arguments": raw_arguments} - elif isinstance(raw_arguments, Mapping): - arguments = dict(raw_arguments) - else: - arguments = {"arguments": raw_arguments} + arguments = self._tool_call_arguments(function.get("arguments", {})) params: list[str] = [] for key, value in arguments.items(): is_string = isinstance(value, str) rendered_value = value if is_string else _json(value) params.append( - f'<{_DSML}parameter name="{key}" ' + f'<{self._parameter_tag} name="{key}" ' f'string="{str(is_string).lower()}">{rendered_value}' - f"" + f"" ) arguments_text = "\n".join(params) - return f'<{_DSML}invoke name="{name}">\n{arguments_text}\n' + return ( + f'<{self._invoke_tag} name="{name}">\n{arguments_text}\n' + f"" + ) def render( self, @@ -388,7 +418,7 @@ def render( tools=tools, add_generation_prompt=add_generation_prompt, add_bos=True, - add_effort_prompt=True, + add_preamble=True, ) def _render( @@ -398,7 +428,7 @@ def _render( tools: list[ToolSpec] | None, add_generation_prompt: bool, add_bos: bool, - add_effort_prompt: bool, + add_preamble: bool, ) -> RenderedTokens: if not messages: raise ValueError("No messages provided.") @@ -413,7 +443,7 @@ def _render( ( index for index, message in enumerate(logical_messages) - if message.role in _QUERY_ROLES + if self._is_turn_start(message) ), default=-1, ) @@ -578,39 +608,27 @@ def emit_text( if text: pending_text.append((text, message_index, sampled, content)) - if add_bos: - emit_special(self._bos, -1) - if add_effort_prompt and self.config.enable_thinking: - emit_text(self._reasoning_effort_prompt(), -1) - last_query_index = -1 for index, message in enumerate(logical_messages): - if message.role in _QUERY_ROLES: + if self._is_turn_start(message): last_query_index = index - tool_target: int | None = None - if tools: - tool_target = next( - ( - index - for index, message in enumerate(logical_messages) - if message.role == "developer" + tool_target = self._tool_target(logical_messages) if tools else None + + if add_bos: + emit_special(self._bos, -1) + if add_preamble: + emit_text( + self._preamble( + opens_with_system=logical_messages[0].role == "system" + or bool(tools and tool_target is None) ), - None, + -1, ) - if tool_target is None: - tool_target = next( - ( - index - for index, message in enumerate(logical_messages) - if message.role == "system" - ), - None, - ) - if tool_target is None: - # Equivalent to an empty synthetic system message carrying - # tools in the reference encoder. - emit_text("\n\n" + self._render_tools(tools), -1) + if tools and tool_target is None: + # Equivalent to an empty synthetic system message carrying + # tools in the reference encoder. + emit_text("\n\n" + self._render_tools(tools), -1) for index, message in enumerate(logical_messages): role = message.role @@ -683,8 +701,8 @@ def emit_text( self._render_tool_call(call) for call in message.tool_calls ) emit_text( - f"\n\n<{_DSML}tool_calls>\n{rendered_calls}\n" - f"", + f"\n\n<{self._calls_tag}>\n{rendered_calls}\n" + f"", msg_idx, sampled=True, content=True, @@ -808,6 +826,9 @@ def parse_response( ), think_end_id=self._think_end, think_start_id=self._think_start, + grammar=self._grammar, + invoke_tag=self._invoke_tag, + parameter_tag=self._parameter_tag, ) def get_stop_token_ids(self) -> list[int]: @@ -863,7 +884,7 @@ def bridge_to_next_turn( tools=None, add_generation_prompt=True, add_bos=False, - add_effort_prompt=False, + add_preamble=False, ) except ValueError: return None diff --git a/renderers/deepseek_v41.py b/renderers/deepseek_v41.py new file mode 100644 index 00000000..c758fc47 --- /dev/null +++ b/renderers/deepseek_v41.py @@ -0,0 +1,141 @@ +"""DeepSeek-V4.1-Flash renderer. + +The checkpoint's source of truth is ``encoding/encoding.py`` in the model +repository. It is V4's encoder with these changes, applied here on top of +:class:`DeepSeekV4Renderer`: + +* a ``<｜System｜>`` token leads the conversation whenever it opens with a + reasoning-effort prefix or a system message (including the empty one that + carries tools when the conversation has no leading system message); +* the reasoning effort is a numeric budget, + ``Reasoning Effort: {budget} (range 1-100, ...)``, rendered in thinking mode + only and defaulting to ``high`` (75); +* thinking mode is the default; +* DSML tags have space-prefixed names: ``<｜DSML｜ calls>``, + ``<｜DSML｜ invoke>`` and ``<｜DSML｜ parameter>``; +* tools attach only to a leading system message; +* a user turn holding only tool results does not start a new turn, so + dropped historical reasoning stops at the last real user message; +* JSON-string tool arguments may be double-encoded, and arguments that do not + decode to an object render as a single ``arguments`` parameter. + +Mid-conversation system messages are not supported yet. +""" + +from __future__ import annotations + +import json +from typing import Any, Mapping + +from renderers.base import Message, RenderedTokens, ToolSpec +from renderers.configs import DeepSeekV41RendererConfig +from renderers.deepseek_v4 import _DSML, DeepSeekV4Renderer, _LogicalMessage +from renderers.tagged import DEEPSEEK_V41_GRAMMAR + +_SYSTEM = "<｜System｜>" +_THINK_START = "" +_THINK_END = "" +_EFFORT_BUDGETS = {"low": 50, "high": 75, "max": 100} + + +def _split_inline_reasoning(message: Message) -> Message: + """Move a leading ``...`` out of an assistant's content into ``reasoning_content``. + + Many SFT datasets keep the reasoning inline in the content; rendered as is, + it would follow an empty think block instead of filling it. + """ + content = message.get("content") + if ( + message["role"] != "assistant" + or message.get("reasoning_content") is not None + or not isinstance(content, str) + ): + return message + if not content.startswith(_THINK_START) or _THINK_END not in content: + return message + reasoning, content = content.removeprefix(_THINK_START).split(_THINK_END, 1) + split = message.copy() + split["reasoning_content"] = reasoning + split["content"] = content + return split + + +class DeepSeekV41Renderer(DeepSeekV4Renderer): + """Renderer for ``deepseek-ai/DeepSeek-V4.1-Flash``.""" + + _calls_tag = f"{_DSML} calls" + _invoke_tag = f"{_DSML} invoke" + _parameter_tag = f"{_DSML} parameter" + _grammar = DEEPSEEK_V41_GRAMMAR + + def __init__(self, tokenizer, config: DeepSeekV41RendererConfig | None = None): + super().__init__(tokenizer, config or DeepSeekV41RendererConfig()) + + def _preamble(self, *, opens_with_system: bool) -> str: + if not self.config.enable_thinking: + return _SYSTEM if opens_with_system else "" + effort = self.config.reasoning_effort + budget = _EFFORT_BUDGETS[effort] if isinstance(effort, str) else effort + return ( + f"{_SYSTEM}Reasoning Effort: {budget} (range 1-100, " + "the higher the value, the more thorough the reasoning)\n\n" + ) + + def _is_turn_start(self, message: _LogicalMessage) -> bool: + if message.role == "user": + return any(block.kind != "tool_result" for block in message.blocks) + return super()._is_turn_start(message) + + def _tool_target(self, messages: list[_LogicalMessage]) -> int | None: + return 0 if messages[0].role == "system" else None + + @staticmethod + def _tool_call_arguments(raw_arguments: Any) -> Any: + arguments = raw_arguments + # A JSON string, possibly double-encoded. + for _ in range(2): + if not isinstance(arguments, str): + break + try: + arguments = json.loads(arguments) + except json.JSONDecodeError: + break + if isinstance(arguments, Mapping): + return dict(arguments) + return {"arguments": raw_arguments} + + def render( + self, + messages: list[Message], + *, + tools: list[ToolSpec] | None = None, + add_generation_prompt: bool = False, + ) -> RenderedTokens: + if any(message["role"] == "system" for message in messages[1:]): + raise NotImplementedError( + "the DeepSeek-V4.1 renderer does not render mid-conversation system messages" + ) + messages = [_split_inline_reasoning(message) for message in messages] + return super().render( + messages, tools=tools, add_generation_prompt=add_generation_prompt + ) + + 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 any(message["role"] == "system" for message in new_messages): + return None + return super().bridge_to_next_turn( + previous_prompt_ids, + previous_completion_ids, + new_messages, + tools=tools, + ) + + +__all__ = ["DeepSeekV41Renderer"] diff --git a/renderers/parsing.py b/renderers/parsing.py index 352ae817..86feefca 100644 --- a/renderers/parsing.py +++ b/renderers/parsing.py @@ -29,6 +29,7 @@ from renderers.tagged import ( DEEPSEEK_V4_GRAMMAR, GLM_GRAMMAR, + TaggedGrammar, TaggedToolBlock, scan_tagged_response, ) @@ -1204,6 +1205,9 @@ def parse_deepseek_v4( thinking_enabled: bool, think_end_id: int, think_start_id: int, + grammar: TaggedGrammar = DEEPSEEK_V4_GRAMMAR, + invoke_tag: str = "｜DSML｜invoke", + parameter_tag: str = "｜DSML｜parameter", ) -> ParsedResponse: """Parse a DeepSeek V4 completion. @@ -1213,6 +1217,8 @@ def parse_deepseek_v4( decoded grammar carries an explicit ``string=`` flag for lossless argument type recovery. ``thinking_enabled`` describes the prompt's initial channel; reasoning opened in the completion is recognized independently. + DeepSeek-V4.1 spells its DSML tags differently and passes its own + ``grammar`` and tag names. """ ids = _strip_stop_tokens(token_ids, stop_ids) parsed = scan_tagged_response( @@ -1220,7 +1226,7 @@ def parse_deepseek_v4( ids, thinking=thinking_enabled, atomic_markers={think_start_id: "think_start", think_end_id: "think_end"}, - grammar=DEEPSEEK_V4_GRAMMAR, + grammar=grammar, ) reasoning = parsed.boundary.text return ParsedResponse( @@ -1230,19 +1236,24 @@ def parse_deepseek_v4( if parsed.boundary.is_open else (reasoning or "").rstrip() or None ), - tool_calls=[_parse_deepseek_v4_tool_call(block) for block in parsed.tools], + tool_calls=[ + _parse_deepseek_v4_tool_call(block, invoke_tag, parameter_tag) + for block in parsed.tools + ], reasoning_complete=not parsed.boundary.is_open, ) -def _parse_deepseek_v4_tool_call(region: TaggedToolBlock) -> ParsedToolCall: +def _parse_deepseek_v4_tool_call( + region: TaggedToolBlock, invoke_tag: str, parameter_tag: str +) -> ParsedToolCall: """Convert one DSML attempt; the shared grammar owns its boundaries.""" import re - invoke_end = "" + invoke_end = f"" + parameter = re.escape(parameter_tag) parameter_pattern = re.compile( - r'<｜DSML｜parameter name="(.*?)" string="(true|false)">' - r"(.*?)", + rf'<{parameter} name="(.*?)" string="(true|false)">(.*?)', re.DOTALL, ) raw = region.text @@ -1255,7 +1266,7 @@ def _parse_deepseek_v4_tool_call(region: TaggedToolBlock) -> ParsedToolCall: else: header = raw[: header_end + 2] name_match = re.fullmatch( - r'<｜DSML｜invoke name="(.*?)">\n', + rf'<{re.escape(invoke_tag)} name="(.*?)">\n', header, flags=re.DOTALL, ) diff --git a/renderers/tagged.py b/renderers/tagged.py index ccb4b7f8..8c850d5a 100644 --- a/renderers/tagged.py +++ b/renderers/tagged.py @@ -190,3 +190,15 @@ def scan_tagged_response( "": "tool_end", }, ) + +# vLLM's ``deepseek_v41`` parser: V4's rules with space-prefixed tag names and +# no misspelled section openers. +DEEPSEEK_V41_GRAMMAR = TaggedGrammar( + DEEPSEEK_V4_GRAMMAR.transitions, + { + "<｜DSML｜ calls>": "section_start", + "": "section_end", + "<｜DSML｜ invoke": "tool_start", + "": "tool_end", + }, +) diff --git a/tests/test_deepseek_v41.py b/tests/test_deepseek_v41.py new file mode 100644 index 00000000..7d77eb58 --- /dev/null +++ b/tests/test_deepseek_v41.py @@ -0,0 +1,392 @@ +"""DeepSeek-V4.1-Flash parity against the checkpoint's shipped Python encoder.""" + +from __future__ import annotations + +import importlib.util +import json +from functools import lru_cache + +import pytest +from huggingface_hub import hf_hub_download +from pydantic import TypeAdapter, ValidationError +from renderers import ( + DeepSeekV41Renderer, + DeepSeekV41RendererConfig, + RendererConfig, + ToolCallParseStatus, + create_renderer, +) +from renderers.base import load_tokenizer +from renderers.configs import config_from_name + +MODEL = "deepseek-ai/DeepSeek-V4.1-Flash" + +CONVERSATIONS = { + "user": [{"role": "user", "content": "What is 2+2?"}], + "system_user": [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "What is 2+2?"}, + ], + "multi_turn": [ + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello!", "reasoning_content": "greet"}, + {"role": "user", "content": "2+2?"}, + ], + "system_multi_turn_ending_in_assistant": [ + {"role": "system", "content": "Be terse."}, + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello!", "reasoning_content": "greet"}, + {"role": "user", "content": "2+2?"}, + {"role": "assistant", "content": "4", "reasoning_content": "add"}, + ], +} + + +@lru_cache(maxsize=1) +def _tokenizer(): + return load_tokenizer(MODEL) + + +@lru_cache(maxsize=1) +def _reference_encoder(): + path = hf_hub_download(MODEL, "encoding/encoding.py") + spec = importlib.util.spec_from_file_location("deepseek_v41_encoding", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _decode(renderer, messages, **kwargs): + return _tokenizer().decode( + renderer.render_ids(messages, **kwargs), skip_special_tokens=False + ) + + +def test_registration_and_native_defaults(): + # Selected by name only: the shared parity suites need tool support for auto-mapping. + renderer = create_renderer(_tokenizer(), config_from_name("deepseek-v4.1")) + + assert isinstance(renderer, DeepSeekV41Renderer) + assert renderer.config.enable_thinking is True + assert renderer.config.reasoning_effort == "high" + + +def test_config_discriminator_and_effort_validation(): + parsed = TypeAdapter(RendererConfig).validate_python( + {"name": "deepseek-v4.1", "reasoning_effort": 30} + ) + assert isinstance(parsed, DeepSeekV41RendererConfig) + assert parsed.reasoning_effort == 30 + + with pytest.raises(ValidationError): + DeepSeekV41RendererConfig(reasoning_effort=101) + + +@pytest.mark.parametrize("name", CONVERSATIONS) +@pytest.mark.parametrize("enable_thinking", [True, False]) +@pytest.mark.parametrize("reasoning_effort", ["low", "high", "max", 30]) +def test_matches_reference_encoder(name, enable_thinking, reasoning_effort): + messages = CONVERSATIONS[name] + add_generation_prompt = messages[-1]["role"] != "assistant" + renderer = DeepSeekV41Renderer( + _tokenizer(), + DeepSeekV41RendererConfig( + enable_thinking=enable_thinking, reasoning_effort=reasoning_effort + ), + ) + + expected = _reference_encoder().encode_messages( + [dict(message) for message in messages], + thinking_mode="thinking" if enable_thinking else "chat", + reasoning_effort=reasoning_effort, + ) + + assert ( + _decode(renderer, messages, add_generation_prompt=add_generation_prompt) + == expected + ) + + +def test_inline_reasoning_moves_into_think_block(): + renderer = DeepSeekV41Renderer(_tokenizer()) + inline = [ + {"role": "user", "content": "2+2?"}, + {"role": "assistant", "content": "add4"}, + ] + split = [ + {"role": "user", "content": "2+2?"}, + {"role": "assistant", "content": "4", "reasoning_content": "add"}, + ] + + assert renderer.render_ids(inline) == renderer.render_ids(split) + + +def test_rejects_mid_conversation_system(): + renderer = DeepSeekV41Renderer(_tokenizer()) + user = [{"role": "user", "content": "Hi"}] + + with pytest.raises(NotImplementedError): + renderer.render([*user, {"role": "system", "content": "Be terse."}]) + + +TOOLS = [ + { + "type": "function", + "function": { + "name": "weather", + "description": "Get weather", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string"}, + "days": {"type": "integer"}, + }, + "required": ["city"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "search", + "description": "Search the web", + "parameters": { + "type": "object", + "properties": { + "query": {"type": "string"}, + "filters": {"type": "object"}, + }, + "required": ["query"], + }, + }, + }, +] + + +def _call(id, name, arguments): + return { + "id": id, + "type": "function", + "function": {"name": name, "arguments": arguments}, + } + + +AGENT_TRAJECTORY = [ + {"role": "system", "content": "You are a travel agent."}, + {"role": "user", "content": "Weather in Paris and Rome?"}, + { + "role": "assistant", + "reasoning_content": "Call weather for both cities.", + "content": "", + "tool_calls": [ + _call("a", "weather", {"city": "Paris"}), + _call("b", "weather", {"city": "Rome", "days": 2}), + ], + }, + {"role": "tool", "tool_call_id": "b", "content": "Rome: sunny"}, + {"role": "tool", "tool_call_id": "a", "content": "Paris: rain"}, + { + "role": "assistant", + "reasoning_content": "Look for indoor activities in Paris.", + "content": "Let me search.", + "tool_calls": [ + _call( + "c", + "search", + {"query": "Paris museums", "filters": {"open": True, "max": 3}}, + ) + ], + }, + {"role": "tool", "tool_call_id": "c", "content": "Louvre, Orsay"}, + { + "role": "assistant", + "reasoning_content": "Summarize.", + "content": "Rome is sunny; Paris is rainy, try the Louvre.", + }, + {"role": "user", "content": "And Berlin?"}, + { + "role": "assistant", + "reasoning_content": "One more call.", + "content": "", + "tool_calls": [_call("d", "weather", {"city": "Berlin"})], + }, + {"role": "tool", "tool_call_id": "d", "content": "Berlin: snow"}, +] + +TOOL_CONVERSATIONS = { + "tool_call_without_system": [ + {"role": "user", "content": "Weather in Oslo?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [_call("x", "weather", {"city": "Oslo"})], + }, + {"role": "tool", "tool_call_id": "x", "content": '{"temp": -3}'}, + ], + "tool_result_then_user_text": [ + {"role": "user", "content": "Weather in Oslo?"}, + { + "role": "assistant", + "reasoning_content": "check", + "tool_calls": [_call("x", "weather", {"city": "Oslo"})], + }, + {"role": "tool", "tool_call_id": "x", "content": "cold"}, + {"role": "user", "content": "Also, be brief."}, + ], + "argument_encodings": [ + {"role": "system", "content": "Call tools."}, + {"role": "user", "content": "Go"}, + { + "role": "assistant", + "reasoning_content": "Many calls.", + "tool_calls": [ + _call("1", "weather", '{"city": "Oslo", "days": 3}'), + _call("2", "weather", json.dumps('{"city": "Lima"}')), + _call("3", "search", "[1, 2]"), + _call("4", "search", "not json"), + _call("5", "search", {"query": "", "filters": None}), + ], + }, + ], + "agent_trajectory": AGENT_TRAJECTORY, + "agent_trajectory_answered": AGENT_TRAJECTORY + + [ + { + "role": "assistant", + "reasoning_content": "Report.", + "content": "Berlin has snow.", + } + ], +} + + +def _reference_messages(messages, tools): + messages = [dict(message) for message in messages] + if tools: + if messages[0]["role"] != "system": + messages.insert(0, {"role": "system", "content": ""}) + messages[0]["tools"] = tools + return messages + + +@pytest.mark.parametrize("name", TOOL_CONVERSATIONS) +@pytest.mark.parametrize("with_tools", [True, False]) +@pytest.mark.parametrize("enable_thinking", [True, False]) +@pytest.mark.parametrize("drop_thinking", [True, False]) +def test_tool_conversations_match_reference_encoder( + name, with_tools, enable_thinking, drop_thinking +): + messages = TOOL_CONVERSATIONS[name] + tools = TOOLS if with_tools else None + add_generation_prompt = messages[-1]["role"] != "assistant" + renderer = DeepSeekV41Renderer( + _tokenizer(), + DeepSeekV41RendererConfig( + enable_thinking=enable_thinking, drop_thinking=drop_thinking + ), + ) + + expected = _reference_encoder().encode_messages( + _reference_messages(messages, tools), + thinking_mode="thinking" if enable_thinking else "chat", + drop_thinking=drop_thinking, + ) + + assert ( + _decode( + renderer, + messages, + tools=tools, + add_generation_prompt=add_generation_prompt, + ) + == expected + ) + + +@pytest.mark.parametrize("enable_thinking", [True, False]) +def test_agent_trajectory_extends_by_prefix_and_bridge(enable_thinking): + renderer = DeepSeekV41Renderer( + _tokenizer(), DeepSeekV41RendererConfig(enable_thinking=enable_thinking) + ) + messages = TOOL_CONVERSATIONS["agent_trajectory_answered"] + assistant_indices = [ + index + for index, message in enumerate(messages) + if message["role"] == "assistant" + ] + + for index, next_index in zip(assistant_indices, assistant_indices[1:]): + prompt = renderer.render_ids( + messages[:index], tools=TOOLS, add_generation_prompt=True + ) + completed = renderer.render_ids(messages[: index + 1], tools=TOOLS) + assert completed[: len(prompt)] == prompt + + new_messages = messages[index + 1 : next_index] + bridged = renderer.bridge_to_next_turn( + prompt, completed[len(prompt) :], new_messages, tools=TOOLS + ) + # Parallel results need the issuing call order; a new user query + # drops tool-cycle reasoning. Both re-render instead. + declines = sum(m["role"] == "tool" for m in new_messages) > 1 or ( + enable_thinking and any(m["role"] == "user" for m in new_messages) + ) + if declines: + assert bridged is None + else: + assert bridged is not None + assert bridged.token_ids == renderer.render_ids( + messages[:next_index], tools=TOOLS, add_generation_prompt=True + ) + + +@pytest.mark.parametrize("enable_thinking", [True, False]) +def test_tool_calls_roundtrip_through_parse_response(enable_thinking): + renderer = DeepSeekV41Renderer( + _tokenizer(), DeepSeekV41RendererConfig(enable_thinking=enable_thinking) + ) + messages = AGENT_TRAJECTORY[:3] + prompt = renderer.render_ids(messages[:2], tools=TOOLS, add_generation_prompt=True) + completed = renderer.render_ids(messages, tools=TOOLS) + completion = completed[len(prompt) :] + + parsed = renderer.parse_response(completion, prompt_ids=prompt) + + expected_reasoning = messages[2]["reasoning_content"] if enable_thinking else None + assert parsed.reasoning_content == expected_reasoning + assert parsed.content == "" + assert [(call.name, call.arguments, call.status) for call in parsed.tool_calls] == [ + ("weather", {"city": "Paris"}, ToolCallParseStatus.OK), + ("weather", {"city": "Rome", "days": 2}, ToolCallParseStatus.OK), + ] + for call in parsed.tool_calls: + start, end = call.token_span + assert ( + _tokenizer() + .decode(completion[start:end], skip_special_tokens=False) + .rstrip("\n") + == call.raw + ) + + +def test_parses_vllm_tool_block_without_leading_newlines(): + text = ( + "<｜DSML｜ calls>\n" + '<｜DSML｜ invoke name="search">\n' + '<｜DSML｜ parameter name="query" string="true">a "b"\n' + '<｜DSML｜ parameter name="filters" string="false">{"k": [1]}\n' + "\n" + "<｜end▁of▁sentence｜>" + ) + renderer = DeepSeekV41Renderer( + _tokenizer(), DeepSeekV41RendererConfig(enable_thinking=False) + ) + + parsed = renderer.parse_response( + _tokenizer().encode(text, add_special_tokens=False) + ) + + assert parsed.content == "" + assert [(call.name, call.arguments, call.status) for call in parsed.tool_calls] == [ + ("search", {"query": 'a "b"', "filters": {"k": [1]}}, ToolCallParseStatus.OK) + ]