abs_reasoning_parsers.py 9.34 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
import importlib
5
import os
6
from abc import abstractmethod
7
from collections.abc import Callable, Sequence
8
from functools import cached_property
9
from typing import TYPE_CHECKING, Any
10

11
from vllm.entrypoints.tool_server import ToolServer
12
from vllm.logger import init_logger
13
from vllm.utils.collection_utils import is_list_of
14
from vllm.utils.import_utils import import_from_path
15

16
if TYPE_CHECKING:
17
18
19
20
21
    from vllm.entrypoints.openai.protocol import (
        ChatCompletionRequest,
        DeltaMessage,
        ResponsesRequest,
    )
22
    from vllm.tokenizers import TokenizerLike
23
24
25
26
else:
    ChatCompletionRequest = Any
    DeltaMessage = Any
    ResponsesRequest = Any
27
    TokenizerLike = Any
28

29
30
31
32
33
logger = init_logger(__name__)


class ReasoningParser:
    """
34
    Abstract reasoning parser class that should not be used directly.
35
36
37
38
39
    Provided and methods should be used in derived classes.

    It is used to extract reasoning content from the model output.
    """

40
    def __init__(self, tokenizer: TokenizerLike, *args, **kwargs):
41
42
43
        self.model_tokenizer = tokenizer

    @cached_property
44
    def vocab(self) -> dict[str, int]:
45
46
47
48
        # NOTE: Only PreTrainedTokenizerFast is guaranteed to have .vocab
        # whereas all tokenizers have .get_vocab()
        return self.model_tokenizer.get_vocab()

49
    @abstractmethod
50
    def is_reasoning_end(self, input_ids: list[int]) -> bool:
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
        """
        Check if the reasoning content ends in the input_ids.

        It is used in structured engines like `xgrammar` to check if the
        reasoning content ends in the model output.

        Parameters:
        input_ids: list[int]
            The input_ids of the model output.

        Returns:
        bool
            True if the reasoning content ends in the input_ids.
        """

    @abstractmethod
    def extract_content_ids(self, input_ids: list[int]) -> list[int]:
        """
        Extract content token ids from the input_ids.
        Parameters:
        input_ids: list[int]
            The input_ids of the model output.
        Returns:
        list[int]
            The extracted content from the input_ids.
        """

    @abstractmethod
79
    def extract_reasoning(
80
81
        self,
        model_output: str,
82
        request: ChatCompletionRequest | ResponsesRequest,
83
    ) -> tuple[str | None, str | None]:
84
85
86
87
88
89
90
91
92
93
94
95
96
97
        """
        Extract reasoning content from a complete model-generated string.

        Used for non-streaming responses where we have the entire model response
        available before sending to the client.

        Parameters:
        model_output: str
            The model-generated string to extract reasoning content from.

        request: ChatCompletionRequest
            The request object that was used to generate the model_output.

        Returns:
98
        tuple[Optional[str], Optional[str]]
99
100
101
            A tuple containing the reasoning content and the content.
        """

102
    @abstractmethod
103
    def extract_reasoning_streaming(
104
105
106
107
108
109
110
        self,
        previous_text: str,
        current_text: str,
        delta_text: str,
        previous_token_ids: Sequence[int],
        current_token_ids: Sequence[int],
        delta_token_ids: Sequence[int],
111
    ) -> DeltaMessage | None:
112
113
114
115
116
117
118
        """
        Instance method that should be implemented for extracting reasoning
        from an incomplete response; for use when handling reasoning calls and
        streaming. Has to be an instance method because  it requires state -
        the current tokens/diffs, but also the information about what has
        previously been parsed and extracted (see constructor)
        """
119

120
121
122
123
124
125
126
127
128
129
130
    def prepare_structured_tag(
        self,
        original_tag: str | None,
        tool_server: ToolServer | None,
    ) -> str:
        """
        Instance method that is implemented for preparing the structured tag
        Otherwise, None is returned
        """
        return None

131
132

class ReasoningParserManager:
133
134
135
136
137
138
139
140
141
142
143
144
    """
    Central registry for ReasoningParser implementations.

    Supports two registration modes:
      - Eager registration via `register_module`
      - Lazy registration via `register_lazy_module`

    Each reasoning parser must inherit from `ReasoningParser`.
    """

    reasoning_parsers: dict[str, type[ReasoningParser]] = {}
    lazy_parsers: dict[str, tuple[str, str]] = {}  # name -> (module_path, class_name)
145
146

    @classmethod
147
    def get_reasoning_parser(cls, name: str) -> type[ReasoningParser]:
