Skip to content

Add continue_final_message to chat completions#431

Open
ethanj801 wants to merge 2 commits into
theroyallab:mainfrom
ethanj801:continue-final-message-pr
Open

Add continue_final_message to chat completions#431
ethanj801 wants to merge 2 commits into
theroyallab:mainfrom
ethanj801:continue-final-message-pr

Conversation

@ethanj801

Copy link
Copy Markdown
Contributor

Is your pull request related to a problem? Please describe.
Yes. There is currently no way to continue a partial assistant message through the chat template. The only mechanism for steering the start of a response is response_prefix, which appends raw text to the rendered prompt after the generation prompt. Because that text never passes through the template, everything the template does to assistant content is skipped. That includes the spacing it places before a reply, any trimming it applies to the text, and the reasoning tags some templates insert around a turn (e.g. the reasoning tags in DeepSeek-R1).

For example the Mistral large 2411 prompt renders an assistant turn as " " + content + eos_token. The space belongs to the template, and the model emits it inside its first token. Suppose the model answered The story begins and a client wants to continue that reply. The prompt the model was trained on looks like this:

[INST] Tell me a story.[/INST] The story begins

With response_prefix the client has to reproduce the template's spacing itself. If it sends the text as stored it produces [/INST]The story begins with the space missing. If it sends the streamed text verbatim, which starts with the space the model emitted, the prompt becomes [/INST] The story begins with a doubled space, which tokenizes into a sequence the model never saw in training. Getting this right requires the client to know how each model's template spaces its turns.

HF transformers and vLLM both solve this with a continue_final_message request option. When set, the final message of the conversation is rendered through the template as a normal turn, and the prompt is then cut back so the turn is left unterminated, with no EOS token and no trailing formatting. Generation continues the message in place. This PR adds the same option to /v1/chat/completions, matching their api shape.

The implementation appends a sentinel string to the final message content before rendering, renders normally, and cuts the prompt at the sentinel. The sentinel mechanism and cut logic are ported from render_jinja_template in huggingface/transformers (src/transformers/utils/chat_template_utils.py, Apache License 2.0), and the port is credited in a code comment at the sentinel constant. The sentinel carries a trailing space so the cut can detect templates that trim trailing whitespace from message content (for example content | trim) and trim the continued text the same way. If the template rewrites or drops the final message so the sentinel does not survive rendering, the request fails with a 422 rather than guessing a cut point. add_generation_prompt must be false when continue_final_message is set, matching vLLM, otherwise the request fails with a 422.

Because HF and vLLM do not have our existing response_prefix parameter, I had to define how response_prefix interacts with continue_final_message. I chose to allow the use of both simultaneously. When both are used response_prefix is appended verbatim after the continued turn. Existing use of response_prefix is unchanged.

This PR has an intentional difference from the transformers implementation.

In this implementation the sentinel is appended as its own final content part rather than into the last text part. format_messages_with_template flattens parts into one string in list order, so the rendered prompt is unchanged for text content, and an image part that follows the last text part stays ahead of the cut instead of being deleted with it. The transformers implementation appends into the last text block and loses a trailing image on continuation.

Finally, there is a known limitation which is that tool calls on the final message are lost, since templates render them after the content and the cut removes them. This limitation is also present in the transformers and vLLM implementations.

Why should this feature be added?
It brings the chat endpoint to parity with HF transformers and vLLM for continuing assistant messages, and it removes the need for clients to carry per-model template knowledge to build correct continuations. The change is contained to the chat completion request type and apply_chat_template. Requests that do not set the option render the same prompts as before.

Examples
Request against mistral large 2411 :

POST /v1/chat/completions
{
  "messages": [
    {"role": "user", "content": "Tell me a story."},
    {"role": "assistant", "content": "The story begins"}
  ],
  "continue_final_message": true,
  "add_generation_prompt": false,
  "max_tokens": 64
}

The rendered prompt ends with:

[INST] Tell me a story.[/INST] The story begins

The template's own space after [/INST] is present, there is no EOS token, and the model continues the sentence.

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