abstract_tool_parser.py 9.73 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 collections.abc import Callable, Sequence
7
from functools import cached_property
8

9
from openai.types.responses import (
10
    ResponseFormatTextJSONSchemaConfig,
11
    ResponseTextConfig,
12
13
)

14
15
16
from vllm.entrypoints.openai.chat_completion.protocol import (
    ChatCompletionRequest,
)
17
from vllm.entrypoints.openai.engine.protocol import (
18
19
    DeltaMessage,
    ExtractedToolCallInformation,
20
21
)
from vllm.entrypoints.openai.responses.protocol import (
22
    ResponsesRequest,
23
)
24
from vllm.logger import init_logger
25
26
27
from vllm.sampling_params import (
    StructuredOutputsParams,
)
28
from vllm.tokenizers import TokenizerLike
29
from vllm.tool_parsers.utils import Tool, get_json_schema_from_tools
30
from vllm.utils.collection_utils import is_list_of
31
from vllm.utils.import_utils import import_from_path
32

33
__all__ = ["Tool"]
34

35
logger = init_logger(__name__)
36

37
38
39
40
41
42
43
44

class ToolParser:
    """
    Abstract ToolParser class that should not be used directly. Provided
    properties and methods should be used in
    derived classes.
    """

45
46
47
48
49
    def __init__(
        self,
        tokenizer: TokenizerLike,
        tools: list[Tool] | None = None,
    ):
50
        self.prev_tool_call_arr: list[dict] = []
51
52
53
        # the index of the tool call that is currently being parsed
        self.current_tool_id: int = -1
        self.current_tool_name_sent: bool = False
54
        self.streamed_args_for_tool: list[str] = []
55
56

        self.model_tokenizer = tokenizer
57
        self.tools = tools
58

59
    @cached_property
60
    def vocab(self) -> dict[str, int]:
61
62
63
64
        # NOTE: Only PreTrainedTokenizerFast is guaranteed to have .vocab
        # whereas all tokenizers have .get_vocab()
        return self.model_tokenizer.get_vocab()

65
66
67
    def adjust_request(
        self, request: ChatCompletionRequest | ResponsesRequest
    ) -> ChatCompletionRequest | ResponsesRequest:
68
69
70
        """
        Static method that used to adjust the request parameters.
        """
71
72
73
74
75
76
77
        if not request.tools:
            return request
        json_schema_from_tool = get_json_schema_from_tools(
            tool_choice=request.tool_choice, tools=request.tools
        )
        # Set structured output params for tool calling
        if json_schema_from_tool is not None:
78
79
80
            if isinstance(request, ChatCompletionRequest):
                # tool_choice: "Forced Function" or "required" will override
                # structured output json settings to make tool calling work correctly
81
                request.structured_outputs = StructuredOutputsParams(
82
                    json=json_schema_from_tool  # type: ignore[call-arg]
83
                )
84
                request.response_format = None
85
86
87
88
89
90
91
92
93
94
            if isinstance(request, ResponsesRequest):
                request.text = ResponseTextConfig()
                request.text.format = ResponseFormatTextJSONSchemaConfig(
                    name="tool_calling_response",
                    schema=json_schema_from_tool,
                    type="json_schema",
                    description="Response format for tool calling",
                    strict=True,
                )

95
96
97
        return request

    def extract_tool_calls(
98
99
        self, model_output: str, request: ChatCompletionRequest
    ) -> ExtractedToolCallInformation:
100
101
102
103
104
105
106
107
        """
        Static method that should be implemented for extracting tool calls from
        a complete model-generated string.
        Used for non-streaming responses where we have the entire model response
        available before sending to the client.
        Static because it's stateless.
        """
        raise NotImplementedError(
108
109
            "AbstractToolParser.extract_tool_calls has not been implemented!"
        )
110
111
112
113
114
115
116
117
118

    def extract_tool_calls_streaming(
        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],
119
        request: ChatCompletionRequest,
120
    ) -> DeltaMessage | None:
121
122
123
124
125
126
127
128
        """
        Instance method that should be implemented for extracting tool calls
        from an incomplete response; for use when handling tool 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)
        """
        raise NotImplementedError(
129
130
            "AbstractToolParser.extract_tool_calls_streaming has not been implemented!"
        )
131
132
133


