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

4
import json
5
from collections.abc import Generator
6
7
8
9
10

import partial_json_parser
import pytest
from partial_json_parser.core.options import Allow

11
from vllm.entrypoints.openai.protocol import DeltaMessage, FunctionCall, ToolCall
12
from vllm.entrypoints.openai.tool_parsers import JambaToolParser
13
from vllm.transformers_utils.detokenizer_utils import detokenize_incrementally
14
15
from vllm.transformers_utils.tokenizer import AnyTokenizer, get_tokenizer

16
17
pytestmark = pytest.mark.cpu_test

18
19
20
21
22
23
24
25
26
27
28
29
30
MODEL = "ai21labs/Jamba-tiny-dev"


@pytest.fixture(scope="module")
def jamba_tokenizer():
    return get_tokenizer(tokenizer_name=MODEL)


@pytest.fixture
def jamba_tool_parser(jamba_tokenizer):
    return JambaToolParser(jamba_tokenizer)


31
32
33
def assert_tool_calls(
    actual_tool_calls: list[ToolCall], expected_tool_calls: list[ToolCall]
):
34
35
    assert len(actual_tool_calls) == len(expected_tool_calls)

36
37
38
    for actual_tool_call, expected_tool_call in zip(
        actual_tool_calls, expected_tool_calls
    ):
39
40
41
42
43
44
45
46
        assert isinstance(actual_tool_call.id, str)
        assert len(actual_tool_call.id) > 16

        assert actual_tool_call.type == "function"
        assert actual_tool_call.function == expected_tool_call.function


def stream_delta_message_generator(
47
48
49
    jamba_tool_parser: JambaToolParser, jamba_tokenizer: AnyTokenizer, model_output: str
) -> Generator[DeltaMessage, None, None]:
    all_token_ids = jamba_tokenizer.encode(model_output, add_special_tokens=False)
50
51
52
53
54
55
56
57

    previous_text = ""
    previous_tokens = None
    prefix_offset = 0
    read_offset = 0
    for i, delta_token in enumerate(all_token_ids):
        delta_token_ids = [delta_token]
        previous_token_ids = all_token_ids[:i]
58
59
60
61
62
63
64
65
66
67
68
69
70
        current_token_ids = all_token_ids[: i + 1]

        (new_tokens, delta_text, new_prefix_offset, new_read_offset) = (
            detokenize_incrementally(
                tokenizer=jamba_tokenizer,
                all_input_ids=current_token_ids,
                prev_tokens=previous_tokens,
                prefix_offset=prefix_offset,
                read_offset=read_offset,
                skip_special_tokens=False,
                spaces_between_special_tokens=True,
            )
        )
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86

        current_text = previous_text + delta_text

        delta_message = jamba_tool_parser.extract_tool_calls_streaming(
            previous_text,
            current_text,
            delta_text,
            previous_token_ids,
            current_token_ids,
            delta_token_ids,
            request=None,  # type: ignore[arg-type]
        )
        if delta_message:
            yield delta_message

        previous_text = current_text
87
88
89
        previous_tokens = (
            previous_tokens + new_tokens if previous_tokens else new_tokens
        )
90
91
92
93
94
95
96
        prefix_offset = new_prefix_offset
        read_offset = new_read_offset


def test_extract_tool_calls_no_tools(jamba_tool_parser):
    model_output = "This is a test"
    extracted_tool_calls = jamba_tool_parser.extract_tool_calls(
97
98
        model_output, request=None
    )  # type: ignore[arg-type]
99
100
101
102
103
104
105
106
107
108
109
110
111
112
    assert not extracted_tool_calls.tools_called
    assert extracted_tool_calls.tool_calls == []
    assert extracted_tool_calls.content == model_output


@pytest.mark.parametrize(
    ids=[
        "single_tool",
        "single_tool_with_content",
        "parallel_tools",
    ],
    argnames=["model_output", "expected_tool_calls", "expected_content"],
    argvalues=[
        (
113
            """ <tool_calls>[\n    {"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}\n]</tool_calls>""",  # noqa: E501
114
            [
115
116
117
118
119
120
121
122
                ToolCall(
                    function=FunctionCall(
                        name="get_current_weather",
                        arguments=json.dumps(
                            {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}
                        ),
                    )
                )
123
            ],
124
125
            None,
        ),
126
        (
127
            """ Sure! let me call the tool for you.<tool_calls>[\n    {"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}\n]</tool_calls>""",  # noqa: E501
128
            [
129
130
131
132
133
134
135
136
                ToolCall(
                    function=FunctionCall(
                        name="get_current_weather",
                        arguments=json.dumps(
                            {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}
                        ),
                    )
                )
137
            ],
138
139
            " Sure! let me call the tool for you.",
        ),
140
        (
141
            """ <tool_calls>[\n    {"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}},\n    {"name": "get_current_weather", "arguments": {"city": "Orlando", "state": "FL", "unit": "fahrenheit"}}\n]</tool_calls>""",  # noqa: E501
142
            [
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
                ToolCall(
                    function=FunctionCall(
                        name="get_current_weather",
                        arguments=json.dumps(
                            {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}
                        ),
                    )
                ),
                ToolCall(
                    function=FunctionCall(
                        name="get_current_weather",
                        arguments=json.dumps(
                            {"city": "Orlando", "state": "FL", "unit": "fahrenheit"}
                        ),
                    )
                ),
159
            ],
160
161
            None,
        ),
162
163
    ],
)
164
165
166
def test_extract_tool_calls(
    jamba_tool_parser, model_output, expected_tool_calls, expected_content
):
167
    extracted_tool_calls = jamba_tool_parser.extract_tool_calls(
168
169
        model_output, request=None
    )  # type: ignore[arg-type]
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
    assert extracted_tool_calls.tools_called

    assert_tool_calls(extracted_tool_calls.tool_calls, expected_tool_calls)

    assert extracted_tool_calls.content == expected_content


