fixup! [Feature]: Add OpenAI server prompt_logprobs support vllm-proj…

…ect#6508

tests/entrypoints/openai/test_completion.py

@@ -130,6 +130,7 @@ async def test_single_completion(client: openai.AsyncOpenAI, model_name: str,
         temperature=0.0,
     )
     assert len(completion.choices[0].text) >= 1
+    assert completion.choices[0].prompt_logprobs is None
 
 
 @pytest.mark.asyncio
@@ -270,10 +271,10 @@ async def test_too_many_completion_logprobs(client: openai.AsyncOpenAI,
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "model_name, prompt_logprobs",
-    [(MODEL_NAME, True), (MODEL_NAME, False), (MODEL_NAME, None)],
+    [(MODEL_NAME, 1), (MODEL_NAME, 0), (MODEL_NAME, -1)],
 )
-async def test_prompt_logprobs(client: openai.AsyncOpenAI, model_name: str,
-                               prompt_logprobs: bool):
+async def test_prompt_logprobs_chat(client: openai.AsyncOpenAI,
+                                    model_name: str, prompt_logprobs: int):
     params: Dict = {
         "messages": [{
             "role": "system",
@@ -291,18 +292,102 @@ async def test_prompt_logprobs(client: openai.AsyncOpenAI, model_name: str,
             "content": "Where was it played?"
         }],
         "model":
-        model_name
+        model_name,
+        "extra_body": {
+            "prompt_logprobs": prompt_logprobs
+        }
     }
 
-    if prompt_logprobs is not None:
-        params["extra_body"] = {"prompt_logprobs": prompt_logprobs}
+    if prompt_logprobs < 0:
+        with pytest.raises(BadRequestError) as err_info:
+            await client.chat.completions.create(**params)
+        expected_err_string = (
+            "Error code: 400 - {'object': 'error', 'message': "
+            "'Prompt_logprobs set to invalid negative value: -1',"
+            " 'type': 'BadRequestError', 'param': None, 'code': 400}")
+        assert str(err_info.value) == expected_err_string
+    else:
+        completion = await client.chat.completions.create(**params)
+        if prompt_logprobs > 0:
+            assert completion.prompt_logprobs is not None
+            assert len(completion.prompt_logprobs) > 0
+        else:
+            assert completion.prompt_logprobs is None
+
 
-    completion = await client.chat.completions.create(**params)
-    if prompt_logprobs:
-        assert completion.prompt_logprobs is not None
-        assert len(completion.prompt_logprobs) > 0
-    elif prompt_logprobs is None or not prompt_logprobs:
-        assert completion.prompt_logprobs is None
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "model_name",
+    [MODEL_NAME],
+)
+async def test_more_than_one_prompt_logprobs_chat(client: openai.AsyncOpenAI,
+                                                  model_name: str):
+    params: Dict = {
+        "messages": [{
+            "role": "system",
+            "content": "You are a helpful assistant."
+        }, {
+            "role": "user",
+            "content": "Who won the world series in 2020?"
+        }, {
+            "role":
+            "assistant",
+            "content":
+            "The Los Angeles Dodgers won the World Series in 2020."
+        }, {
+            "role": "user",
+            "content": "Where was it played?"
+        }],
+        "model":
+        model_name,
+        "extra_body": {
+            "prompt_logprobs": 1
+        }
+    }
+
+    completion_1 = await client.chat.completions.create(**params)
+
+    params["extra_body"] = {"prompt_logprobs": 2}
+    completion_2 = await client.chat.completions.create(**params)
+
+    assert len(completion_1.prompt_logprobs[3]) == 1
+    assert len(completion_2.prompt_logprobs[3]) == 2
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name, prompt_logprobs", [(MODEL_NAME, -1),
+                                                         (MODEL_NAME, 0),
+                                                         (MODEL_NAME, 1)])
+async def test_prompt_logprobs_completion(client: openai.AsyncOpenAI,
+                                          model_name: str,
+                                          prompt_logprobs: int):
+    params: Dict = {
+        "prompt": ["A robot may not injure another robot", "My name is"],
+        "model": model_name,
+        "extra_body": {
+            "prompt_logprobs": prompt_logprobs
+        }
+    }
+
+    if prompt_logprobs < 0:
+        with pytest.raises(BadRequestError) as err_info:
+            await client.completions.create(**params)
+        expected_err_string = (
+            "Error code: 400 - {'object': 'error', 'message': "
+            "'Prompt_logprobs set to invalid negative value: -1',"
+            " 'type': 'BadRequestError', 'param': None, 'code': 400}")
+        assert str(err_info.value) == expected_err_string
+    else:
+        completion = await client.completions.create(**params)
+        if prompt_logprobs > 0:
+            assert completion.choices[0].prompt_logprobs is not None
+            assert len(completion.choices[0].prompt_logprobs) > 0
+
+            assert completion.choices[1].prompt_logprobs is not None
+            assert len(completion.choices[1].prompt_logprobs) > 0
+
+        else:
+            assert completion.choices[0].prompt_logprobs is None
 
 
 @pytest.mark.asyncio

