main
py 414 lines 14.6 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 """Tests for Gemma4 chat template rendering."""
5
6 from pathlib import Path
7
8 import jinja2.sandbox
9 import pytest
10
11 TEMPLATE_PATH = (
12 Path(__file__).resolve().parent.parent.parent
13 / "examples"
14 / "tool_chat_template_gemma4.jinja"
15 )
16
17
18 @pytest.fixture(scope="module")
19 def gemma4_template():
20 """Load and compile the Gemma4 chat template."""
21 template_str = TEMPLATE_PATH.read_text()
22 env = jinja2.sandbox.ImmutableSandboxedEnvironment()
23 return env.from_string(template_str)
24
25
26 def _render(template, messages, **kwargs):
27 """Render the template with sensible defaults."""
28 kwargs.setdefault("bos_token", "<bos>")
29 kwargs.setdefault("add_generation_prompt", False)
30 return template.render(messages=messages, **kwargs)
31
32
33 class TestGemma4ChatTemplate:
34 def test_basic_multiturn_thinking_disabled(self, gemma4_template):
35 """With enable_thinking=False (default), generation prompt ends with
36 an empty thought channel to suppress thinking."""
37 messages = [
38 {"role": "user", "content": "Hello"},
39 {"role": "assistant", "content": "Hi there!"},
40 {"role": "user", "content": "How are you?"},
41 ]
42 result = _render(gemma4_template, messages, add_generation_prompt=True)
43 assert "<|turn>user\n" in result
44 assert "<|turn>model\n" in result
45 assert "Hello" in result
46 assert "Hi there!" in result
47 assert "How are you?" in result
48 assert result.rstrip("\n").endswith("<|channel>thought\n<channel|>")
49
50 def test_basic_multiturn_thinking_enabled(self, gemma4_template):
51 """With enable_thinking=True, generation prompt ends with model
52 turn opener (no thought suppression)."""
53 messages = [
54 {"role": "user", "content": "Hello"},
55 {"role": "assistant", "content": "Hi there!"},
56 {"role": "user", "content": "How are you?"},
57 ]
58 result = _render(
59 gemma4_template,
60 messages,
61 add_generation_prompt=True,
62 enable_thinking=True,
63 )
64 assert "<|turn>user\n" in result
65 assert "<|turn>model\n" in result
66 assert "Hello" in result
67 assert "Hi there!" in result
68 assert "How are you?" in result
69 assert result.rstrip("\n").endswith("<|turn>model")
70
71 def test_system_message(self, gemma4_template):
72 messages = [
73 {"role": "system", "content": "You are helpful."},
74 {"role": "user", "content": "Hi"},
75 ]
76 result = _render(gemma4_template, messages)
77 assert "<|turn>system\n" in result
78 assert "You are helpful." in result
79
80 def test_thinking_enabled(self, gemma4_template):
81 messages = [{"role": "user", "content": "Think about this"}]
82 result = _render(
83 gemma4_template,
84 messages,
85 add_generation_prompt=True,
86 enable_thinking=True,
87 )
88 assert "<|think|>" in result
89 assert "<|turn>system\n" in result
90
91 def test_tool_declarations(self, gemma4_template):
92 tools = [
93 {
94 "type": "function",
95 "function": {
96 "name": "get_weather",
97 "description": "Get weather for a city",
98 "parameters": {
99 "type": "object",
100 "properties": {
101 "city": {
102 "type": "string",
103 "description": "City name",
104 }
105 },
106 "required": ["city"],
107 },
108 },
109 }
110 ]
111 messages = [{"role": "user", "content": "What is the weather?"}]
112 result = _render(
113 gemma4_template,
114 messages,
115 tools=tools,
116 add_generation_prompt=True,
117 )
118 assert "<|tool>" in result
119 assert "declaration:get_weather" in result
120 assert "<tool|>" in result
121 assert '<|"|>City name<|"|>' in result
122
123 def test_tool_calls_in_assistant(self, gemma4_template):
124 messages = [
125 {"role": "user", "content": "Weather in London?"},
126 {
127 "role": "assistant",
128 "content": "",
129 "tool_calls": [
130 {
131 "id": "call_1",
132 "function": {
133 "name": "get_weather",
134 "arguments": {"city": "London"},
135 },
136 }
137 ],
138 },
139 ]
140 result = _render(gemma4_template, messages)
141 assert "<|tool_call>call:get_weather{" in result
142 assert "}<tool_call|>" in result
143 assert '<|"|>London<|"|>' in result
144
145 def test_tool_responses_openai_style(self, gemma4_template):
146 """role='tool' messages are formatted as <|tool_response> blocks
147 with content dumped as-is."""
148 messages = [
149 {"role": "user", "content": "Weather?"},
150 {
151 "role": "assistant",
152 "content": "",
153 "tool_calls": [
154 {
155 "id": "call_1",
156 "function": {
157 "name": "get_weather",
158 "arguments": {"city": "London"},
159 },
160 }
161 ],
162 },
163 {
164 "role": "tool",
165 "tool_call_id": "call_1",
166 "content": '{"temperature": 15, "condition": "sunny"}',
167 },
168 ]
169 result = _render(gemma4_template, messages, add_generation_prompt=True)
170 assert "<|tool_response>" in result
171 assert "response:get_weather{" in result
172 assert "<tool_response|>" in result
173 assert '"temperature": 15' in result
174
175 def test_tool_responses_legacy_style(self, gemma4_template):
176 """tool_responses embedded on the assistant message."""
177 messages = [
178 {"role": "user", "content": "Weather?"},
179 {
180 "role": "assistant",
181 "content": "",
182 "tool_calls": [
183 {
184 "function": {
185 "name": "get_weather",
186 "arguments": {"city": "London"},
187 },
188 }
189 ],
190 "tool_responses": [
191 {
192 "name": "get_weather",
193 "response": {"temperature": 20},
194 }
195 ],
196 },
197 ]
198 result = _render(gemma4_template, messages)
199 assert "<|tool_response>" in result
200 assert "response:get_weather{" in result
201 assert "temperature:" in result
202
203 def test_generation_prompt_not_after_tool_response(self, gemma4_template):
204 """add_generation_prompt=True should NOT add <|turn>model when the
205 last message type was tool_response (the model turn continues)."""
206 messages = [
207 {"role": "user", "content": "Weather?"},
208 {
209 "role": "assistant",
210 "content": "",
211 "tool_calls": [
212 {
213 "id": "call_1",
214 "function": {
215 "name": "get_weather",
216 "arguments": {"city": "London"},
217 },
218 }
219 ],
220 },
221 {
222 "role": "tool",
223 "tool_call_id": "call_1",
224 "content": "sunny",
225 },
226 ]
227 result = _render(gemma4_template, messages, add_generation_prompt=True)
228 assert not result.strip().endswith("<|turn>model\n")
229
230 def test_reasoning_in_tool_chains(self, gemma4_template):
231 """Reasoning field on assistant with tool_calls after last user
232 message emits <|channel>thought\\n...<channel|>."""
233 messages = [
234 {"role": "user", "content": "Calculate something"},
235 {
236 "role": "assistant",
237 "content": "",
238 "reasoning": "Let me think about this...",
239 "tool_calls": [
240 {
241 "function": {
242 "name": "calculator",
243 "arguments": {"expr": "2+2"},
244 },
245 }
246 ],
247 },
248 ]
249 result = _render(gemma4_template, messages)
250 assert "<|channel>thought\n" in result
251 assert "Let me think about this..." in result
252 assert "<channel|>" in result
253
254 def test_reasoning_not_before_last_user(self, gemma4_template):
255 """Reasoning on assistant BEFORE the last user message is dropped."""
256 messages = [
257 {"role": "user", "content": "First"},
258 {
259 "role": "assistant",
260 "content": "Response",
261 "reasoning": "Old reasoning that should be dropped",
262 "tool_calls": [
263 {
264 "function": {
265 "name": "fn",
266 "arguments": {},
267 },
268 }
269 ],
270 },
271 {"role": "user", "content": "Second"},
272 ]
273 result = _render(gemma4_template, messages, add_generation_prompt=True)
274 assert "Old reasoning" not in result
275
276 def test_strip_thinking_in_model_content(self, gemma4_template):
277 """<|channel>...<channel|> in model content is stripped by the
278 strip_thinking macro."""
279 messages = [
280 {"role": "user", "content": "Hi"},
281 {
282 "role": "assistant",
283 "content": ("<|channel>internal thought<channel|>Visible answer"),
284 },
285 ]
286 result = _render(gemma4_template, messages)
287 assert "internal thought" not in result
288 assert "Visible answer" in result
289
290 def test_multi_turn_tool_chain(self, gemma4_template):
291 """assistant->tool->assistant->tool produces exactly one
292 <|turn>model (later assistants continue the same turn)."""
293 messages = [
294 {"role": "user", "content": "Do two things"},
295 {
296 "role": "assistant",
297 "content": "",
298 "tool_calls": [
299 {
300 "id": "c1",
301 "function": {"name": "step1", "arguments": {}},
302 },
303 ],
304 },
305 {"role": "tool", "tool_call_id": "c1", "content": "result1"},
306 {
307 "role": "assistant",
308 "content": "",
309 "tool_calls": [
310 {
311 "id": "c2",
312 "function": {"name": "step2", "arguments": {}},
313 },
314 ],
315 },
316 {"role": "tool", "tool_call_id": "c2", "content": "result2"},
317 ]
318 result = _render(gemma4_template, messages, add_generation_prompt=True)
319 assert result.count("<|turn>model\n") == 1
320
321 def test_format_argument_types(self, gemma4_template):
322 """Strings wrapped in <|"|>, booleans as true/false, numbers bare."""
323 messages = [
324 {"role": "user", "content": "Test"},
325 {
326 "role": "assistant",
327 "content": "",
328 "tool_calls": [
329 {
330 "function": {
331 "name": "test_fn",
332 "arguments": {
333 "name": "Alice",
334 "active": True,
335 "count": 42,
336 },
337 },
338 }
339 ],
340 },
341 ]
342 result = _render(gemma4_template, messages)
343 assert '<|"|>Alice<|"|>' in result
344 assert "active:true" in result
345 assert "count:42" in result
346
347 def test_tool_response_with_multimodal_content(self, gemma4_template):
348 """Multimodal placeholders in tool messages are emitted after the
349 tool_response block."""
350 messages = [
351 {"role": "user", "content": "Download the image and describe it."},
352 {
353 "role": "assistant",
354 "content": "",
355 "tool_calls": [
356 {
357 "id": "call_1",
358 "type": "function",
359 "function": {
360 "name": "download_image",
361 "arguments": {"url": "https://example.com/x.png"},
362 },
363 },
364 ],
365 },
366 {
367 "role": "tool",
368 "tool_call_id": "call_1",
369 "content": [
370 {"type": "text", "text": "Image downloaded successfully."},
371 {"type": "image"},
372 ],
373 },
374 ]
375 result = _render(gemma4_template, messages, add_generation_prompt=True)
376 assert "<|tool_response>" in result
377 assert "response:download_image{" in result
378 assert "<tool_response|>" in result
379 assert "<|image|>" in result
380
381 def test_tool_response_with_all_modalities(self, gemma4_template):
382 """All multimodal types (image, audio, video) in a single tool
383 response are rendered."""
384 messages = [
385 {"role": "user", "content": "Process media"},
386 {
387 "role": "assistant",
388 "content": "",
389 "tool_calls": [
390 {
391 "id": "c1",
392 "type": "function",
393 "function": {
394 "name": "process",
395 "arguments": {},
396 },
397 },
398 ],
399 },
400 {
401 "role": "tool",
402 "tool_call_id": "c1",
403 "content": [
404 {"type": "text", "text": "Results."},
405 {"type": "image"},
406 {"type": "audio"},
407 {"type": "video"},
408 ],
409 },
410 ]
411 result = _render(gemma4_template, messages, add_generation_prompt=True)
412 assert "<|image|>" in result
413 assert "<|audio|>" in result
414 assert "<|video|>" in result