@pytest.mark.parametrize(
    ids=[
        "no_tools",
        "single_tool",
        "single_tool_with_content",
        "parallel_tools",
    ],
    argnames=["model_output", "expected_tool_calls", "expected_content"],
    argvalues=[
186
        ("""This is a test""", [], """This is a test"""),
187
        (
188
            """ <tool_calls>[\n    {"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}\n]</tool_calls>""",  # noqa: E501
189
            [
190
191
192
193
194
195
196
197
                ToolCall(
                    function=FunctionCall(
                        name="get_current_weather",
                        arguments=json.dumps(
                            {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}
                        ),
                    )
                )
198
            ],
199
200
            " ",
        ),
201
        (
202
            """ Sure! let me call the tool for you.<tool_calls>[\n    {"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}\n]</tool_calls>""",  # noqa: E501
203
            [
204
205
206
207
208
209
210
211
                ToolCall(
                    function=FunctionCall(
                        name="get_current_weather",
                        arguments=json.dumps(
                            {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}
                        ),
                    )
                )
212
            ],
213
214
            " Sure! let me call the tool for you.",
        ),
215
        (
216
            """ <tool_calls>[\n    {"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}},\n    {"name": "get_current_weather", "arguments": {"city": "Orlando", "state": "FL", "unit": "fahrenheit"}}\n]</tool_calls>""",  # noqa: E501
217
            [
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
                ToolCall(
                    function=FunctionCall(
                        name="get_current_weather",
                        arguments=json.dumps(
                            {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}
                        ),
                    )
                ),
                ToolCall(
                    function=FunctionCall(
                        name="get_current_weather",
                        arguments=json.dumps(
                            {"city": "Orlando", "state": "FL", "unit": "fahrenheit"}
                        ),
                    )
                ),
234
            ],
235
236
            " ",
        ),
237
238
    ],
)
239
240
241
242
243
244
245
246
def test_extract_tool_calls_streaming(
    jamba_tool_parser,
    jamba_tokenizer,
    model_output,
    expected_tool_calls,
    expected_content,
):
    other_content: str = ""
247
248
    function_names: list[str] = []
    function_args_strs: list[str] = []
249
    tool_call_idx: int = -1
250
    tool_call_ids: list[str | None] = []
251
252

    for delta_message in stream_delta_message_generator(
253
254
        jamba_tool_parser, jamba_tokenizer, model_output
    ):
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
        # role should never be streamed from tool parser
        assert not delta_message.role

        if delta_message.content:
            other_content += delta_message.content

        streamed_tool_calls = delta_message.tool_calls

        if streamed_tool_calls and len(streamed_tool_calls) > 0:
            # make sure only one diff is present - correct even for parallel
            assert len(streamed_tool_calls) == 1
            tool_call = streamed_tool_calls[0]

            # if a new tool is being called, set up empty arguments
            if tool_call.index != tool_call_idx:
                tool_call_idx = tool_call.index
                function_args_strs.append("")
                tool_call_ids.append(None)

            # if a tool call ID is streamed, make sure one hasn't been already
            if tool_call.id and not tool_call_ids[tool_call.index]:
                tool_call_ids[tool_call.index] = tool_call.id

            # if parts of the function start being streamed
            if tool_call.function:
                # if the function name is defined, set it. it should be streamed
                # IN ENTIRETY, exactly one time.
                if tool_call.function.name:
                    assert isinstance(tool_call.function.name, str)
                    function_names.append(tool_call.function.name)

                if tool_call.function.arguments:
                    # make sure they're a string and then add them to the list
                    assert isinstance(tool_call.function.arguments, str)

290
                    function_args_strs[tool_call.index] += tool_call.function.arguments
291
292
293
294

    assert other_content == expected_content

    actual_tool_calls = [
295
296
297
298
299
300
301
302
303
        ToolCall(
            id=tool_call_id,
            function=FunctionCall(
                name=function_name,
                arguments=partial_json_parser.ensure_json(
                    function_args_str, Allow.OBJ | Allow.STR
                ),
            ),
        )
304
        for tool_call_id, function_name, function_args_str in zip(
305
306
            tool_call_ids, function_names, function_args_strs
        )
307
308
    ]
    assert_tool_calls(actual_tool_calls, expected_tool_calls)