[Bugfix][Frontend] Only echo the assistant turn in batched chat completions (#52529)
Signed-off-by: Kaif <kaifkohari10@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Kaif Kohari committed
Aug 29, 2026 at 17:44 UTC
7fbfca2670bbd23061fb6aec66db78e7004a38bd
2 files changed
+106
-1
tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py
+101
@@ -2,11 +2,20 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import json
5
+from collections.abc import AsyncGenerator
6
7
import httpx
8
import pytest
9
10
from tests.utils import RemoteOpenAIServer
11
+from vllm.entrypoints.openai.chat_completion.batch_serving import (
12
+ OpenAIServingChatBatch,
13
+)
14
+from vllm.entrypoints.openai.chat_completion.protocol import (
15
+ BatchChatCompletionRequest,
16
+)
17
+from vllm.entrypoints.openai.engine.protocol import RequestResponseMetadata
18
+from vllm.outputs import CompletionOutput, RequestOutput
19
20
# any model with a chat template defined in tokenizer_config should work here
21
MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
@@ -221,3 +230,95 @@ async def test_batched_chat_completions_logprob_token_ids(
230
"token_id:5000",
231
sampled_token,
232
}
233
+
234
+
235
+def _make_request_output(prompt_idx: int, text: str) -> RequestOutput:
236
+ return RequestOutput(
237
+ request_id=f"req-{prompt_idx}",
238
+ prompt=None,
239
+ prompt_token_ids=[1, 2, 3],
240
+ prompt_logprobs=None,
241
+ outputs=[
242
+ CompletionOutput(
243
+ index=0,
244
+ text=text,
245
+ token_ids=[4, 5],
246
+ cumulative_logprob=None,
247
+ logprobs=None,
248
+ finish_reason="stop",
249
+ )
250
+ ],
251
+ finished=True,
252
+ )
253
+
254
+
255
+async def _generator(prompt_idx: int, text: str) -> AsyncGenerator[RequestOutput, None]:
256
+ yield _make_request_output(prompt_idx, text)
257
+
258
+
259
+@pytest.mark.asyncio
260
+@pytest.mark.skip_global_cleanup
261
+async def test_batched_echo_does_not_prepend_user_prompt() -> None:
262
+ """echo must only echo the assistant turn, never the user's prompt.
263
+
264
+ With `add_generation_prompt` (the default) the response role is `assistant`
265
+ while the last conversation message is the user's, so there is nothing to
266
+ echo; without the role check the user prompt leaks into the answer.
267
+ """
268
+ serving = OpenAIServingChatBatch.__new__(OpenAIServingChatBatch)
269
+ serving.response_role = "assistant"
270
+ serving.system_fingerprint = None
271
+
272
+ conversations = [[{"role": "user", "content": "USER PROMPT"}]]
273
+ request = BatchChatCompletionRequest(
274
+ model="test-model",
275
+ messages=conversations,
276
+ echo=True,
277
+ )
278
+
279
+ response = await serving.chat_completion_full_generator_batch(
280
+ request=request,
281
+ generators=[_generator(0, "ASSISTANT ANSWER")],
282
+ request_id="req-echo",
283
+ model_name="test-model",
284
+ all_conversations=conversations,
285
+ tokenizer=None,
286
+ request_metadata=RequestResponseMetadata(request_id="req-echo"),
287
+ )
288
+
289
+ assert response.choices[0].message.content == "ASSISTANT ANSWER"
290
+
291
+
292
+@pytest.mark.asyncio
293
+@pytest.mark.skip_global_cleanup
294
+async def test_batched_echo_prepends_matching_assistant_prefix() -> None:
295
+ """A trailing assistant turn is a real prefix and must still be echoed."""
296
+ serving = OpenAIServingChatBatch.__new__(OpenAIServingChatBatch)
297
+ serving.response_role = "assistant"
298
+ serving.system_fingerprint = None
299
+
300
+ conversations = [
301
+ [
302
+ {"role": "user", "content": "USER PROMPT"},
303
+ {"role": "assistant", "content": "PREFIX "},
304
+ ]
305
+ ]
306
+ request = BatchChatCompletionRequest(
307
+ model="test-model",
308
+ messages=conversations,
309
+ echo=True,
310
+ add_generation_prompt=False,
311
+ continue_final_message=True,
312
+ )
313
+
314
+ response = await serving.chat_completion_full_generator_batch(
315
+ request=request,
316
+ generators=[_generator(0, "ASSISTANT ANSWER")],
317
+ request_id="req-echo-2",
318
+ model_name="test-model",
319
+ all_conversations=conversations,
320
+ tokenizer=None,
321
+ request_metadata=RequestResponseMetadata(request_id="req-echo-2"),
322
+ )
323
+
324
+ assert response.choices[0].message.content == "PREFIX ASSISTANT ANSWER"
vllm/entrypoints/openai/chat_completion/batch_serving.py
+5
-1
@@ -291,7 +291,11 @@ class OpenAIServingChatBatch(OpenAIServingChat):
291
if request.echo:
292
conversation = all_conversations[prompt_idx]
293
last_msg_content: str | list[dict[str, str]] = ""
294
- if conversation and "content" in conversation[-1]:
294
+ if (
295
+ conversation
296
+ and "content" in conversation[-1]
297
+ and conversation[-1].get("role") == role
298
+ ):
299
last_msg_content = conversation[-1]["content"] or ""
300
if isinstance(last_msg_content, list):
301
last_msg_content = "\n".join(