@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 144 lines 5.55 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5 import json
6
7 import pytest
8
9 from tests.tool_parsers.common_tests import (
10 ToolParserTestConfig,
11 ToolParserTests,
12 )
13 from tests.tool_parsers.utils import (
14 run_tool_extraction,
15 run_tool_extraction_streaming,
16 split_string_into_token_deltas,
17 )
18 from vllm.tokenizers import get_tokenizer
19 from vllm.tool_parsers.granite_engine_tool_parser import GraniteEngineToolParser
20
21
22 class TestGraniteToolParser(ToolParserTests):
23 @pytest.fixture
24 def test_config(self) -> ToolParserTestConfig:
25 return ToolParserTestConfig(
26 parser_name="granite",
27 # Test data
28 no_tool_calls_output="This is a regular response without any tool calls.",
29 single_tool_call_output=(
30 '<|tool_call|> [{"name": "get_weather", '
31 '"arguments": {"city": "Tokyo"}}]'
32 ),
33 parallel_tool_calls_output="""<|tool_call|> [
34 {"name": "get_weather", "arguments": {"city": "Tokyo"}},
35 {"name": "get_time", "arguments": {"timezone": "Asia/Tokyo"}}
36 ]""",
37 various_data_types_output="""<tool_call> [{
38 "name": "test_function",
39 "arguments": {
40 "string_field": "hello",
41 "int_field": 42,
42 "float_field": 3.14,
43 "bool_field": true,
44 "null_field": null,
45 "array_field": ["a", "b", "c"],
46 "object_field": {"nested": "value"},
47 "empty_array": [],
48 "empty_object": {}
49 }
50 }]""",
51 empty_arguments_output=(
52 '<|tool_call|> [{"name": "refresh", "arguments": {}}]'
53 ),
54 surrounding_text_output="""Let me check the weather for you.
55 <|tool_call|> [{"name": "get_weather", "arguments": {"city": "Tokyo"}}]
56 I'll get that information.""",
57 escaped_strings_output="""<tool_call> [{
58 "name": "test_function",
59 "arguments": {
60 "quoted": "He said \\"hello\\"",
61 "path": "C:\\\\Users\\\\file.txt",
62 "newline": "line1\\nline2",
63 "unicode": "emoji: 🎉"
64 }
65 }]""",
66 malformed_input_outputs=[
67 '<|tool_call|> [{"name": "func", "arguments": {',
68 '<|tool_call|> {"name": "func", "arguments": {}}', # Not an array
69 '[{"name": "func", "arguments": "not a dict"}]',
70 'Some text [{"name": "func"}]', # JSON but not tool call format
71 ],
72 # Expected results
73 single_tool_call_expected_name="get_weather",
74 single_tool_call_expected_args={"city": "Tokyo"},
75 # Granite strips content when tool calls present
76 single_tool_call_expected_content=None,
77 parallel_tool_calls_count=2,
78 parallel_tool_calls_names=["get_weather", "get_time"],
79 )
80
81 # Granite-Specific Tests
82
83 @pytest.mark.parametrize("streaming", [True, False])
84 def test_granite_token_prefix_format(self, tool_parser, streaming):
85 """Verify parser handles Granite 3.0 <|tool_call|> token format."""
86 single_tool_call_token = (
87 '<|tool_call|> [{"name": "get_weather", "arguments": {"city": "Tokyo"}}]'
88 )
89 content, tool_calls = run_tool_extraction(
90 tool_parser, single_tool_call_token, streaming=streaming
91 )
92 assert len(tool_calls) == 1, (
93 f"Expected 1 tool call from token format, got {len(tool_calls)}"
94 )
95 assert tool_calls[0].function.name == "get_weather"
96
97 @pytest.mark.parametrize("streaming", [True, False])
98 def test_granite_string_prefix_format(self, tool_parser, streaming):
99 """Verify parser handles Granite 3.1 <tool_call> string format."""
100 single_tool_call_string = (
101 '<tool_call> [{"name": "get_weather", "arguments": {"city": "Tokyo"}}]'
102 )
103 content, tool_calls = run_tool_extraction(
104 tool_parser, single_tool_call_string, streaming=streaming
105 )
106 assert len(tool_calls) == 1, (
107 f"Expected 1 tool call from string format, got {len(tool_calls)}"
108 )
109 assert tool_calls[0].function.name == "get_weather"
110
111
112 # granite emits arguments before name and its own tokenizer (not gpt2) is used
113 # here so the token boundaries match production; get_tokenizer only fetches the
114 # small tokenizer files, not the model weights.
115 @pytest.fixture(scope="module")
116 def granite_tokenizer():
117 return get_tokenizer(tokenizer_name="ibm-granite/granite-3.1-8b-instruct")
118
119
120 @pytest.mark.parametrize("chunk_size", [2, 3, 4, 5])
121 def test_streaming_parallel_calls_batched_deltas(granite_tokenizer, chunk_size):
122 """A batched delta (multiple tokens) spanning the boundary between two
123 parallel calls must not drop the first call's name. granite streams
124 arguments before name, so the name only completes as the next call appears.
125 """
126 parser = GraniteEngineToolParser(granite_tokenizer)
127 model_output = (
128 '<|tool_call|> [{"arguments": {"city": "Tokyo"}, "name": "get_weather"}, '
129 '{"arguments": {"timezone": "Asia/Tokyo"}, "name": "get_time"}]'
130 )
131 token_deltas = split_string_into_token_deltas(granite_tokenizer, model_output)
132 batched = [
133 "".join(token_deltas[i : i + chunk_size])
134 for i in range(0, len(token_deltas), chunk_size)
135 ]
136 reconstructor = run_tool_extraction_streaming(
137 parser, batched, assert_one_tool_per_delta=False
138 )
139 names = [tc.function.name for tc in reconstructor.tool_calls]
140 assert names == ["get_weather", "get_time"]
141 # trailing args of the final call are flushed by the serving layer
142 assert json.loads(reconstructor.tool_calls[0].function.arguments) == {
143 "city": "Tokyo"
144 }