diff --git a/responses_api_models/vllm_model/app.py b/responses_api_models/vllm_model/app.py index 5767261ec5..4fe3701427 100644 --- a/responses_api_models/vllm_model/app.py +++ b/responses_api_models/vllm_model/app.py @@ -470,11 +470,7 @@ def _preprocess_chat_completion_create_params(self, request: Request, body_dict: top_logprobs=0, # Typically passed via OpenAI client extra_body. return_tokens_as_token_ids=True, - # TODO add this when NeMo RL upgrades to vLLM 0.10.2 support for prompt token ids - # For prompt and generation token IDs - # return_token_ids=True, - # For prompt token IDs - # prompt_logprobs=0, + return_token_ids=True, ) if self.config.uses_reasoning_parser and not self.config.preserve_reasoning_in_assistant_content: @@ -744,48 +740,32 @@ async def chat_completions( ) log_probs = logprobs_block["content"] generation_log_probs = [log_prob["logprob"] for log_prob in log_probs] - - """ - START TODO remove this when NeMo RL upgrades to vLLM 0.10.2 support for prompt token ids - """ - # Looks like `"token_id:151667"` generation_token_ids = [log_prob["token"].removeprefix("token_id:") for log_prob in log_probs] - # The tokenize endpoint doesn't accept any sampling parameters - # The only relevant params are model, messages, and tools. - # - # IMPORTANT: pass through chat-template knobs (e.g. enable_thinking) - # when tokenizing, otherwise `prompt_token_ids` (and therefore logged - # `prompt_str`) can be built with different chat template settings than - # the actual generation request. - tokenize_body_dict = dict() - for key in ("model", "messages", "tools", "chat_template_kwargs"): - if key in body_dict: - tokenize_body_dict[key] = body_dict[key] - - # The base url has /v1 at the end but vLLM's tokenize endpoint does not have v1, hence the .. - tokenize_response = await client.create_tokenize(**tokenize_body_dict) - """ - END - """ + # Fast path: use the IDs from the original generation request so the + # prompt does not need to be rendered and tokenized a second time. + prompt_token_ids = chat_completion_dict.pop("prompt_token_ids", None) + + # Older vLLM versions may accept return_token_ids without returning + # prompt_token_ids. Preserve /tokenize as a compatibility fallback. + if prompt_token_ids is None: + tokenize_body_dict = { + key: body_dict[key] + for key in ("model", "messages", "tools", "chat_template_kwargs") + if key in body_dict + } + tokenize_response = await client.create_tokenize(**tokenize_body_dict) + prompt_token_ids = tokenize_response["tokens"] - message_dict = choice_dict["message"] - message_dict.update( - dict( - # TODO add this when NeMo RL upgrades to vLLM 0.10.2 support for prompt token ids - # prompt_token_ids=chat_completion_dict["prompt_token_ids"], - prompt_token_ids=tokenize_response["tokens"], - # generation_token_ids=choice_dict["token_ids"], - generation_token_ids=generation_token_ids, - generation_log_probs=generation_log_probs, - ) + choice_dict["message"].update( + prompt_token_ids=prompt_token_ids, + generation_token_ids=generation_token_ids, + generation_log_probs=generation_log_probs, ) - # Clean the duplicated information + # Remove vLLM-specific duplicate fields before validating the response. choice_dict.pop("logprobs") - # TODO add this when NeMo RL upgrades to vLLM 0.10.2 support for prompt token ids - # chat_completion_dict.pop("prompt_token_ids") - # choice_dict.pop("token_ids") + choice_dict.pop("token_ids", None) return NeMoGymChatCompletion.model_validate(chat_completion_dict) diff --git a/responses_api_models/vllm_model/tests/test_app.py b/responses_api_models/vllm_model/tests/test_app.py index f6c65acc1f..2aadab7bc6 100644 --- a/responses_api_models/vllm_model/tests/test_app.py +++ b/responses_api_models/vllm_model/tests/test_app.py @@ -4597,6 +4597,7 @@ def test_capture_path_pins_top_logprobs_to_zero(self) -> None: assert result["top_logprobs"] == 0 assert result["logprobs"] is True assert result["return_tokens_as_token_ids"] is True + assert result["return_token_ids"] is True # Inbound non-zero value must also be overridden, not inherited. result = model._preprocess_chat_completion_create_params( @@ -4661,6 +4662,43 @@ def _capture_chat_completion_dict( ], } + @mark.parametrize("prompt_token_ids", [[10, 20, 30], []]) + def test_capture_path_uses_native_prompt_token_ids(self, prompt_token_ids: list[int]) -> None: + """Native prompt IDs, including an empty list, must bypass /tokenize.""" + model = _make_top_logprobs_model(return_token_id_information=True) + app = model.setup_webserver() + captured_kwargs: dict[str, Any] = {} + + async def mock_create_chat_completion(**kwargs): + captured_kwargs.update(kwargs) + response = self._capture_chat_completion_dict( + logprobs={ + "content": [ + {"token": "token_id:123", "logprob": -0.1, "bytes": None, "top_logprobs": []}, + ] + } + ) + response["prompt_token_ids"] = prompt_token_ids + response["choices"][0]["token_ids"] = [123] + return response + + mock_client = MagicMock(spec=NeMoGymAsyncOpenAI) + mock_client.create_chat_completion = AsyncMock(side_effect=mock_create_chat_completion) + mock_client.create_tokenize = AsyncMock() + model._clients = [mock_client] + + response = TestClient(app).post( + "/v1/chat/completions", + json={"messages": [{"role": "user", "content": "hi"}]}, + ) + + assert response.status_code == 200 + assert captured_kwargs["return_token_ids"] is True + assert response.json()["choices"][0]["message"]["prompt_token_ids"] == prompt_token_ids + assert "prompt_token_ids" not in response.json() + assert "token_ids" not in response.json()["choices"][0] + mock_client.create_tokenize.assert_not_called() + def test_capture_path_succeeds_with_inbound_null_top_logprobs(self) -> None: """End-to-end regression: a request with top_logprobs=null no longer empties capture; token ids and logprobs come back populated (and coerced to int).""" @@ -4668,6 +4706,7 @@ def test_capture_path_succeeds_with_inbound_null_top_logprobs(self) -> None: app = model.setup_webserver() captured_kwargs: dict[str, Any] = {} + captured_tokenize_kwargs: dict[str, Any] = {} async def mock_create_chat_completion(**kwargs): captured_kwargs.update(kwargs) @@ -4681,6 +4720,7 @@ async def mock_create_chat_completion(**kwargs): ) async def mock_create_tokenize(**kwargs): + captured_tokenize_kwargs.update(kwargs) return {"tokens": [10, 20, 30]} mock_client = MagicMock(spec=NeMoGymAsyncOpenAI) @@ -4691,18 +4731,92 @@ async def mock_create_tokenize(**kwargs): client = TestClient(app) response = client.post( "/v1/chat/completions", - json={"messages": [{"role": "user", "content": "hi"}], "top_logprobs": None}, + json={ + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "description": "Look something up", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + "metadata": {"chat_template_kwargs": '{"enable_thinking": false}'}, + "top_logprobs": None, + }, ) assert response.status_code == 200 # The request forwarded to vLLM had top_logprobs pinned to 0, not the inbound null. assert captured_kwargs["top_logprobs"] == 0 + assert captured_kwargs["return_token_ids"] is True + assert captured_tokenize_kwargs == { + key: captured_kwargs[key] for key in ("model", "messages", "tools", "chat_template_kwargs") + } message = response.json()["choices"][0]["message"] assert message["generation_token_ids"] == [123, 456] assert message["generation_log_probs"] == [-0.1, -0.2] assert message["prompt_token_ids"] == [10, 20, 30] + def test_capture_path_falls_back_for_null_prompt_token_ids(self) -> None: + model = _make_top_logprobs_model(return_token_id_information=True) + app = model.setup_webserver() + + async def mock_create_chat_completion(**kwargs): + response = self._capture_chat_completion_dict( + logprobs={ + "content": [ + {"token": "token_id:123", "logprob": -0.1, "bytes": None, "top_logprobs": []}, + ] + } + ) + response["prompt_token_ids"] = None + return response + + mock_client = MagicMock(spec=NeMoGymAsyncOpenAI) + mock_client.create_chat_completion = AsyncMock(side_effect=mock_create_chat_completion) + mock_client.create_tokenize = AsyncMock(return_value={"tokens": [10, 20]}) + model._clients = [mock_client] + + response = TestClient(app).post( + "/v1/chat/completions", + json={"messages": [{"role": "user", "content": "hi"}]}, + ) + + assert response.status_code == 200 + assert response.json()["choices"][0]["message"]["prompt_token_ids"] == [10, 20] + mock_client.create_tokenize.assert_awaited_once() + + def test_capture_path_propagates_tokenize_failure(self) -> None: + model = _make_top_logprobs_model(return_token_id_information=True) + app = model.setup_webserver() + tokenize_error = RuntimeError("tokenize failed") + + async def mock_create_chat_completion(**kwargs): + return self._capture_chat_completion_dict( + logprobs={ + "content": [ + {"token": "token_id:123", "logprob": -0.1, "bytes": None, "top_logprobs": []}, + ] + } + ) + + mock_client = MagicMock(spec=NeMoGymAsyncOpenAI) + mock_client.create_chat_completion = AsyncMock(side_effect=mock_create_chat_completion) + mock_client.create_tokenize = AsyncMock(side_effect=tokenize_error) + model._clients = [mock_client] + + with raises(RuntimeError, match="tokenize failed") as exc_info: + TestClient(app).post( + "/v1/chat/completions", + json={"messages": [{"role": "user", "content": "hi"}]}, + ) + + assert exc_info.value is tokenize_error + def test_capture_path_preserves_routed_experts(self) -> None: model = _make_top_logprobs_model(return_token_id_information=True) app = model.setup_webserver()