@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 164 lines 6.13 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import json
5 from unittest.mock import MagicMock
6
7 import pytest
8
9 from tests.tool_parsers.common_tests import (
10 ToolParserTestConfig,
11 ToolParserTests,
12 )
13 from vllm.tokenizers import TokenizerLike
14 from vllm.tool_parsers.internlm2_tool_parser import Internlm2ToolParser
15
16
17 class TestInternLM2ToolParser(ToolParserTests):
18 @pytest.fixture
19 def tokenizer(self, default_tokenizer: TokenizerLike) -> TokenizerLike:
20 """Add some internlm2 specific tokens to the default vocab."""
21 tokenizer_vocab = default_tokenizer.get_vocab()
22 default_tokenizer.get_vocab = MagicMock()
23 tokenizer_vocab.update(
24 {
25 "<|action_start|>": 92540,
26 "<|plugin|>": 92541,
27 "<|action_end|>": 92542,
28 }
29 )
30 default_tokenizer.get_vocab.return_value = tokenizer_vocab
31 return default_tokenizer
32
33 @pytest.fixture
34 def test_config(self) -> ToolParserTestConfig:
35 return ToolParserTestConfig(
36 parser_name="internlm",
37 # Test data
38 no_tool_calls_output="This is a regular response without any tool calls.",
39 single_tool_call_output=(
40 '<|action_start|><|plugin|>{"name": "get_weather", '
41 '"parameters": {"city": "Tokyo"}}<|action_end|>'
42 ),
43 # InternLM2 doesn't support parallel calls
44 parallel_tool_calls_output=(
45 '<|action_start|><|plugin|>{"name": "get_weather", '
46 '"parameters": {"city": "Tokyo"}}<|action_end|>'
47 ),
48 various_data_types_output="""<|action_start|><|plugin|>{
49 "name": "test_function",
50 "parameters": {
51 "string_field": "hello",
52 "int_field": 42,
53 "float_field": 3.14,
54 "bool_field": true,
55 "null_field": null,
56 "array_field": ["a", "b", "c"],
57 "object_field": {"nested": "value"},
58 "empty_array": [],
59 "empty_object": {}
60 }
61 }<|action_end|>""",
62 empty_arguments_output=(
63 '<|action_start|><|plugin|>{"name": "refresh", '
64 '"parameters": {}}<|action_end|>'
65 ),
66 surrounding_text_output=(
67 "Let me check the weather for you. "
68 '<|action_start|><|plugin|>{"name": "get_weather", '
69 '"parameters": {"city": "Tokyo"}}<|action_end|>'
70 ),
71 escaped_strings_output="""<|action_start|><|plugin|>{
72 "name": "test_function",
73 "parameters": {
74 "quoted": "He said \\"hello\\"",
75 "path": "C:\\\\Users\\\\file.txt",
76 "newline": "line1\\nline2",
77 "unicode": "emoji: 🎉"
78 }
79 }<|action_end|>""",
80 malformed_input_outputs=[
81 '<|action_start|><|plugin|>{"name": "func", "parameters": {',
82 (
83 '<|action_start|><|plugin|>{"name": "func", '
84 '"parameters": "not a dict"}<|action_end|>'
85 ),
86 "<|action_start|><|plugin|>not json<|action_end|>",
87 "<|action_start|><|plugin|>",
88 '<|action_start|>{"name": "func"}',
89 ],
90 # Expected results
91 single_tool_call_expected_name="get_weather",
92 single_tool_call_expected_args={"city": "Tokyo"},
93 single_tool_call_expected_content=None,
94 parallel_tool_calls_count=1, # InternLM2 only supports single tool calls
95 parallel_tool_calls_names=["get_weather"],
96 # Parser-specific settings
97 allow_empty_or_json_empty_args=True,
98 # xfail markers
99 xfail_streaming={
100 "test_single_tool_call_simple_args": (
101 "InternLM2 streaming not fully implemented"
102 ),
103 "test_parallel_tool_calls": (
104 "InternLM2 streaming not fully implemented"
105 ),
106 "test_various_data_types": (
107 "InternLM2 streaming not fully implemented"
108 ),
109 "test_empty_arguments": ("InternLM2 streaming not fully implemented"),
110 "test_surrounding_text": ("InternLM2 streaming not fully implemented"),
111 "test_escaped_strings": ("InternLM2 streaming not fully implemented"),
112 "test_streaming_reconstruction": (
113 "InternLM2 streaming parser returns '<|action_start|' as "
114 "content instead of None - streaming/non-streaming inconsistency"
115 ),
116 },
117 xfail_nonstreaming={
118 "test_malformed_input": (
119 "InternLM2 parser raises JSONDecodeError on malformed JSON "
120 "instead of gracefully handling it"
121 ),
122 },
123 )
124
125
126 def test_streaming_arguments_in_single_delta(default_tokenizer: TokenizerLike) -> None:
127 """Arguments arriving whole in one delta must not be dropped."""
128 tokenizer_vocab = default_tokenizer.get_vocab()
129 default_tokenizer.get_vocab = MagicMock()
130 tokenizer_vocab.update(
131 {
132 "<|action_start|>": 92540,
133 "<|plugin|>": 92541,
134 "<|action_end|>": 92542,
135 }
136 )
137 default_tokenizer.get_vocab.return_value = tokenizer_vocab
138 parser = Internlm2ToolParser(default_tokenizer)
139
140 deltas = [
141 '<|action_start|><|plugin|>{"name": "get_weather"',
142 ', "parameters": {"city": "Dallas", "state": "TX"}}<|action_end|>',
143 ]
144
145 streamed = ""
146 current_text = ""
147 for delta_text in deltas:
148 previous_text = current_text
149 current_text += delta_text
150 delta_message = parser.extract_tool_calls_streaming(
151 previous_text=previous_text,
152 current_text=current_text,
153 delta_text=delta_text,
154 previous_token_ids=[],
155 current_token_ids=[],
156 delta_token_ids=[],
157 request=None,
158 )
159 if delta_message and delta_message.tool_calls:
160 arguments = delta_message.tool_calls[0].function.arguments
161 if arguments:
162 streamed += arguments
163
164 assert json.loads(streamed) == {"city": "Dallas", "state": "TX"}