vllm/entrypoints/openai/protocol.py

@@ -153,6 +153,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
     skip_special_tokens: bool = True
     spaces_between_special_tokens: bool = True
     truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None
+    prompt_logprobs: int = 0
     # doc: end-chat-completion-sampling-params
 
     # doc: begin-chat-completion-extra-params
@@ -162,13 +163,6 @@ class ChatCompletionRequest(OpenAIBaseModel):
             "If true, the new message will be prepended with the last message "
             "if they belong to the same role."),
     )
-
-    prompt_logprobs: bool = Field(
-        default=None,
-        description=
-        "If true, the logprobs from the decoded prompt will be output.",
-    )
-
     add_generation_prompt: bool = Field(
         default=True,
         description=
@@ -271,8 +265,8 @@ def to_sampling_params(
             stop=self.stop,
             stop_token_ids=self.stop_token_ids,
             logprobs=self.top_logprobs if self.logprobs else None,
-            prompt_logprobs=self.prompt_logprobs if self.prompt_logprobs else
-            (self.top_logprobs if self.echo else None),
+            prompt_logprobs=self.prompt_logprobs if self.prompt_logprobs > 0
+            else (self.top_logprobs if self.echo else None),
             ignore_eos=self.ignore_eos,
             max_tokens=max_tokens,
             min_tokens=self.min_tokens,
@@ -377,6 +371,7 @@ class CompletionRequest(OpenAIBaseModel):
     spaces_between_special_tokens: bool = True
     truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None
     allowed_token_ids: Optional[List[int]] = None
+    prompt_logprobs: int = 0
     # doc: end-completion-sampling-params
 
     # doc: begin-completion-extra-params
@@ -463,7 +458,8 @@ def to_sampling_params(
             min_tokens=self.min_tokens,
             use_beam_search=self.use_beam_search,
             early_stopping=self.early_stopping,
-            prompt_logprobs=self.logprobs if self.echo else None,
+            prompt_logprobs=self.prompt_logprobs if self.prompt_logprobs > 0
+            else self.logprobs if self.echo else None,
             skip_special_tokens=self.skip_special_tokens,
             spaces_between_special_tokens=self.spaces_between_special_tokens,
             include_stop_str_in_output=self.include_stop_str_in_output,
@@ -541,6 +537,8 @@ class CompletionResponseChoice(OpenAIBaseModel):
             "to stop, None if the completion finished for some other reason "
             "including encountering the EOS token"),
     )
+    prompt_logprobs: Optional[List[Optional[Dict[int, Logprob]]]] = Field(
+        default=None)
 
 
 class CompletionResponse(OpenAIBaseModel):

vllm/entrypoints/openai/serving_chat.py

@@ -83,17 +83,14 @@ async def create_chat_completion(
         if error_check_ret is not None:
             return error_check_ret
 
-        try:
-            if request.stream:
-                assert (request.prompt_logprobs is None
-                        or not request.prompt_logprobs), (
-                            "Prompt_logprobs are not "
-                            "available when stream is "
-                            "enabled")
-        except ValueError as e:
-            logger.error(
-                "Mutually exclusive options in chat completion request: %s", e)
-            return self.create_error_response(str(e))
+        if request.stream and request.prompt_logprobs > 0:
+            return self.create_error_response(
+                "Prompt_logprobs are not available when stream is enabled")
+
+        if request.prompt_logprobs < 0:
+            return self.create_error_response(
+                f"Prompt_logprobs set to invalid "
+                f"negative value: {request.prompt_logprobs}")
 
         try:
             (

vllm/entrypoints/openai/serving_completion.py

@@ -83,6 +83,13 @@ async def create_completion(self, request: CompletionRequest,
         model_name = self.served_model_names[0]
         request_id = f"cmpl-{random_uuid()}"
         created_time = int(time.time())
+        if request.stream and request.prompt_logprobs > 0:
+            return self.create_error_response(
+                "Prompt_logprobs are not available when stream is enabled")
+        elif request.prompt_logprobs < 0:
+            return self.create_error_response(
+                f"Prompt_logprobs set to invalid negative "
+                f"value: {request.prompt_logprobs}")
 
         # Schedule the request and get the result generator.
         generators: List[AsyncGenerator[RequestOutput, None]] = []
@@ -377,7 +384,7 @@ def request_output_to_completion_response(
                     logprobs=logprobs,
                     finish_reason=output.finish_reason,
                     stop_reason=output.stop_reason,
-                )
+                    prompt_logprobs=final_res.prompt_logprobs)
                 choices.append(choice_data)
 
             num_prompt_tokens += len(prompt_token_ids)