fix(chat_template): emit tool_response multimodal placeholders inside the block

#140

Problem

Multimodal placeholders in tool messages are emitted after format_tool_response_block() has already closed the block with <tool_response|>:

{{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}   {# text goes inside #}
{%- for part in tool_body -%}
    {%- if part.get('type') in ['image', 'image_url'] -%}
        {{- '<|image|>' -}}                                    {# image lands outside #}

When the tool message is the final message, the rendered prompt therefore ends with a bare multimodal token:

... <|"|>}<tool_response|><|image|>

In the tested setup, generation from this prompt shape is unstable. The observed failure rate depends on the prompt and image: the minimal reproduction below fails in 58/60 runs with a dark checkerboard and 25/60 with a light checkerboard.

Same issue as google/gemma-4-26B-A4B-it#55 — the chat_template.jinja in this repository is byte-identical, and the failure reproduces here at a higher rate.

Fix

Move content-type dispatch into format_tool_response_block() and emit multimodal placeholders before the closing tag:

before   <|tool_response>response:search{value:<|"|>text<|"|>}<tool_response|><|image|>
after    <|tool_response>response:search{value:<|"|>text<|"|>}<|image|><tool_response|>

This preserves the template's current text-first ordering. The same tokens are emitted; only the position of <|image|> relative to <tool_response|> changes.

Reproduction

The reproduction script (repro_gemma4_tool_image.py, attached to google/gemma-4-26B-A4B-it#55) pins a tool call with tool_choice, so the only correct output is that call with no prose. Results on google/gemma-4-31B-it, vLLM 0.26.0, temperature 1.0, n=60 per cell:

Condition Current template This PR
tool-message image, dark board 58/60 (96.7%) 0/60
tool-message image, light board 25/60 (41.7%) 0/60
no tool-message image (control) 0/20

Because the prompt ends on a bare <|image|> outside any block, the model loses track of block state and fails to open the next block. Raw token dumps (logprobs) show the dominant failure: generation starts directly with call: — no opening <|tool_call> — and closes a block it never opened:

call:pick_category{...}<tool_call|><|tool_response>

This is not a parser issue: the opening tag is absent from the sampled tokens, the failure reproduces without streaming, and the same parser extracts every call once the template is fixed. Repetition loops as in the 26B report (24/60 light, 49/60 dark) also occur, less often.

vllm serve google/gemma-4-31B-it \
    --enable-auto-tool-choice --tool-call-parser gemma4 --reasoning-parser gemma4 \
    --trust-request-chat-template

python repro_gemma4_tool_image.py --base http://localhost:8000/v1 \
    --model google/gemma-4-31B-it -n 60 [--invert] [--chat-template patched.jinja]

Also fixes a crash

In Jinja a dict is also is sequence. The current call site tests tool_body is string and then is sequence, so a tool message whose content is a dict falls into the content-parts branch, where part.get('type') raises UndefinedError.

tool content current this PR
string same same
mapping render error renders correctly
parts (text + image) image outside the block image inside the block
parts (text only) same same
Ready to merge
This branch is ready to get merged automatically.

Sign up or log in to comment