diff --git a/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py b/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py index b529f48188..0c5698cd14 100644 --- a/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py +++ b/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py @@ -2,11 +2,20 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json +from collections.abc import AsyncGenerator import httpx import pytest from tests.utils import RemoteOpenAIServer +from vllm.entrypoints.openai.chat_completion.batch_serving import ( + OpenAIServingChatBatch, +) +from vllm.entrypoints.openai.chat_completion.protocol import ( + BatchChatCompletionRequest, +) +from vllm.entrypoints.openai.engine.protocol import RequestResponseMetadata +from vllm.outputs import CompletionOutput, RequestOutput # any model with a chat template defined in tokenizer_config should work here MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" @@ -221,3 +230,95 @@ async def test_batched_chat_completions_logprob_token_ids( "token_id:5000", sampled_token, } + + +def _make_request_output(prompt_idx: int, text: str) -> RequestOutput: + return RequestOutput( + request_id=f"req-{prompt_idx}", + prompt=None, + prompt_token_ids=[1, 2, 3], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text=text, + token_ids=[4, 5], + cumulative_logprob=None, + logprobs=None, + finish_reason="stop", + ) + ], + finished=True, + ) + + +async def _generator(prompt_idx: int, text: str) -> AsyncGenerator[RequestOutput, None]: + yield _make_request_output(prompt_idx, text) + + +@pytest.mark.asyncio +@pytest.mark.skip_global_cleanup +async def test_batched_echo_does_not_prepend_user_prompt() -> None: + """echo must only echo the assistant turn, never the user's prompt. + + With `add_generation_prompt` (the default) the response role is `assistant` + while the last conversation message is the user's, so there is nothing to + echo; without the role check the user prompt leaks into the answer. + """ + serving = OpenAIServingChatBatch.__new__(OpenAIServingChatBatch) + serving.response_role = "assistant" + serving.system_fingerprint = None + + conversations = [[{"role": "user", "content": "USER PROMPT"}]] + request = BatchChatCompletionRequest( + model="test-model", + messages=conversations, + echo=True, + ) + + response = await serving.chat_completion_full_generator_batch( + request=request, + generators=[_generator(0, "ASSISTANT ANSWER")], + request_id="req-echo", + model_name="test-model", + all_conversations=conversations, + tokenizer=None, + request_metadata=RequestResponseMetadata(request_id="req-echo"), + ) + + assert response.choices[0].message.content == "ASSISTANT ANSWER" + + +@pytest.mark.asyncio +@pytest.mark.skip_global_cleanup +async def test_batched_echo_prepends_matching_assistant_prefix() -> None: + """A trailing assistant turn is a real prefix and must still be echoed.""" + serving = OpenAIServingChatBatch.__new__(OpenAIServingChatBatch) + serving.response_role = "assistant" + serving.system_fingerprint = None + + conversations = [ + [ + {"role": "user", "content": "USER PROMPT"}, + {"role": "assistant", "content": "PREFIX "}, + ] + ] + request = BatchChatCompletionRequest( + model="test-model", + messages=conversations, + echo=True, + add_generation_prompt=False, + continue_final_message=True, + ) + + response = await serving.chat_completion_full_generator_batch( + request=request, + generators=[_generator(0, "ASSISTANT ANSWER")], + request_id="req-echo-2", + model_name="test-model", + all_conversations=conversations, + tokenizer=None, + request_metadata=RequestResponseMetadata(request_id="req-echo-2"), + ) + + assert response.choices[0].message.content == "PREFIX ASSISTANT ANSWER" diff --git a/vllm/entrypoints/openai/chat_completion/batch_serving.py b/vllm/entrypoints/openai/chat_completion/batch_serving.py index 275ad13f06..1b0934e501 100644 --- a/vllm/entrypoints/openai/chat_completion/batch_serving.py +++ b/vllm/entrypoints/openai/chat_completion/batch_serving.py @@ -291,7 +291,11 @@ class OpenAIServingChatBatch(OpenAIServingChat): if request.echo: conversation = all_conversations[prompt_idx] last_msg_content: str | list[dict[str, str]] = "" - if conversation and "content" in conversation[-1]: + if ( + conversation + and "content" in conversation[-1] + and conversation[-1].get("role") == role + ): last_msg_content = conversation[-1]["content"] or "" if isinstance(last_msg_content, list): last_msg_content = "\n".join(