Skip to content

fix(server): keep the reasoning between tool calls - #981

Open
NVolcz wants to merge 1 commit into
ROCm:mainfrom
NVolcz:keep-reasoning-between-tool-calls
Open

NVolcz wants to merge 1 commit into
ROCm:mainfrom
NVolcz:keep-reasoning-between-tool-calls

Conversation

@NVolcz

@NVolcz NVolcz commented Nov 7, 2026

Copy link
Copy Markdown

normalize_messages removes reasoning_content from every assistant message. When a model reasons, calls a tool and goes on from the result, its reasoning is gone from the prompt by the time it reads the result. The Gemma 4 docs say:

If a single model turn involves function or tool calls, thoughts must NOT be removed between the function calls.

Before, for a request that sends the reasoning back with the tool call:

<|turn>model
<|tool_call>call:find_person{name:<|"|>Alice<|"|>}<tool_call|><|tool_response>response:find_person{...}<tool_response|>

After:

<|turn>model
<|channel>thought
...the model's reasoning...
<channel|><|tool_call>call:find_person{name:<|"|>Alice<|"|>}<tool_call|><|tool_response>response:find_person{...}<tool_response|>

Reasoning is now kept for the assistant messages after the last user message and removed from earlier ones, as before. The Gemma 4 and Qwen3 templates render it only for those messages. Clients that send it as reasoning or thinking instead of reasoning_content have it moved to reasoning_content, which the templates read.

The prompt is longer now, and the request after a tool call prefills the reasoning again. The engine already holds the reasoning it generated, but the prompt cache reuses only saved prompt prefixes, and the template also writes a newline before <channel|> that the model did not generate. In a two-step tool chain with thinking on (gemma4-it:e4b, about 240 tokens of reasoning):

Request Before After
After the first tool call 176 tokens, 147 cached, prefill 0.81 s 412 tokens, 147 cached, prefill 0.97 s
After the second tool call 213 tokens, 147 cached, prefill 0.82 s 449 tokens, 412 cached, prefill 0.81 s

Each tool step pays once for its reasoning. The next request can then reuse the longer prompt, because it is now large enough to be saved as a prefix.

Once a new user message follows, the template drops the earlier turn's reasoning again, so the new prompt differs from the cached one from the first tool call on. With prefix slots (the default on NPU) the follow-up in the test above reused the same 147 tokens as before. Where prefix slots are off, on OGA backends and for Gemma 4 conversations with images or audio, the cache is cleared on any difference, so that request prefills the whole conversation again; this is not measured.

The docs now tell clients to send the reasoning back, and test_tool_calls.reasoning_kept_within_turn checks it with prompt_tokens.

How to reproduce

Send a tool call with its result, once without and once with reasoning_content on the tool call, and compare prompt_tokens:

for reasoning in '' ', "reasoning_content": "The user wants the weather in Paris, so I call get_weather."'; do
curl -s http://localhost:52625/v1/chat/completions -H 'Content-Type: application/json' -d '{
  "model": "gemma4-it:e4b", "max_tokens": 1,
  "tools": [{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object",
    "properties": {"city": {"type": "string"}}, "required": ["city"]}}}],
  "messages": [
    {"role": "user", "content": "What is the weather in Paris?"},
    {"role": "assistant", "content": null, "tool_calls": [{"id": "call_1", "type": "function",
      "function": {"name": "get_weather", "arguments": "{\"city\": \"Paris\"}"}}]'"$reasoning"'},
    {"role": "tool", "tool_call_id": "call_1", "content": "{\"temperature_c\": 23}"}
  ]}' | jq '.usage.prompt_tokens'
done

On main both requests report the same count, because the reasoning is removed. On this branch the second one is larger by the reasoning's tokens.

I found this while comparing the prompts flm renders with Google's template, using a small debug branch that logs every prompt to a file (log-prompt, FLM_LOG_PROMPT=<file>). It also shows the rendered prompt directly, if you want to see the reasoning in it.

normalize_messages removed reasoning_content from every assistant message, so when a model
reasoned and then called a tool, its reasoning was gone by the time it read the tool's result.
The Gemma 4 docs say the reasoning must not be removed between the tool calls of one model
turn, and the Gemma 4 and Qwen3 templates render it for the assistant messages after the last
user message.

normalize_messages now takes the first string among reasoning_content, reasoning and thinking,
removes all three, and puts it back as reasoning_content for the assistant messages after the
last user message, or for all of them when there is none. The docs tell clients to send the
reasoning back, and test_tool_calls.reasoning_kept_within_turn checks it with prompt_tokens.
@tawei-amd
tawei-amd requested a review from zaneni6 November 7, 2026 21:54

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant