[Bugfix][Frontend] Return logprobs in Cohere Chat v2 responses (#59072)

Signed-off-by: Mohammad Hijjawi <mohammad.hijjawi1997@gmail.com>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

tests/entrypoints/cohere/test_serving_conversion.py

@@ -25,6 +25,7 @@ from vllm.entrypoints.cohere.cohere_chat_message import (
 from vllm.entrypoints.cohere.protocol import (
     CohereChatV2Request,
     CohereChatV2Response,
+    CohereLogprobItem,
 )
 from vllm.entrypoints.cohere.serving import (
     _FINISH_REASON_MAP,
@@ -130,6 +131,8 @@ def _build_chat_completion_response(
     citations: list[Any] | None = None,
     usage: dict[str, Any] | None = None,
     kv_transfer_params: dict[str, Any] | None = None,
+    logprobs: list[tuple[str, float]] | None = None,
+    token_ids: list[int] | None = None,
 ) -> ChatCompletionResponse:
     message: dict[str, Any] = {"role": "assistant"}
     if content is not None:
@@ -151,6 +154,12 @@ def _build_chat_completion_response(
                 "message": message,
                 "finish_reason": finish_reason,
                 "stop_reason": stop_reason,
+                "logprobs": (
+                    {"content": [{"token": t, "logprob": lp} for t, lp in logprobs]}
+                    if logprobs is not None
+                    else None
+                ),
+                "token_ids": token_ids,
             }
         ],
         # ``usage`` is a required field on ChatCompletionResponse, but the
@@ -689,6 +698,8 @@ class TestBuildBaseChatCompletion:
         assert result.top_k == 50
         assert result.top_p == 0.95
         assert result.logprobs is True
+        # Needed to fill ``token_ids`` in the Cohere logprobs response.
+        assert result.return_token_ids is True
         assert result.priority == 2
         assert result.kv_transfer_params == {"x": 1}
         # ``chat_template_kwargs`` may be expanded by _apply_cohere_*; the
@@ -1907,6 +1918,7 @@ class TestChatCompletionToV2:
         assert v2.message.tool_calls is None
         assert v2.message.tool_plan is None
         assert v2.usage is None
+        assert v2.logprobs is None
 
     def test_stop_sequence_finish_reason(self):
         serving = _serving()
@@ -1990,6 +2002,19 @@ class TestChatCompletionToV2:
         v2 = serving._chat_completion_to_v2(resp, _make_request())
         assert v2.kv_transfer_params == {"k": 1}
 
+    def test_logprobs_propagated(self):
+        serving = _serving()
+        resp = _build_chat_completion_response(
+            content="Hello",
+            logprobs=[("Hel", -0.25), ("lo", -0.5)],
+            token_ids=[101, 102],
+        )
+        v2 = serving._chat_completion_to_v2(resp, _make_request(logprobs=True))
+        assert v2.logprobs == [
+            CohereLogprobItem(text="Hel", token_ids=[101], logprobs=[-0.25]),
+            CohereLogprobItem(text="lo", token_ids=[102], logprobs=[-0.5]),
+        ]
+
 
 # ======================================================================
 # _build_usage

vllm/entrypoints/cohere/serving.py

@@ -58,6 +58,7 @@ from vllm.entrypoints.cohere.protocol import (
     CohereChatV2Request,
     CohereChatV2Response,
     CohereFinishReason,
+    CohereLogprobItem,
     CohereUsage,
     CohereUsageBilledUnits,
     CohereUsageTokens,
@@ -79,6 +80,7 @@ from vllm.entrypoints.generate.base.protocol import (
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionRequest,
     ChatCompletionResponse,
+    ChatCompletionResponseChoice,
     ChatCompletionStreamResponse,
     ChatCompletionToolsParam,
     ChatMessage,
@@ -562,6 +564,8 @@ class CohereServingChatV2(OpenAIServingChat):
             frequency_penalty=request.frequency_penalty,
             presence_penalty=request.presence_penalty,
             logprobs=request.logprobs,
+            # Cohere logprob items carry token ids alongside the logprobs.
+            return_token_ids=request.logprobs,
             priority=request.priority or 0,
             kv_transfer_params=request.kv_transfer_params,
             chat_template_kwargs=request.chat_template_kwargs,
@@ -1091,9 +1095,29 @@ class CohereServingChatV2(OpenAIServingChat):
             ),
             message=assistant_msg,
             usage=usage,
+            logprobs=self._build_logprobs(choice),
             kv_transfer_params=response.kv_transfer_params,
         )
 
+    @staticmethod
+    def _build_logprobs(
+        choice: ChatCompletionResponseChoice,
+    ) -> list[CohereLogprobItem] | None:
+        if (
+            choice.logprobs is None
+            or not choice.logprobs.content
+            or choice.token_ids is None
+        ):
+            return None
+        return [
+            CohereLogprobItem(
+                text=entry.token, token_ids=[token_id], logprobs=[entry.logprob]
+            )
+            for entry, token_id in zip(
+                choice.logprobs.content, choice.token_ids, strict=True
+            )
+        ]
+
     def _create_chat_message(self, *args: Any, **kwargs: Any) -> ChatMessage:
         """Route response construction through the citation-carrying subclass.