| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import string |
| 5 | from collections.abc import Sequence |
| 6 | from typing import cast |
| 7 | |
| 8 | import pytest |
| 9 | |
| 10 | from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest |
| 11 | from vllm.parser.parser_manager import ParserManager |
| 12 | from vllm.reasoning import ReasoningParserManager |
| 13 | from vllm.reasoning.k2_horizon_reasoning_parser import K2HorizonReasoningParser |
| 14 | from vllm.tokenizers import TokenizerLike |
| 15 | |
| 16 | pytestmark = pytest.mark.skip_global_cleanup |
| 17 | |
| 18 | EFFORT_TOKENS = { |
| 19 | "high": ("<ifm|think>", "</ifm|think>"), |
| 20 | "medium": ("<ifm|think_fast>", "</ifm|think_fast>"), |
| 21 | "low": ("<ifm|think_faster>", "</ifm|think_faster>"), |
| 22 | } |
| 23 | TOOL_CALLS_START = "<ifm|tool_calls>" |
| 24 | |
| 25 | |
| 26 | class K2Tokenizer: |
| 27 | def __init__(self) -> None: |
| 28 | special_tokens = [ |
| 29 | token for token_pair in EFFORT_TOKENS.values() for token in token_pair |
| 30 | ] + [TOOL_CALLS_START] |
| 31 | tokens = special_tokens + list(string.printable) |
| 32 | self._vocab = {token: index for index, token in enumerate(tokens, start=1)} |
| 33 | |
| 34 | def get_vocab(self) -> dict[str, int]: |
| 35 | return dict(self._vocab) |
| 36 | |
| 37 | def encode(self, text: str) -> list[int]: |
| 38 | tokens: list[int] = [] |
| 39 | position = 0 |
| 40 | specials = sorted(EFFORT_TOKENS["high"] + (TOOL_CALLS_START,), key=len) |
| 41 | while position < len(text): |
| 42 | for special in specials: |
| 43 | if text.startswith(special, position): |
| 44 | tokens.append(self._vocab[special]) |
| 45 | position += len(special) |
| 46 | break |
| 47 | else: |
| 48 | tokens.append(self._vocab[text[position]]) |
| 49 | position += 1 |
| 50 | return tokens |
| 51 | |
| 52 | def decode( |
| 53 | self, |
| 54 | ids: Sequence[int] | int, |
| 55 | skip_special_tokens: bool = False, |
| 56 | ) -> str: |
| 57 | del skip_special_tokens |
| 58 | if isinstance(ids, int): |
| 59 | ids = [ids] |
| 60 | inverse_vocab = {token_id: token for token, token_id in self._vocab.items()} |
| 61 | return "".join(inverse_vocab[token_id] for token_id in ids) |
| 62 | |
| 63 | |
| 64 | def _k2_tokenizer() -> TokenizerLike: |
| 65 | return cast(TokenizerLike, K2Tokenizer()) |
| 66 | |
| 67 | |
| 68 | def _request(**kwargs) -> ChatCompletionRequest: |
| 69 | return ChatCompletionRequest( |
| 70 | model="k2-horizon", |
| 71 | messages=[{"role": "user", "content": "test"}], |
| 72 | **kwargs, |
| 73 | ) |
| 74 | |
| 75 | |
| 76 | @pytest.mark.parametrize("effort", ["high", "medium", "low"]) |
| 77 | def test_reasoning_effort_selects_tokens(effort: str): |
| 78 | parser = K2HorizonReasoningParser( |
| 79 | _k2_tokenizer(), |
| 80 | chat_template_kwargs={"reasoning_effort": effort}, |
| 81 | ) |
| 82 | |
| 83 | assert (parser.start_token, parser.end_token) == EFFORT_TOKENS[effort] |
| 84 | |
| 85 | |
| 86 | def test_reasoning_effort_precedence_and_default(): |
| 87 | tokenizer = _k2_tokenizer() |
| 88 | request = _request( |
| 89 | reasoning_effort="medium", |
| 90 | chat_template_kwargs={"reasoning_effort": "low"}, |
| 91 | ) |
| 92 | resolved_kwargs = ( |
| 93 | request.build_chat_params(None, "auto") |
| 94 | .with_defaults({"reasoning_effort": "high"}) |
| 95 | .chat_template_kwargs |
| 96 | ) |
| 97 | |
| 98 | parser = K2HorizonReasoningParser( |
| 99 | tokenizer, |
| 100 | chat_template_kwargs=resolved_kwargs, |
| 101 | ) |
| 102 | default_parser = K2HorizonReasoningParser(tokenizer) |
| 103 | |
| 104 | assert parser.start_token == EFFORT_TOKENS["medium"][0] |
| 105 | assert default_parser.start_token == EFFORT_TOKENS["high"][0] |
| 106 | |
| 107 | |
| 108 | @pytest.mark.parametrize("effort", ["none", "unknown", None]) |
| 109 | def test_invalid_reasoning_effort_raises(effort: str | None): |
| 110 | with pytest.raises(ValueError, match="Unsupported reasoning_effort"): |
| 111 | K2HorizonReasoningParser( |
| 112 | _k2_tokenizer(), |
| 113 | chat_template_kwargs={"reasoning_effort": effort}, |
| 114 | ) |
| 115 | |
| 116 | |
| 117 | def test_non_streaming_explicit_and_implicit_boundaries(): |
| 118 | parser = K2HorizonReasoningParser(_k2_tokenizer()) |
| 119 | request = _request() |
| 120 | |
| 121 | assert parser.extract_reasoning("<ifm|think>plan</ifm|think>answer", request) == ( |
| 122 | "plan", |
| 123 | "answer", |
| 124 | ) |
| 125 | assert parser.extract_reasoning( |
| 126 | f"plan{TOOL_CALLS_START}<ifm|tool_call>ping</ifm|tool_call>", request |
| 127 | ) == ( |
| 128 | "plan", |
| 129 | f"{TOOL_CALLS_START}<ifm|tool_call>ping</ifm|tool_call>", |
| 130 | ) |
| 131 | |
| 132 | |
| 133 | @pytest.mark.parametrize("reasoning", ["plan", " ", "\n", " \n"]) |
| 134 | def test_character_split_reasoning_stream(reasoning: str): |
| 135 | parser = K2HorizonReasoningParser(_k2_tokenizer()) |
| 136 | reasoning_parts: list[str] = [] |
| 137 | content_parts: list[str] = [] |
| 138 | |
| 139 | for char in f"{reasoning}</ifm|think>answer": |
| 140 | delta = parser.extract_reasoning_streaming( |
| 141 | previous_text="", |
| 142 | current_text="", |
| 143 | delta_text=char, |
| 144 | previous_token_ids=[], |
| 145 | current_token_ids=[], |
| 146 | delta_token_ids=[], |
| 147 | ) |
| 148 | if delta is not None: |
| 149 | if delta.reasoning is not None: |
| 150 | reasoning_parts.append(delta.reasoning) |
| 151 | if delta.content is not None: |
| 152 | content_parts.append(delta.content) |
| 153 | |
| 154 | assert "".join(reasoning_parts) == reasoning |
| 155 | assert "".join(content_parts) == "answer" |
| 156 | |
| 157 | |
| 158 | def test_latest_boundary_controls_reasoning_state(): |
| 159 | tokenizer = _k2_tokenizer() |
| 160 | parser = K2HorizonReasoningParser(tokenizer) |
| 161 | vocab = tokenizer.get_vocab() |
| 162 | |
| 163 | assert parser.is_reasoning_end([vocab["<ifm|think>"], vocab[TOOL_CALLS_START]]) |
| 164 | assert not parser.is_reasoning_end([vocab[TOOL_CALLS_START], vocab["<ifm|think>"]]) |
| 165 | |
| 166 | |
| 167 | def test_composed_streaming_reasoning_to_tool_handoff(): |
| 168 | tokenizer = _k2_tokenizer() |
| 169 | tools = [ |
| 170 | { |
| 171 | "type": "function", |
| 172 | "function": { |
| 173 | "name": "ping", |
| 174 | "parameters": {"type": "object", "properties": {}}, |
| 175 | }, |
| 176 | } |
| 177 | ] |
| 178 | request = _request(tools=tools, tool_choice="auto") |
| 179 | parser_cls = ParserManager.get_parser( |
| 180 | tool_parser_name="k2_horizon", |
| 181 | reasoning_parser_name="k2_horizon", |
| 182 | enable_auto_tools=True, |
| 183 | ) |
| 184 | assert parser_cls is not None |
| 185 | parser = parser_cls( |
| 186 | tokenizer, |
| 187 | request.tools, |
| 188 | chat_template_kwargs={"reasoning_effort": "high"}, |
| 189 | ) |
| 190 | output = "plan<ifm|tool_calls><ifm|tool_call>ping</ifm|tool_call></ifm|tool_calls>" |
| 191 | reasoning_parts: list[str] = [] |
| 192 | tool_names: list[str] = [] |
| 193 | |
| 194 | for index, char in enumerate(output): |
| 195 | delta = parser.parse_delta( |
| 196 | delta_text=char, |
| 197 | delta_token_ids=tokenizer.encode(char), |
| 198 | request=request, |
| 199 | prompt_token_ids=[] if index == 0 else None, |
| 200 | finished=index == len(output) - 1, |
| 201 | ) |
| 202 | if delta is not None: |
| 203 | if delta.reasoning is not None: |
| 204 | reasoning_parts.append(delta.reasoning) |
| 205 | tool_names.extend( |
| 206 | tool_call.function.name |
| 207 | for tool_call in delta.tool_calls |
| 208 | if tool_call.function is not None |
| 209 | and tool_call.function.name is not None |
| 210 | ) |
| 211 | |
| 212 | assert "".join(reasoning_parts) == "plan" |
| 213 | assert tool_names == ["ping"] |
| 214 | |
| 215 | |
| 216 | def test_reasoning_parser_registered(): |
| 217 | assert ( |
| 218 | ReasoningParserManager.get_reasoning_parser("k2_horizon") |
| 219 | is K2HorizonReasoningParser |
| 220 | ) |