diff --git a/tests/entrypoints/cohere/test_serving_conversion.py b/tests/entrypoints/cohere/test_serving_conversion.py index 58fa7fe793..e6696c26b4 100644 --- a/tests/entrypoints/cohere/test_serving_conversion.py +++ b/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 diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index bdaa098659..98242cc55d 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/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.