148
        """
149
150
151
152
        Retrieve a registered or lazily registered ReasoningParser class.

        If the parser is lazily registered, it will be imported and cached
        on first access.
153

154
155
        Raises:
            KeyError: if no parser is found under the given name.
156
157
158
159
        """
        if name in cls.reasoning_parsers:
            return cls.reasoning_parsers[name]

160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
        if name in cls.lazy_parsers:
            return cls._load_lazy_parser(name)

        raise KeyError(f"Reasoning parser '{name}' not found.")

    @classmethod
    def list_registered(cls) -> list[str]:
        """Return names of all eagerly and lazily registered reasoning parsers."""
        return sorted(set(cls.reasoning_parsers.keys()) | set(cls.lazy_parsers.keys()))

    @classmethod
    def _load_lazy_parser(cls, name: str) -> type[ReasoningParser]:
        """Import and register a lazily loaded reasoning parser."""
        module_path, class_name = cls.lazy_parsers[name]
        try:
            mod = importlib.import_module(module_path)
            parser_cls = getattr(mod, class_name)
            if not issubclass(parser_cls, ReasoningParser):
                raise TypeError(
                    f"{class_name} in {module_path} is not a ReasoningParser subclass."
                )

            cls.reasoning_parsers[name] = parser_cls  # cache
            return parser_cls
        except Exception as e:
            logger.exception(
                "Failed to import lazy reasoning parser '%s' from %s: %s",
                name,
                module_path,
                e,
            )
            raise
192
193

    @classmethod
194
195
    def _register_module(
        cls,
196
        module: type[ReasoningParser],
197
        module_name: str | list[str] | None = None,
198
199
        force: bool = True,
    ) -> None:
200
        """Register a ReasoningParser class immediately."""
201
        if not issubclass(module, ReasoningParser):
202
203
204
            raise TypeError(
                f"module must be subclass of ReasoningParser, but got {type(module)}"
            )
205

206
        if module_name is None:
207
208
209
210
211
212
213
214
215
            module_names = [module.__name__]
        elif isinstance(module_name, str):
            module_names = [module_name]
        elif is_list_of(module_name, str):
            module_names = module_name
        else:
            raise TypeError("module_name must be str, list[str], or None.")

        for name in module_names:
216
            if not force and name in cls.reasoning_parsers:
217
218
                existed = cls.reasoning_parsers[name]
                raise KeyError(f"{name} is already registered at {existed.__module__}")
219
220
            cls.reasoning_parsers[name] = module

221
222
223
224
225
226
227
228
229
230
231
232
233
234
    @classmethod
    def register_lazy_module(cls, name: str, module_path: str, class_name: str) -> None:
        """
        Register a lazy module mapping for delayed import.

        Example:
            ReasoningParserManager.register_lazy_module(
                name="qwen3",
                module_path="vllm.reasoning.parsers.qwen3_reasoning_parser",
                class_name="Qwen3ReasoningParser",
            )
        """
        cls.lazy_parsers[name] = (module_path, class_name)

235
236
    @classmethod
    def register_module(
237
        cls,
238
        name: str | list[str] | None = None,
239
        force: bool = True,
240
241
242
243
        module: type[ReasoningParser] | None = None,
    ) -> (
        type[ReasoningParser] | Callable[[type[ReasoningParser]], type[ReasoningParser]]
    ):
244
245
        """
        Register module with the given name or name list. it can be used as a
246
        decoder(with module as None) or normal function(with module as not
247
248
249
250
251
        None).
        """
        if not isinstance(force, bool):
            raise TypeError(f"force must be a boolean, but got {type(force)}")

252
        # Immediate registration (explicit call)
253
254
255
256
        if module is not None:
            cls._register_module(module=module, module_name=name, force=force)
            return module

257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
        # Decorator usage
        def _decorator(obj: type[ReasoningParser]) -> type[ReasoningParser]:
            module_path = obj.__module__
            class_name = obj.__name__

            if isinstance(name, str):
                names = [name]
            elif is_list_of(name, str):
                names = name
            else:
                names = [class_name]

            for n in names:
                cls.lazy_parsers[n] = (module_path, class_name)

            return obj
273

274
        return _decorator
275
276
277
278

    @classmethod
    def import_reasoning_parser(cls, plugin_path: str) -> None:
        """
279
        Import a user-defined reasoning parser by the path
280
281
282
283
284
285
286
        of the reasoning parser define file.
        """
        module_name = os.path.splitext(os.path.basename(plugin_path))[0]

        try:
            import_from_path(module_name, plugin_path)
        except Exception:
287
288
289
            logger.exception(
                "Failed to load module '%s' from %s.", module_name, plugin_path
            )
290
            return