main
py 364 lines 12.6 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import importlib
5 from abc import abstractmethod
6 from collections.abc import Callable, Iterable, Sequence
7 from functools import cached_property
8 from typing import TYPE_CHECKING, cast
9
10 from vllm.entrypoints.mcp.tool_server import ToolServer
11 from vllm.logger import init_logger
12 from vllm.utils.collection_utils import is_list_of
13 from vllm.utils.import_utils import import_plugin
14
15 if TYPE_CHECKING:
16 from vllm.config import ModelConfig
17 from vllm.entrypoints.generate.base.protocol import DeltaMessage
18 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
19 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
20 from vllm.tokenizers import TokenizerLike
21
22 logger = init_logger(__name__)
23
24
25 class ReasoningParser:
26 """Abstract reasoning parser class that should not be used directly.
27 Provided and methods should be used in derived classes.
28
29 It is used to extract reasoning content from the model output.
30 """
31
32 engine_based_streaming: bool = False
33
34 def __init__(self, tokenizer: "TokenizerLike", *args, **kwargs):
35 self.model_tokenizer = tokenizer
36 # Optional vLLM ModelConfig from the server. Use get (not pop) so composite
37 # parsers can forward **kwargs to nested parsers.
38 self._model_config: ModelConfig | None = kwargs.get("model_config")
39
40 @cached_property
41 def vocab(self) -> dict[str, int]:
42 # NOTE: Only TokenizersBackend is guaranteed to have .vocab
43 # whereas all tokenizers have .get_vocab()
44 return self.model_tokenizer.get_vocab()
45
46 @property
47 def reasoning_start_str(self) -> str | None:
48 """Set `reasoning_start_str` to the strings that delimit
49 the reasoning block (e.g. `""<seed:think>""` and `"<think>"`).
50 """
51 return None
52
53 @property
54 def reasoning_end_str(self) -> str | None:
55 """Set `reasoning_end_str` to the strings that delimit
56 the reasoning block (e.g. `""</seed:think>""` and `"</think>"`).
57 """
58 return None
59
60 def has_engine_confirmed_reasoning_end(self) -> bool:
61 """Whether the engine has confirmed the reasoning end transition.
62
63 Engine-based parsers may defer terminal processing when the
64 detokenizer holds back text. This method returns the engine's
65 *processed* state, not a raw token-ID check.
66
67 Only called for parsers with ``engine_based_streaming = True``.
68 """
69 return False
70
71 @abstractmethod
72 def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
73 """Check if the reasoning content ends in the input_ids.
74
75 It is used in structured engines like `xgrammar` to check if the
76 reasoning content ends in the model output.
77
78 Parameters
79 ----------
80 input_ids: list[int]
81 The input_ids of the model output.
82
83 Returns
84 -------
85 bool
86 True if the reasoning content ends in the input_ids.
87
88 """
89
90 def is_reasoning_end_streaming(
91 self, input_ids: Sequence[int], delta_ids: Iterable[int]
92 ) -> bool:
93 """Check if the reasoning content ends in the input_ids on a
94 decode step.
95
96 It is used in structured engines like `xgrammar` to check if the
97 reasoning content ends in the model output during a decode step.
98 `input_ids` the entire model output and `delta_ids` are the last few
99 computed tokens of the model output (like during a decode step).
100
101 Parameters
102 ----------
103 input_ids: list[int]
104 The entire model output.
105 delta_ids: list[int]
106 The last few computed tokens of the model output at the current decode step.
107
108 Returns
109 -------
110 bool
111 True if the reasoning content ends in the `delta_ids` on a
112 decode step.
113
114 """
115 return self.is_reasoning_end(input_ids)
116
117 @abstractmethod
118 def extract_content_ids(self, input_ids: list[int]) -> list[int]:
119 """Extract content token ids from the input_ids.
120
121 Parameters
122 ----------
123 input_ids: list[int]
124 The input_ids of the model output.
125
126 Returns
127 -------
128 list[int]
129 The extracted content from the input_ids.
130
131 """
132
133 def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
134 """Count the number of reasoning tokens in a sequence.
135
136 Text-based reasoning models typically wrap their chain-of-thought
137 between special start/end tokens (e.g., ``<think> ... </think>``).
138 Implementations that support reasoning token counting should override
139 this method. The default implementation returns ``0`` so existing
140 parsers remain unchanged unless they explicitly opt in.
141
142 Args:
143 token_ids: Sequence of generated token ids (excluding prompt).
144
145 Returns:
146 int: Number of tokens that belong to reasoning content.
147
148 """
149 # By default, assume the parser cannot detect reasoning spans.
150 return 0
151
152 @abstractmethod
153 def extract_reasoning(
154 self,
155 model_output: str,
156 request: "ChatCompletionRequest | ResponsesRequest",
157 ) -> tuple[str | None, str | None]:
158 """Extract reasoning content from a complete model-generated string.
159
160 Used for non-streaming responses where we have the entire model response
161 available before sending to the client.
162
163 Args:
164 model_output: The model-generated string to extract reasoning content from.
165 request: The request object that was used to generate the model_output.
166
167 Returns:
168 A tuple containing the reasoning content and the content.
169
170 """
171
172 @abstractmethod
173 def extract_reasoning_streaming(
174 self,
175 previous_text: str,
176 current_text: str,
177 delta_text: str,
178 previous_token_ids: Sequence[int],
179 current_token_ids: Sequence[int],
180 delta_token_ids: Sequence[int],
181 ) -> "DeltaMessage | None":
182 """Instance method that should be implemented for extracting reasoning
183 from an incomplete response; for use when handling reasoning calls and
184 streaming. Has to be an instance method because it requires state -
185 the current tokens/diffs, but also the information about what has
186 previously been parsed and extracted (see constructor)
187 """
188
189 def adjust_request(
190 self, request: "ChatCompletionRequest | ResponsesRequest"
191 ) -> "ChatCompletionRequest | ResponsesRequest":
192 """Adjust request parameters; override in subclasses as needed."""
193 return request
194
195 def adjust_initial_state_from_prompt(self, prompt_token_ids: Sequence[int]) -> None:
196 """Hook called once at the start of streaming with the prompt tokens.
197
198 Gives parsers a chance to adjust their initial parsing state based on
199 the prompt — for example, when the chat template leaves the prompt
200 inside an open reasoning channel and the engine's default initial
201 state would otherwise misclassify the first generated tokens.
202
203 Default is a no-op; override in subclasses as needed.
204 """
205 return
206
207 def prepare_structured_tag(
208 self,
209 original_tag: str | None,
210 tool_server: ToolServer | None,
211 ) -> str | None:
212 """Instance method that is implemented for preparing the structured tag"""
213 return original_tag
214
215
216 class ReasoningParserManager:
217 """Central registry for ReasoningParser implementations.
218
219 Supports two registration modes:
220 - Eager registration via `register_module`
221 - Lazy registration via `register_lazy_module`
222
223 Each reasoning parser must inherit from `ReasoningParser`.
224 """
225
226 reasoning_parsers: dict[str, type[ReasoningParser]] = {}
227 lazy_parsers: dict[str, tuple[str, str]] = {} # name -> (module_path, class_name)
228
229 @classmethod
230 def get_reasoning_parser(cls, name: str) -> type[ReasoningParser]:
231 """Retrieve a registered or lazily registered ReasoningParser class.
232
233 If the parser is lazily registered, it will be imported and cached
234 on first access.
235
236 Raises:
237 KeyError: if no parser is found under the given name.
238
239 """
240 if name in cls.reasoning_parsers:
241 return cls.reasoning_parsers[name]
242
243 if name in cls.lazy_parsers:
244 return cls._load_lazy_parser(name)
245
246 registered = ", ".join(cls.list_registered())
247 raise KeyError(
248 f"Reasoning parser '{name}' not found. Available parsers: {registered}"
249 )
250
251 @classmethod
252 def list_registered(cls) -> list[str]:
253 """Return names of all eagerly and lazily registered reasoning parsers."""
254 return sorted(set(cls.reasoning_parsers.keys()) | set(cls.lazy_parsers.keys()))
255
256 @classmethod
257 def _load_lazy_parser(cls, name: str) -> type[ReasoningParser]:
258 """Import and register a lazily loaded reasoning parser."""
259 module_path, class_name = cls.lazy_parsers[name]
260 try:
261 mod = importlib.import_module(module_path)
262 parser_cls = getattr(mod, class_name)
263 if not issubclass(parser_cls, ReasoningParser):
264 raise TypeError(
265 f"{class_name} in {module_path} is not a ReasoningParser subclass."
266 )
267
268 cls.reasoning_parsers[name] = parser_cls # cache
269 return parser_cls
270 except Exception as e:
271 logger.exception(
272 "Failed to import lazy reasoning parser '%s' from %s: %s",
273 name,
274 module_path,
275 e,
276 )
277 raise
278
279 @classmethod
280 def _register_module(
281 cls,
282 module: type[ReasoningParser],
283 module_name: str | list[str] | None = None,
284 force: bool = True,
285 ) -> None:
286 """Register a ReasoningParser class immediately."""
287 if not issubclass(module, ReasoningParser):
288 raise TypeError(
289 f"module must be subclass of ReasoningParser, but got {type(module)}"
290 )
291
292 if module_name is None:
293 module_names = [module.__name__]
294 elif isinstance(module_name, str):
295 module_names = [module_name]
296 elif is_list_of(module_name, str):
297 module_names = module_name
298 else:
299 raise TypeError("module_name must be str, list[str], or None.")
300
301 for name in module_names:
302 if not force and name in cls.reasoning_parsers:
303 existed = cls.reasoning_parsers[name]
304 raise KeyError(f"{name} is already registered at {existed.__module__}")
305 cls.reasoning_parsers[name] = module
306
307 @classmethod
308 def register_lazy_module(cls, name: str, module_path: str, class_name: str) -> None:
309 """Register a lazy module mapping for delayed import.
310
311 Example:
312 ReasoningParserManager.register_lazy_module(
313 name="qwen3",
314 module_path="vllm.reasoning.qwen3_engine_reasoning_parser",
315 class_name="Qwen3ParserReasoningAdapter",
316 )
317
318 """
319 cls.lazy_parsers[name] = (module_path, class_name)
320
321 @classmethod
322 def register_module(
323 cls,
324 name: str | list[str] | None = None,
325 force: bool = True,
326 module: type[ReasoningParser] | None = None,
327 ) -> (
328 type[ReasoningParser] | Callable[[type[ReasoningParser]], type[ReasoningParser]]
329 ):
330 """Register module with the given name or name list. it can be used as a
331 decoder(with module as None) or normal function(with module as not
332 None).
333 """
334 if not isinstance(force, bool):
335 raise TypeError(f"force must be a boolean, but got {type(force)}")
336
337 # Immediate registration (explicit call)
338 if module is not None:
339 cls._register_module(module=module, module_name=name, force=force)
340 return module
341
342 # Decorator usage
343 def _decorator(obj: type[ReasoningParser]) -> type[ReasoningParser]:
344 module_path = obj.__module__
345 class_name = obj.__name__
346
347 if isinstance(name, str):
348 names = [name]
349 elif is_list_of(name, str):
350 names = cast(list[str], name)
351 else:
352 names = [class_name]
353
354 for n in names:
355 cls.lazy_parsers[n] = (module_path, class_name)
356
357 return obj
358
359 return _decorator
360
361 @classmethod
362 def import_reasoning_parser(cls, plugin_path: str) -> None:
363 """Import a user-defined reasoning parser."""
364 import_plugin(plugin_path)