diff --git a/scripts/response_regeneration/script.py b/scripts/response_regeneration/script.py index 94a4a61ad..ad351c0a1 100644 --- a/scripts/response_regeneration/script.py +++ b/scripts/response_regeneration/script.py @@ -472,6 +472,8 @@ def _sample_from_response( message = choice["message"] content = message.get("content") tool_calls = message.get("tool_calls") + # Older vLLM builds returned this as `reasoning_content`. + reasoning = message.get("reasoning") or message.get("reasoning_content") # A tool call legitimately has empty content; only a wholly empty generation # corrupts the next prefix and must fail the conversation. @@ -487,8 +489,7 @@ def _sample_from_response( input_ids, loss_mask = build_boundary_sample(prompt_token_ids, completion_token_ids) if tool_calls: - # History keeps the parsed call; any generated is supervised in - # this row's completion tokens, not re-rendered. + # History keeps the parsed call, not the raw text it was parsed from. assistant_msg = { "role": "assistant", "content": content or "", @@ -496,6 +497,10 @@ def _sample_from_response( } else: assistant_msg = {"role": "assistant", "content": content} + if reasoning: + # Templates re-render a turn's reasoning from this field. Kimi K2.5 has no + # -in-content fallback, so dropping it changes the next prompt. + assistant_msg["reasoning_content"] = reasoning sample = { "id": f"{conv_id}_gen{sample_index}", diff --git a/tests/unit/scripts/test_response_regeneration.py b/tests/unit/scripts/test_response_regeneration.py index 41a003264..2857e7738 100644 --- a/tests/unit/scripts/test_response_regeneration.py +++ b/tests/unit/scripts/test_response_regeneration.py @@ -566,13 +566,23 @@ def _tool_call(call_id="call_1", name="get_weather", arguments='{"city": "Tokyo" def _response( - *, prompt_token_ids, token_ids, content=None, tool_calls=None, finish="stop" + *, + prompt_token_ids, + token_ids, + content=None, + tool_calls=None, + finish="stop", + reasoning=None, ): """A vLLM chat-completion response with ``return_token_ids`` populated.""" return { "choices": [ { - "message": {"content": content, "tool_calls": tool_calls}, + "message": { + "content": content, + "tool_calls": tool_calls, + "reasoning": reasoning, + }, "finish_reason": finish, "token_ids": token_ids, } @@ -729,6 +739,31 @@ def test_sample_from_response_rejects_empty_and_missing_token_ids(): ) +def test_reasoning_is_carried_into_the_tool_call_history(): + # Kimi K2.5 preserves reasoning for exactly the in-flight tool loop and reads + # it only from this field, so dropping it silently changes the next prompt. + _, assistant_msg, _ = regen._sample_from_response( + _response( + prompt_token_ids=[1, 2], + token_ids=[3], + tool_calls=[_tool_call()], + reasoning="Need the weather tool.", + ), + detokenize=_detok, + conv_id="c", + sample_index=0, + idx=0, + endpoint="ep", + sampling_params={}, + ) + assert assistant_msg == { + "role": "assistant", + "content": "", + "tool_calls": [_tool_call()], + "reasoning_content": "Need the weather tool.", + } + + # --- the tool-call loop: splice, truncate, and the unchanged plain path ---