class ToolParserManager:
134
135
136
137
138
139
140
141
142
143
    """
    Central registry for ToolParser implementations.

    Supports two modes:
      - Eager (immediate) registration via `register_module`
      - Lazy registration via `register_lazy_module`
    """

    tool_parsers: dict[str, type[ToolParser]] = {}
    lazy_parsers: dict[str, tuple[str, str]] = {}  # name -> (module_path, class_name)
144
145

    @classmethod
146
    def get_tool_parser(cls, name: str) -> type[ToolParser]:
147
        """
148
        Retrieve a registered or lazily registered ToolParser class.
149

150
151
152
        If the parser is lazily registered,
        it will be imported and cached on first access.
        Raises KeyError if not found.
153
154
155
156
        """
        if name in cls.tool_parsers:
            return cls.tool_parsers[name]

157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
        if name in cls.lazy_parsers:
            return cls._load_lazy_parser(name)

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

    @classmethod
    def _load_lazy_parser(cls, name: str) -> type[ToolParser]:
        """Import and register a lazily loaded 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, ToolParser):
                raise TypeError(
                    f"{class_name} in {module_path} is not a ToolParser subclass."
                )
            cls.tool_parsers[name] = parser_cls  # cache
            return parser_cls
        except Exception as e:
            logger.exception(
                "Failed to import lazy tool parser '%s' from %s: %s",
                name,
                module_path,
                e,
            )
            raise
183
184

    @classmethod
185
186
    def _register_module(
        cls,
187
        module: type[ToolParser],
188
        module_name: str | list[str] | None = None,
189
190
        force: bool = True,
    ) -> None:
191
        """Register a ToolParser class immediately."""
192
193
        if not issubclass(module, ToolParser):
            raise TypeError(
194
                f"module must be subclass of ToolParser, but got {type(module)}"
195
            )
196

197
198
        if module_name is None:
            module_name = module.__name__
199

200
        if isinstance(module_name, str):
201
202
203
204
205
206
207
            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:
208
            if not force and name in cls.tool_parsers:
209
210
                existed = cls.tool_parsers[name]
                raise KeyError(f"{name} is already registered at {existed.__module__}")
211
212
            cls.tool_parsers[name] = module

213
214
215
216
217
218
219
220
    @classmethod
    def register_lazy_module(cls, name: str, module_path: str, class_name: str) -> None:
        """
        Register a lazy module mapping.

        Example:
            ToolParserManager.register_lazy_module(
                name="kimi_k2",
221
                module_path="vllm.tool_parsers.kimi_k2_parser",
222
223
224
225
226
                class_name="KimiK2ToolParser",
            )
        """
        cls.lazy_parsers[name] = (module_path, class_name)

227
228
    @classmethod
    def register_module(
229
        cls,
230
        name: str | list[str] | None = None,
231
        force: bool = True,
232
233
        module: type[ToolParser] | None = None,
    ) -> type[ToolParser] | Callable[[type[ToolParser]], type[ToolParser]]:
234
        """
235
236
237
238
239
240
241
242
243
        Register module immediately or lazily (as a decorator).

        Usage:
            @ToolParserManager.register_module("kimi_k2")
            class KimiK2ToolParser(ToolParser):
                ...

        Or:
            ToolParserManager.register_module(module=SomeToolParser)
244
245
        """
        if not isinstance(force, bool):
246
            raise TypeError(f"force must be a boolean, but got {type(force)}")
247

248
        # Immediate registration
249
250
251
252
        if module is not None:
            cls._register_module(module=module, module_name=name, force=force)
            return module

253
254
255
256
        # Decorator usage
        def _decorator(obj: type[ToolParser]) -> type[ToolParser]:
            module_path = obj.__module__
            class_name = obj.__name__
257

258
259
            if isinstance(name, str):
                names = [name]
260
            elif name is not None and is_list_of(name, str):
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
                names = name
            else:
                names = [class_name]

            for n in names:
                # Lazy mapping only: do not import now
                cls.lazy_parsers[n] = (module_path, class_name)

            return obj

        return _decorator

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

    @classmethod
    def import_tool_parser(cls, plugin_path: str) -> None:
280
        """Import a user-defined parser file from arbitrary path."""
281

282
        module_name = os.path.splitext(os.path.basename(plugin_path))[0]
283
284
285
        try:
            import_from_path(module_name, plugin_path)
        except Exception:
286
287
288
            logger.exception(
                "Failed to load module '%s' from %s.", module_name, plugin_path
            )