diff --git a/components/src/dynamo/frontend/sglang_processor.py b/components/src/dynamo/frontend/sglang_processor.py index f5b859a47911..26c179e58294 100644 --- a/components/src/dynamo/frontend/sglang_processor.py +++ b/components/src/dynamo/frontend/sglang_processor.py @@ -41,7 +41,13 @@ detect_force_reasoning_from_template, preprocess_chat_request, ) -from .utils import PreprocessError, extract_mm_urls, random_uuid, worker_warmup +from .utils import ( + PreprocessError, + extract_mm_urls, + nvext_extra_field_requested, + random_uuid, + worker_warmup, +) logger = logging.getLogger(__name__) @@ -197,13 +203,17 @@ def _build_dynamo_preproc( max_tokens = request.get("max_completion_tokens") or request.get("max_tokens") stop = request.get("stop") + stop_token_ids = request.get("stop_token_ids", []) if isinstance(stop, str): stop = [stop] + elif isinstance(stop, list) and all( + isinstance(item, int) and not isinstance(item, bool) for item in stop + ): + stop_token_ids = [*stop_token_ids, *stop] + stop = [] elif stop is None: stop = [] - stop_token_ids = request.get("stop_token_ids", []) - # Handle logprobs logprobs_val = None logprobs = request.get("logprobs") @@ -246,6 +256,7 @@ def _build_dynamo_preproc( # (e.g. <|tool_call|>) to detect calls. Mirrors the # post-processor's _skip_special_tokens logic. "skip_special_tokens": tool_call_parser is None, + "return_tokens_as_token_ids": request.get("return_tokens_as_token_ids"), }, "eos_token_ids": [eos_token_id] if eos_token_id is not None else [], "annotations": [], @@ -519,6 +530,7 @@ async def _generate_and_stream( new_ids = engine_response["token_ids"] raw_finish = engine_response.get("finish_reason") finish_reason = _map_finish_reason(raw_finish) + stop_reason = engine_response.get("stop_reason") if usage := engine_response.get("completion_usage"): pending_usage = usage @@ -554,6 +566,10 @@ async def _generate_and_stream( } if pending_usage: dynamo_out["usage"] = pending_usage + if stop_reason is not None and nvext_extra_field_requested( + request, "stop_reason" + ): + dynamo_out["nvext"] = {"stop_reason": stop_reason} yield dynamo_out diff --git a/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py b/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py index 36b1c33ab4a3..d6c9b289cf53 100644 --- a/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py +++ b/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py @@ -10,6 +10,7 @@ """ +import asyncio import json import sys import types @@ -34,12 +35,18 @@ ) from dynamo.frontend.sglang_processor import ( SglangPreprocessWorkerResult, + SglangProcessor, _build_dynamo_preproc, _init_worker, _map_finish_reason, _runtime_config_parser_name, ) -from dynamo.frontend.utils import PreprocessError, random_call_id, random_uuid +from dynamo.frontend.utils import ( + PreprocessError, + nvext_extra_field_requested, + random_call_id, + random_uuid, +) # Needs sglang packages (gpu_1 container). No need for parallel marker. pytestmark = [ @@ -252,6 +259,62 @@ def test_model_name_and_token_ids(self): assert result["model"] == "my-model" assert result["token_ids"] == [10, 20, 30] + def test_stop_token_id_array_maps_to_stop_token_ids(self): + """Integer stop arrays are token-id stops, not string stops.""" + result = _build_dynamo_preproc( + {"model": "test", "stop": [32, 34]}, + [1], + "test", + None, + ) + + assert result["stop_conditions"]["stop"] == [] + assert result["stop_conditions"]["stop_token_ids"] == [32, 34] + + def test_string_stops_remain_string_stops(self): + """String stops are forwarded as string stops.""" + result = _build_dynamo_preproc( + {"model": "test", "stop": " The"}, + [1], + "test", + None, + ) + + assert result["stop_conditions"]["stop"] == [" The"] + assert result["stop_conditions"]["stop_token_ids"] == [] + + result = _build_dynamo_preproc( + {"model": "test", "stop": ["A", "B"]}, + [1], + "test", + None, + ) + + assert result["stop_conditions"]["stop"] == ["A", "B"] + assert result["stop_conditions"]["stop_token_ids"] == [] + + def test_token_id_display_string_remains_string_stop(self): + """token_id:N strings are output display strings, not token-id stops.""" + result = _build_dynamo_preproc( + {"model": "test", "stop": "token_id:576"}, + [1], + "test", + None, + ) + + assert result["stop_conditions"]["stop"] == ["token_id:576"] + assert result["stop_conditions"]["stop_token_ids"] == [] + + result = _build_dynamo_preproc( + {"model": "test", "stop": ["token_id:576"]}, + [1], + "test", + None, + ) + + assert result["stop_conditions"]["stop"] == ["token_id:576"] + assert result["stop_conditions"]["stop_token_ids"] == [] + # --------------------------------------------------------------------------- # _map_finish_reason @@ -1567,6 +1630,58 @@ def test_finish_reason_only(self, tokenizer): assert choice is not None assert choice["finish_reason"] == "stop" + def test_stop_reason_not_emitted_on_choice(self, tokenizer): + """Backend stop_reason is not part of the OpenAI choice shape.""" + post = SglangStreamingPostProcessor( + tokenizer=tokenizer, tool_call_parser=None, reasoning_parser=None + ) + + choice = post.process_output( + {"token_ids": [], "finish_reason": "stop", "stop_reason": "END"} + ) + + assert choice is not None + assert "stop_reason" not in choice + + def test_stop_reason_emits_in_nvext_when_requested(self, tokenizer): + """Frontend emits backend stop_reason under nvext when requested.""" + + class FakeRouter: + async def generate(self, *args, **kwargs): + yield { + "token_ids": [], + "finish_reason": "stop", + "stop_reason": "END", + } + + async def collect(): + processor = SglangProcessor( + tokenizer=tokenizer, + router=FakeRouter(), + tool_call_parser_name=None, + reasoning_parser_name=None, + eos_token_id=None, + ) + post = SglangStreamingPostProcessor( + tokenizer=tokenizer, tool_call_parser=None, reasoning_parser=None + ) + request = { + "model": "test-model", + "nvext": {"extra_fields": ["stop_reason"]}, + } + return [ + item + async for item in processor._generate_and_stream( + "req-stop", request, {}, [], post + ) + ] + + items = asyncio.run(collect()) + + assert len(items) == 1 + assert items[0]["nvext"]["stop_reason"] == "END" + assert "stop_reason" not in items[0]["choices"][0] + def test_lookback_trimming(self, tokenizer): """Verify _all_token_ids doesn't grow unbounded.""" post = SglangStreamingPostProcessor( @@ -1684,6 +1799,13 @@ def test_preprocess_error(self): # FRONTEND.8 err = PreprocessError("n=2 unsupported") assert "n=2" in str(err) + def test_nvext_extra_field_requested(self): + assert nvext_extra_field_requested( + {"nvext": {"extra_fields": ["stop_reason"]}}, "stop_reason" + ) + assert not nvext_extra_field_requested({"nvext": {}}, "stop_reason") + assert not nvext_extra_field_requested({}, "stop_reason") + # --------------------------------------------------------------------------- # SglangPreprocessWorkerResult picklability diff --git a/components/src/dynamo/frontend/utils.py b/components/src/dynamo/frontend/utils.py index 2ef3f3f0ed71..310cf5c2ae9e 100644 --- a/components/src/dynamo/frontend/utils.py +++ b/components/src/dynamo/frontend/utils.py @@ -20,6 +20,15 @@ def random_call_id() -> str: return f"call_{uuid.uuid4().int & _MASK_64_BITS:016x}" +def nvext_extra_field_requested(request: dict[str, Any], field: str) -> bool: + """Return whether a request opted into a response nvext field.""" + nvext = request.get("nvext") + if not isinstance(nvext, dict): + return False + extra_fields = nvext.get("extra_fields") + return isinstance(extra_fields, list) and field in extra_fields + + def worker_warmup() -> bool: """Dummy task to ensure a ProcessPoolExecutor worker is fully initialized.""" return True diff --git a/components/src/dynamo/sglang/CLAUDE.md b/components/src/dynamo/sglang/CLAUDE.md index b273d5de3341..ffff7abd3a81 100644 --- a/components/src/dynamo/sglang/CLAUDE.md +++ b/components/src/dynamo/sglang/CLAUDE.md @@ -228,6 +228,12 @@ absolute sequence position where logprob computation starts: `-1` (default) = ou only (`len(prompt) - 1`), `0` = from prompt start. We set it to 0 when `prompt_logprobs` is requested. +**Top-logprobs gate**: `logprobs >= 1` (or `prompt_logprobs >= 1`) raises `ValueError` +by default. SGLang's tokenizer manager detokenizes top-k tokens per-position serially, +causing severe latency degradation (O(N) per generated token). Callers must use +`logprobs=0` for chosen-token-only logprobs. Set `DYN_SGL_ALLOW_TOP_LOGPROBS=1` to +override once upstream batches `detokenize_top_logprobs_tokens`. + **Streaming behavior** (`_extract_logprobs`): Dynamo forces `stream_output=True` (args.py:374), making `output_ids` disjoint per chunk. diff --git a/components/src/dynamo/sglang/protocol.py b/components/src/dynamo/sglang/protocol.py index 18dccabad3b1..b359dc4b37f3 100644 --- a/components/src/dynamo/sglang/protocol.py +++ b/components/src/dynamo/sglang/protocol.py @@ -19,6 +19,7 @@ class StopConditions(BaseModel): max_tokens: Optional[int] = None stop: Optional[List[str]] = None + stop_token_ids: Optional[List[TokenIdType]] = None stop_token_ids_hidden: Optional[List[TokenIdType]] = None min_tokens: Optional[int] = None ignore_eos: Optional[bool] = None diff --git a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py index 3295d622b254..55d9748e8ec9 100644 --- a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py +++ b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py @@ -3,6 +3,7 @@ import asyncio import logging +import os import time from typing import Any, AsyncGenerator, Dict, Optional @@ -18,6 +19,26 @@ from dynamo.sglang.publisher import DynamoSglangPublisher from dynamo.sglang.request_handlers.handler_base import BaseWorkerHandler +# Escape hatch: set to "1" (or any truthy value) to allow top_logprobs_num >= 1. +# Default-off because SGLang's tokenizer manager detokenizes top-k tokens +# per-position serially (O(N) per generated token), causing severe latency +# degradation. Flip once upstream lands batched top-logprob detokenization: +# https://github.com/sgl-project/sglang/pull/24447 +_ALLOW_TOP_LOGPROBS_ENV = "DYN_SGL_ALLOW_TOP_LOGPROBS" + +_TOP_LOGPROBS_UNSUPPORTED_MSG = ( + "Dynamo's SGLang backend does not currently support logprobs >= 1 due to " + "an O(N) per-position detokenization in the upstream sglang tokenizer " + "manager. Use logprobs=0 for chosen-token logprobs, or set " + "DYN_SGL_ALLOW_TOP_LOGPROBS=1 to override at your own risk. " + "Track the upstream fix at https://github.com/sgl-project/sglang/pull/24447." +) + + +def _top_logprobs_allowed() -> bool: + """Return True if the DYN_SGL_ALLOW_TOP_LOGPROBS escape hatch is enabled.""" + return os.environ.get(_ALLOW_TOP_LOGPROBS_ENV, "").lower() not in ("", "0", "false") + def _extract_media_urls(mm_data: Dict[str, Any], media_key: str) -> list[str] | None: """Normalize multimodal URL items from the frontend wire format.""" @@ -40,6 +61,90 @@ def _extract_media_urls(mm_data: Dict[str, Any], media_key: str) -> list[str] | return urls or None +def _nvext_extra_field_requested(request: Dict[str, Any], field: str) -> bool: + nvext = request.get("nvext") + if not isinstance(nvext, dict): + return False + extra_fields = nvext.get("extra_fields") + if not isinstance(extra_fields, list): + return False + return field in extra_fields + + +def _user_stop_token_ids(request: Dict[str, Any]) -> set[int]: + stop_conditions = request.get("stop_conditions") + if isinstance(stop_conditions, dict): + return { + token_id + for token_id in (stop_conditions.get("stop_token_ids") or []) + if isinstance(token_id, int) and not isinstance(token_id, bool) + } + + stop = request.get("stop") + if isinstance(stop, list) and all( + isinstance(item, int) and not isinstance(item, bool) for item in stop + ): + return set(stop) + + return { + token_id + for token_id in (request.get("stop_token_ids") or []) + if isinstance(token_id, int) and not isinstance(token_id, bool) + } + + +def _openai_stop_sampling_params(request: Dict[str, Any]) -> Dict[str, Any]: + stop = request.get("stop") + if isinstance(stop, str): + return {"stop": stop} + if isinstance(stop, list): + if stop and all( + isinstance(item, int) and not isinstance(item, bool) for item in stop + ): + return {"stop_token_ids": stop} + if stop and all(isinstance(item, str) for item in stop): + return {"stop": stop} + + stop_token_ids = [ + token_id + for token_id in (request.get("stop_token_ids") or []) + if isinstance(token_id, int) and not isinstance(token_id, bool) + ] + if stop_token_ids: + return {"stop_token_ids": stop_token_ids} + return {} + + +def _extract_sglang_stop_reason( + finish_reason: Dict[str, Any] | None, + user_stop_token_ids: set[int] | None = None, +) -> Any | None: + """Extract SGLang's matched stop value for Dynamo's stop_reason field.""" + + if not finish_reason: + return None + + matched = finish_reason.get("matched") + if isinstance(matched, bool): + return None + if isinstance(matched, str): + return matched + if isinstance(matched, int): + if user_stop_token_ids is not None and matched not in user_stop_token_ids: + return None + return matched + if isinstance(matched, list) and all( + isinstance(item, int) and not isinstance(item, bool) for item in matched + ): + if user_stop_token_ids is not None and any( + item not in user_stop_token_ids for item in matched + ): + return None + return matched + + return None + + class DecodeWorkerHandler(BaseWorkerHandler): """Handler for decode workers in both aggregated and disaggregated serving modes.""" @@ -143,6 +248,7 @@ def _build_sampling_params(self, request: Dict[str, Any]) -> Dict[str, Any]: "top_k": request.get("top_k"), "n": request.get("n"), "max_new_tokens": request.get("max_tokens"), + **_openai_stop_sampling_params(request), **self._get_guided_decoding_params(request.get("guided_decoding")), } @@ -181,54 +287,58 @@ def _build_logprob_kwargs(request: Dict[str, Any]) -> Dict[str, Any]: if not output_options: return kwargs - logprobs_value = output_options.get("logprobs") - if logprobs_value is not None: + allow_top = _top_logprobs_allowed() + + def _parse(name: str, value: Any) -> Optional[int]: try: - parsed = int(logprobs_value) - if parsed < 0: - logging.warning( - f"Invalid logprobs value: {logprobs_value} " - "(must be non-negative), ignoring" - ) - else: - kwargs["return_logprob"] = True - kwargs["top_logprobs_num"] = parsed + parsed = int(value) except (ValueError, TypeError): logging.warning( - f"Invalid logprobs value: {logprobs_value} " - "(must be integer), ignoring" + f"Invalid {name} value: {value} (must be integer), ignoring" ) + return None + if parsed < 0: + logging.warning( + f"Invalid {name} value: {value} (must be non-negative), ignoring" + ) + return None + if parsed >= 1 and not allow_top: + raise ValueError(_TOP_LOGPROBS_UNSUPPORTED_MSG) + return parsed + + logprobs_value = output_options.get("logprobs") + if logprobs_value is not None: + parsed = _parse("logprobs", logprobs_value) + if parsed is not None: + kwargs["return_logprob"] = True + kwargs["top_logprobs_num"] = parsed prompt_logprobs_value = output_options.get("prompt_logprobs") if prompt_logprobs_value is not None: - try: - parsed = int(prompt_logprobs_value) - if parsed < 0: - logging.warning( - f"Invalid prompt_logprobs value: {prompt_logprobs_value} " - "(must be non-negative), ignoring" - ) - else: - kwargs["return_logprob"] = True - # SGLang has a single top_logprobs_num for both prompt - # and output tokens, so take the max of the two. - kwargs["top_logprobs_num"] = max( - kwargs.get("top_logprobs_num", 0), parsed - ) - # logprob_start_len=0 computes from prompt start; - # omitting it (or -1) computes output tokens only. - kwargs["logprob_start_len"] = 0 - except (ValueError, TypeError): - logging.warning( - f"Invalid prompt_logprobs value: {prompt_logprobs_value} " - "(must be integer), ignoring" + parsed = _parse("prompt_logprobs", prompt_logprobs_value) + if parsed is not None: + kwargs["return_logprob"] = True + # SGLang has a single top_logprobs_num for both prompt + # and output tokens, so take the max of the two. + kwargs["top_logprobs_num"] = max( + kwargs.get("top_logprobs_num", 0), parsed ) + # logprob_start_len=0 computes from prompt start; + # omitting it (or -1) computes output tokens only. + kwargs["logprob_start_len"] = 0 + + # Belt-and-suspenders: if return_logprob was requested and the gate is + # not open, pin top_logprobs_num=0 so no future code path can flip it on. + if kwargs.get("return_logprob") and not allow_top: + kwargs["top_logprobs_num"] = 0 return kwargs @staticmethod def _extract_logprobs( - meta_info: Dict[str, Any], num_output_logprobs_so_far: int + meta_info: Dict[str, Any], + num_output_logprobs_so_far: int, + return_tokens_as_token_ids: bool = False, ) -> tuple: """Extract logprobs from SGLang meta_info for new tokens. @@ -272,11 +382,17 @@ def _extract_logprobs( continue position_list = [] for rank_idx, entry in enumerate(position_entries): + tok_id = entry[1] + token_str = ( + f"token_id:{tok_id}" + if return_tokens_as_token_ids + else entry[2] + ) position_list.append( { "rank": rank_idx + 1, - "token_id": entry[1], - "token": entry[2], + "token_id": tok_id, + "token": token_str, "logprob": float(entry[0]), } ) @@ -307,6 +423,12 @@ async def generate( priority = (request.get("routing") or {}).get("priority") logprob_kwargs = self._build_logprob_kwargs(request) + output_options = request.get("output_options", {}) + return_tokens_as_token_ids = bool( + output_options.get("return_tokens_as_token_ids") + ) + user_stop_token_ids = _user_stop_token_ids(request) + lora_path = self._resolve_lora(request) if lora_path: logging.debug(f"Request {context.id()} will use LoRA adapter: {lora_path}") @@ -351,10 +473,20 @@ async def generate( ) if not self.use_sglang_tokenizer: - async for out in self._process_token_stream(decode, context): + async for out in self._process_token_stream( + decode, + context, + return_tokens_as_token_ids, + user_stop_token_ids=user_stop_token_ids, + ): yield out else: - async for out in self._process_text_stream(decode, context): + async for out in self._process_text_stream( + decode, + context, + request=request, + user_stop_token_ids=user_stop_token_ids, + ): yield out else: # Extract image/video URLs for multimodal requests. SGLang's mm_data_processor @@ -385,16 +517,28 @@ async def generate( **self._priority_kwargs(priority), ) if not self.use_sglang_tokenizer: - async for out in self._process_token_stream(agg, context): + async for out in self._process_token_stream( + agg, + context, + return_tokens_as_token_ids, + user_stop_token_ids=user_stop_token_ids, + ): yield out else: - async for out in self._process_text_stream(agg, context): + async for out in self._process_text_stream( + agg, + context, + request=request, + user_stop_token_ids=user_stop_token_ids, + ): yield out async def _process_token_stream( self, stream_source: AsyncGenerator[Dict[str, Any], None], context: Context, + return_tokens_as_token_ids: bool = False, + user_stop_token_ids: set[int] | None = None, ) -> AsyncGenerator[Dict[str, Any], None]: """Process token-based stream output. @@ -438,6 +582,11 @@ async def _process_token_stream( out["finish_reason"] = normalize_finish_reason( finish_reason["type"] ) + stop_reason = _extract_sglang_stop_reason( + finish_reason, user_stop_token_ids + ) + if stop_reason is not None: + out["stop_reason"] = stop_reason # With stream_output=True, output_ids contains only new tokens (disjoint) output_ids = res.get("output_ids", []) @@ -457,7 +606,9 @@ async def _process_token_stream( top_logprobs, next_logprobs_total, ) = self._extract_logprobs( - res["meta_info"], output_logprobs_per_choice.get(output_idx, 0) + res["meta_info"], + output_logprobs_per_choice.get(output_idx, 0), + return_tokens_as_token_ids=return_tokens_as_token_ids, ) output_logprobs_per_choice[output_idx] = next_logprobs_total if log_probs is not None: @@ -493,6 +644,8 @@ async def _process_text_stream( self, stream_source: AsyncGenerator[Dict[str, Any], None], context: Context, + request: Dict[str, Any] | None = None, + user_stop_token_ids: set[int] | None = None, ) -> AsyncGenerator[Dict[str, Any], None]: """Process text-based stream output in OpenAI format. @@ -503,6 +656,7 @@ async def _process_text_stream( Yields: OpenAI-formatted chat completion chunk dicts. """ + request = request or {} # SGLang text chunks are cumulative per choice. Keep independent text # offsets so interleaved n>1 choices do not compute deltas from each # other's previous text. @@ -544,6 +698,9 @@ async def _process_text_stream( "delta": {"role": "assistant", "content": delta}, "finish_reason": finish_reason_type, } + stop_reason = _extract_sglang_stop_reason( + finish_reason, user_stop_token_ids + ) response = { "id": res["meta_info"]["id"], @@ -552,13 +709,20 @@ async def _process_text_stream( "model": self.config.server_args.served_model_name, "object": "chat.completion.chunk", } + response_nvext: dict[str, Any] = {} + if stop_reason is not None and _nvext_extra_field_requested( + request, "stop_reason" + ): + response_nvext["stop_reason"] = stop_reason routed_experts = res["meta_info"].get("routed_experts") if routed_experts is not None: # Base64-encode tensor bytes to match sglang's output format. routed_experts = pybase64.b64encode( routed_experts.numpy().tobytes() ).decode("utf-8") - response["nvext"] = {"routed_experts": routed_experts} + response_nvext["routed_experts"] = routed_experts + if response_nvext: + response["nvext"] = response_nvext if not context.is_stopped(): yield response text_counts_per_choice[index] = next_count diff --git a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py index be8c34a85fe3..aae6a08c8554 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py +++ b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py @@ -10,6 +10,9 @@ from dynamo.sglang.request_handlers.llm.decode_handler import ( DecodeWorkerHandler, _extract_media_urls, + _extract_sglang_stop_reason, + _openai_stop_sampling_params, + _user_stop_token_ids, ) from dynamo.sglang.request_handlers.multimodal.worker_handler import StreamProcessor @@ -44,6 +47,80 @@ def test_extract_media_urls_returns_none_for_missing_or_invalid_items(): ) +@pytest.mark.parametrize( + ("finish_reason", "expected"), + [ + ({"type": "stop", "matched": "END"}, "END"), + ({"type": "stop", "matched": 128001}, 128001), + ({"type": "stop", "matched": [128001, 128009]}, [128001, 128009]), + ({"type": "stop", "matched": True}, None), + ({"type": "stop", "matched": ["END"]}, None), + ({"type": "length"}, None), + (None, None), + ], +) +def test_extract_sglang_stop_reason(finish_reason, expected): + assert _extract_sglang_stop_reason(finish_reason) == expected + + +def test_extract_sglang_stop_reason_filters_hidden_token_ids(): + finish_reason = {"type": "stop", "matched": 128001} + + assert _extract_sglang_stop_reason(finish_reason, {576}) is None + assert _extract_sglang_stop_reason(finish_reason, {128001}) == 128001 + + +def test_extract_sglang_stop_reason_filters_hidden_token_id_arrays(): + finish_reason = {"type": "stop", "matched": [128001, 128009]} + + assert _extract_sglang_stop_reason(finish_reason, {128001}) is None + assert _extract_sglang_stop_reason(finish_reason, {128001, 128009}) == [ + 128001, + 128009, + ] + + +def test_user_stop_token_ids_ignores_hidden_ids(): + assert _user_stop_token_ids( + { + "stop_conditions": { + "stop_token_ids": [576], + "stop_token_ids_hidden": [128001], + } + } + ) == {576} + + +def test_user_stop_token_ids_handles_null_fields(): + assert _user_stop_token_ids({"stop_conditions": {"stop_token_ids": None}}) == set() + assert _user_stop_token_ids({"stop_token_ids": None}) == set() + + +def test_user_stop_token_ids_accepts_stop_token_id_array(): + assert _user_stop_token_ids({"stop": [32, 34]}) == {32, 34} + + +def test_user_stop_token_ids_treats_token_id_display_as_string_stop(): + assert _user_stop_token_ids({"stop": ["token_id:576"]}) == set() + + +def test_openai_stop_sampling_params_preserves_string_stops(): + assert _openai_stop_sampling_params({"stop": "END"}) == {"stop": "END"} + assert _openai_stop_sampling_params({"stop": ["END"]}) == {"stop": ["END"]} + assert _openai_stop_sampling_params({"stop": ["token_id:576"]}) == { + "stop": ["token_id:576"] + } + + +def test_openai_stop_sampling_params_maps_token_id_stop_array(): + assert _openai_stop_sampling_params({"stop": [32, 34]}) == { + "stop_token_ids": [32, 34] + } + assert _openai_stop_sampling_params({"stop_token_ids": [32, 34]}) == { + "stop_token_ids": [32, 34] + } + + def _new_decode_handler(*, use_sglang_tokenizer: bool = False): handler = DecodeWorkerHandler.__new__(DecodeWorkerHandler) handler.use_sglang_tokenizer = use_sglang_tokenizer @@ -88,12 +165,66 @@ def test_build_sampling_params_passes_n_for_sglang_tokenizer_requests(): handler = _new_decode_handler(use_sglang_tokenizer=True) sampling_params = handler._build_sampling_params( - {"temperature": 0.2, "top_p": 0.9, "n": 2, "max_tokens": 8} + { + "temperature": 0.2, + "top_p": 0.9, + "n": 2, + "max_tokens": 8, + "stop": [32, 34], + } ) assert sampling_params["n"] == 2 assert sampling_params["temperature"] == 0.2 assert sampling_params["max_new_tokens"] == 8 + assert sampling_params["stop_token_ids"] == [32, 34] + + +def test_build_logprob_kwargs_allows_chosen_token_logprobs(monkeypatch): + monkeypatch.delenv("DYN_SGL_ALLOW_TOP_LOGPROBS", raising=False) + + kwargs = DecodeWorkerHandler._build_logprob_kwargs( + {"output_options": {"logprobs": 0}} + ) + + assert kwargs == {"return_logprob": True, "top_logprobs_num": 0} + + +def test_build_logprob_kwargs_rejects_top_logprobs_by_default(monkeypatch): + monkeypatch.delenv("DYN_SGL_ALLOW_TOP_LOGPROBS", raising=False) + + with pytest.raises(ValueError, match="does not currently support logprobs >= 1"): + DecodeWorkerHandler._build_logprob_kwargs({"output_options": {"logprobs": 1}}) + + +def test_build_logprob_kwargs_allows_top_logprobs_with_escape_hatch(monkeypatch): + monkeypatch.setenv("DYN_SGL_ALLOW_TOP_LOGPROBS", "1") + + kwargs = DecodeWorkerHandler._build_logprob_kwargs( + {"output_options": {"logprobs": 2}} + ) + + assert kwargs == {"return_logprob": True, "top_logprobs_num": 2} + + +def test_extract_logprobs_formats_top_tokens_as_token_ids(): + log_probs, top_logprobs, total = DecodeWorkerHandler._extract_logprobs( + { + "output_token_logprobs": [(-0.1, 101, "a")], + "output_top_logprobs": [[(-0.1, 101, "a"), (-0.2, 102, "b")]], + }, + 0, + return_tokens_as_token_ids=True, + ) + + assert log_probs == [-0.1] + assert top_logprobs == [ + [ + {"rank": 1, "token_id": 101, "token": "token_id:101", "logprob": -0.1}, + {"rank": 2, "token_id": 102, "token": "token_id:102", "logprob": -0.2}, + ] + ] + assert total == 1 @pytest.mark.asyncio @@ -189,6 +320,88 @@ async def test_process_text_stream_tracks_delta_per_choice_index(): ] +@pytest.mark.asyncio +async def test_process_text_stream_stop_reason_uses_response_nvext(): + handler = _new_decode_handler() + + chunks = await _collect( + handler._process_text_stream( + _stream( + [ + { + "index": 0, + "text": "Hello", + "meta_info": { + "id": "request-1", + "finish_reason": {"type": "stop", "matched": "END"}, + }, + } + ] + ), + _Context(), + request={"nvext": {"extra_fields": ["stop_reason"]}}, + ) + ) + + assert "stop_reason" not in chunks[0]["choices"][0] + assert chunks[0]["nvext"]["stop_reason"] == "END" + + +@pytest.mark.asyncio +async def test_process_text_stream_stop_reason_requires_nvext_extra_field(): + handler = _new_decode_handler() + + chunks = await _collect( + handler._process_text_stream( + _stream( + [ + { + "index": 0, + "text": "Hello", + "meta_info": { + "id": "request-1", + "finish_reason": {"type": "stop", "matched": "END"}, + }, + } + ] + ), + _Context(), + ) + ) + + assert "stop_reason" not in chunks[0]["choices"][0] + assert "nvext" not in chunks[0] + + +@pytest.mark.asyncio +async def test_process_token_stream_suppresses_hidden_stop_token_reason(): + handler = _new_decode_handler() + + chunks = await _collect( + handler._process_token_stream( + _stream( + [ + { + "index": 0, + "output_ids": [128001], + "meta_info": { + "id": "request-1", + "finish_reason": {"type": "stop", "matched": 128001}, + "prompt_tokens": 1, + "completion_tokens": 1, + "cached_tokens": None, + }, + } + ] + ), + _Context(), + user_stop_token_ids={576}, + ) + ) + + assert "stop_reason" not in chunks[0] + + @pytest.mark.asyncio async def test_multimodal_stream_keeps_reading_after_one_choice_finishes(): chunks = await _collect( diff --git a/docs/components/frontend/nvext.md b/docs/components/frontend/nvext.md index 68b311f3595d..65144aad632c 100644 --- a/docs/components/frontend/nvext.md +++ b/docs/components/frontend/nvext.md @@ -35,13 +35,24 @@ Include `nvext` as a top-level field alongside standard OpenAI-compatible fields | `backend_instance_id` | `u64` | `None` | Router | Routes the request to a specific backend instance. | | `token_data` | `u32[]` | `None` | Preprocessor | Pre-tokenized prompt tokens. When provided with `backend_instance_id`, tokenization is skipped. | | `max_thinking_tokens` | `u32` | `None` | Backend | Maximum thinking tokens allowed (passed through to backends). | -| `extra_fields` | `string[]` | `None` | Response builder | Fields to include in the response `nvext`. Supported: `"worker_id"`, `"timing"`, `"routed_experts"`. | +| `extra_fields` | `string[]` | `None` | Response builder | Fields to include in the response `nvext`. Supported: `"worker_id"`, `"timing"`, `"routed_experts"`, `"engine_data"`, `"stop_reason"`. | | `prefill_worker_id` | `u64` | `None` | Router | Routes the request to a specific prefill worker (disaggregated serving). | | `decode_worker_id` | `u64` | `None` | Router | Routes the request to a specific decode worker (disaggregated serving). | | `agent_context` | object | `None` | Preprocessor | Passive session and trajectory identity for agent traces. See [Agent Context](#agent-context) below and [Agent Tracing](../../agents/agent-tracing.md). | | `agent_hints` | object | `None` | Router | Per-request hints for scheduling and load balancing. See [Agent Hints](#agent-hints). | | `session_control` | object | `None` | Router | Session lifecycle and sticky routing for subagent KV isolation. See [Session Control](#session-control). | +Related root-level Dynamo output option: + +| Field | Type | Default | Consumed By | Description | +|-------|------|---------|-------------|-------------| +| `return_tokens_as_token_ids` | `bool` | `false` | Response builder | Formats logprob token strings as `token_id:` instead of decoded text. | + +`return_tokens_as_token_ids` only changes returned logprob token display. To stop on +token IDs, pass integer IDs in the normal `stop` array, for example +`"stop": [576]`. Strings such as `"token_id:576"` remain literal string stop +sequences and are not parsed as token IDs. + ### Header Overrides Routing fields can also be set via HTTP headers, which take priority over `nvext` values: @@ -195,6 +206,8 @@ When the client requests response metadata via `extra_fields`, the response incl | `worker_id` | `extra_fields: ["worker_id"]` | Prefill/decode worker IDs and data parallel ranks that processed the request. | | `timing` | `extra_fields: ["timing"]` | Per-request timing information (TTFT, ITL, queue time, etc.). | | `routed_experts` | `extra_fields: ["routed_experts"]` | Routed expert capture payload returned by SGLang-backed requests. | +| `engine_data` | `extra_fields: ["engine_data"]` | Opaque backend-provided engine metadata. | +| `stop_reason` | `extra_fields: ["stop_reason"]` | Backend-specific matched stop condition, returned under `nvext` because it is not part of the OpenAI completions schema. Dynamo currently serves this as a response-level field for single-choice requests; supporting `n > 1` will require an indexed per-choice shape. | | `token_ids` | Automatic (GAIE Stage 1) | Tokenized prompt for reuse in Stage 2 query-only mode. | ### Example response `nvext` diff --git a/lib/llm/src/audit/stream.rs b/lib/llm/src/audit/stream.rs index 2bb1d7f64f3e..b3dd5267504a 100644 --- a/lib/llm/src/audit/stream.rs +++ b/lib/llm/src/audit/stream.rs @@ -221,7 +221,6 @@ pub fn final_response_to_one_chunk_stream( index: idx as u32, delta, finish_reason: ch.finish_reason, - stop_reason: ch.stop_reason.clone(), logprobs: ch.logprobs.clone(), }; choices.push(choice); @@ -278,7 +277,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }; @@ -319,7 +317,6 @@ mod tests { reasoning_content: None, }, finish_reason: Some(FinishReason::Stop), - stop_reason: None, logprobs: None, }; @@ -461,7 +458,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, } }], diff --git a/lib/llm/src/backend.rs b/lib/llm/src/backend.rs index f72ee007bc56..2df3cdf374c9 100644 --- a/lib/llm/src/backend.rs +++ b/lib/llm/src/backend.rs @@ -262,6 +262,13 @@ impl // System EOS token - no stop_reason (user didn't request this stop) (Some(FinishReason::Stop), None) } + Some(StopTrigger::UserStopTokenDetected(token_id)) => { + // User-provided token stop (hidden from output) + ( + Some(FinishReason::Stop), + Some(StopReason::Int((*token_id).into())), + ) + } Some(StopTrigger::HiddenStopSequenceDetected(seq)) => { // User-provided stop sequence (hidden from output) ( @@ -322,7 +329,7 @@ impl // which we don't want to propagate to `data.finish_reason`. if finish_reason.is_some() { data.finish_reason = finish_reason; - data.stop_reason = stop_reason; + data.stop_reason = stop_reason.or(data.stop_reason); } data.text = text; data.tokens = Some(tokens); @@ -435,6 +442,10 @@ pub struct Decoder { // minimum number of tokens have been generated hidden_stop_ids: HashSet, + // user-provided token stop IDs, kept separate from system/EOS stop IDs so + // stop_reason can report user-triggered token stops without reporting EOS. + user_stop_ids: HashSet, + // text sequences that if found in the response will trigger a stop condition after the // minimum number of tokens have been generated (excluded from output) hidden_stop_sequences: Vec, @@ -461,6 +472,7 @@ pub struct Decoder { pub enum StopTrigger { MaxTokensLimit, HiddenStopTokenDetected(TokenIdType), + UserStopTokenDetected(TokenIdType), HiddenStopSequenceDetected(String), VisibleStopSequenceDetected(String), } @@ -501,12 +513,19 @@ impl Decoder { include_stop_str_in_output: bool, tracker: Option>, ) -> Self { - let hidden_stop_ids: HashSet = stop_condition + let user_stop_ids: HashSet = stop_condition + .stop_token_ids + .unwrap_or_default() + .iter() + .copied() + .collect(); + let system_stop_ids: HashSet = stop_condition .stop_token_ids_hidden .unwrap_or_default() .iter() .copied() .collect(); + let hidden_stop_ids = user_stop_ids.union(&system_stop_ids).copied().collect(); // Categorize stop sequences based on include_stop_str_in_output: // - When true: user-provided stop sequences go to visible (included in output) @@ -529,6 +548,7 @@ impl Decoder { decode_stream, tracker, hidden_stop_ids, + user_stop_ids, hidden_stop_sequences, visible_stop_sequences, min_tokens: stop_condition.min_tokens.unwrap_or(0), @@ -565,12 +585,15 @@ impl Decoder { return Ok(StepResult::ok(token)); } - // check for hidden stop tokens - eos takes precedence + // Check token stops. User-provided token IDs take precedence over + // system/EOS IDs so stop_reason only reports stops the caller requested. if self.hidden_stop_ids.contains(&token_id) { - return Ok(StepResult::with_stop_trigger( - None, - StopTrigger::HiddenStopTokenDetected(token_id), - )); + let trigger = if self.user_stop_ids.contains(&token_id) { + StopTrigger::UserStopTokenDetected(token_id) + } else { + StopTrigger::HiddenStopTokenDetected(token_id) + }; + return Ok(StepResult::with_stop_trigger(None, trigger)); } // check stop sequences - the jail will always hold at least the largest stop sequence diff --git a/lib/llm/src/engines.rs b/lib/llm/src/engines.rs index 2b07a99cb161..87530d3d2b2d 100644 --- a/lib/llm/src/engines.rs +++ b/lib/llm/src/engines.rs @@ -183,7 +183,7 @@ impl break; } tokio::time::sleep(*TOKEN_ECHO_DELAY).await; - let response = deltas.create_choice(0, Some(c.to_string()), None, None, None); + let response = deltas.create_choice(0, Some(c.to_string()), None, None); yield Annotated { id: Some(id.to_string()), data: Some(response), @@ -200,7 +200,6 @@ impl None, Some(dynamo_protocols::types::FinishReason::Stop), None, - None, ); yield Annotated { id: Some(id.to_string()), @@ -254,12 +253,13 @@ impl for c in prompt.chars() { // we are returning characters not tokens, so there will be some postprocessing overhead tokio::time::sleep(*TOKEN_ECHO_DELAY).await; - let response = deltas.create_choice(0, Some(c.to_string()), None, None, None); + let response = deltas.create_choice(0, Some(c.to_string()), None, None); yield Annotated{ id: Some(id.to_string()), data: Some(response), event: None, comment: None, error: None }; id += 1; } - let response = deltas.create_choice(0, None, Some(dynamo_protocols::types::FinishReason::Stop), None, None); + let response = + deltas.create_choice(0, None, Some(dynamo_protocols::types::FinishReason::Stop), None); yield Annotated { id: Some(id.to_string()), data: Some(response), event: None, comment: None, error: None }; }; @@ -291,7 +291,12 @@ impl yield Annotated{ id: Some(id.to_string()), data: Some(response), event: None, comment: None, error: None }; id += 1; } - let response = deltas.create_choice(0, None, Some(dynamo_protocols::types::CompletionFinishReason::Stop), None); + let response = deltas.create_choice( + 0, + None, + Some(dynamo_protocols::types::CompletionFinishReason::Stop), + None, + ); yield Annotated { id: Some(id.to_string()), data: Some(response), event: None, comment: None, error: None }; }; diff --git a/lib/llm/src/entrypoint/input/text.rs b/lib/llm/src/entrypoint/input/text.rs index 1c0138fd34b3..5821176496a9 100644 --- a/lib/llm/src/entrypoint/input/text.rs +++ b/lib/llm/src/entrypoint/input/text.rs @@ -115,6 +115,7 @@ async fn main_loop( nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; diff --git a/lib/llm/src/grpc/service/openai.rs b/lib/llm/src/grpc/service/openai.rs index f1555e98a56f..6156ecbc83e4 100644 --- a/lib/llm/src/grpc/service/openai.rs +++ b/lib/llm/src/grpc/service/openai.rs @@ -344,6 +344,7 @@ impl TryFrom for NvCreateCompletionRequest { common: Default::default(), nvext: None, metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }) } diff --git a/lib/llm/src/http/service/openai.rs b/lib/llm/src/http/service/openai.rs index 91b2c0d6f982..f0a76b5f5ade 100644 --- a/lib/llm/src/http/service/openai.rs +++ b/lib/llm/src/http/service/openai.rs @@ -2874,6 +2874,7 @@ mod tests { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_chat_completion_required_fields(&request); @@ -2906,6 +2907,7 @@ mod tests { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_chat_completion_required_fields(&request); @@ -2943,6 +2945,7 @@ mod tests { common: Default::default(), nvext: None, metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; @@ -2967,6 +2970,7 @@ mod tests { common: Default::default(), nvext: None, metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_completion_fields_generic(&request); @@ -2990,6 +2994,7 @@ mod tests { common: Default::default(), nvext: None, metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_completion_fields_generic(&request); @@ -3013,6 +3018,7 @@ mod tests { common: Default::default(), nvext: None, metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_completion_fields_generic(&request); @@ -3038,6 +3044,7 @@ mod tests { .unwrap(), nvext: None, metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_completion_fields_generic(&request); @@ -3061,6 +3068,7 @@ mod tests { common: Default::default(), nvext: None, metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_completion_fields_generic(&request); @@ -3092,6 +3100,7 @@ mod tests { "session": {"id": "session-1", "timestamp": 1640995200} }) .into(), + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; @@ -3122,6 +3131,7 @@ mod tests { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; @@ -3152,6 +3162,7 @@ mod tests { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_chat_completion_fields_generic(&request); @@ -3181,6 +3192,7 @@ mod tests { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_chat_completion_fields_generic(&request); @@ -3210,6 +3222,7 @@ mod tests { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_chat_completion_fields_generic(&request); @@ -3241,6 +3254,7 @@ mod tests { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_chat_completion_fields_generic(&request); @@ -3270,6 +3284,7 @@ mod tests { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_chat_completion_fields_generic(&request); @@ -3745,7 +3760,6 @@ mod tests { reasoning_content: reasoning.map(|s| s.to_string()), }, finish_reason: finish, - stop_reason: None, logprobs: None, } } @@ -3777,7 +3791,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, } } @@ -3879,7 +3892,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }; @@ -3951,7 +3963,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }; @@ -3988,7 +3999,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }; diff --git a/lib/llm/src/perf/logprobs.rs b/lib/llm/src/perf/logprobs.rs index 02c93cdf1ecf..6873ad49a27c 100644 --- a/lib/llm/src/perf/logprobs.rs +++ b/lib/llm/src/perf/logprobs.rs @@ -963,7 +963,6 @@ mod tests { reasoning_content: None, }, finish_reason: Some(FinishReason::Stop), - stop_reason: None, logprobs: Some(ChatChoiceLogprobs { content: Some(token_logprobs), refusal: None, @@ -999,7 +998,6 @@ mod tests { reasoning_content: None, }, finish_reason: Some(FinishReason::Stop), - stop_reason: None, logprobs: Some(ChatChoiceLogprobs { content: Some(token_logprobs), refusal: None, @@ -1353,7 +1351,6 @@ mod tests { reasoning_content: None, }, finish_reason: Some(FinishReason::Stop), - stop_reason: None, logprobs: None, // No logprobs }], created: 1234567890, diff --git a/lib/llm/src/protocols/anthropic/stream_converter.rs b/lib/llm/src/protocols/anthropic/stream_converter.rs index 426ff3e7164b..bb8220445a79 100644 --- a/lib/llm/src/protocols/anthropic/stream_converter.rs +++ b/lib/llm/src/protocols/anthropic/stream_converter.rs @@ -756,7 +756,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }], created: 0, @@ -799,7 +798,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }], created: 0, @@ -946,7 +944,6 @@ mod tests { reasoning_content: Some(text.into()), }, finish_reason: None, - stop_reason: None, logprobs: None, }], created: 0, diff --git a/lib/llm/src/protocols/anthropic/types.rs b/lib/llm/src/protocols/anthropic/types.rs index 5214ee10b0a7..a199713ec90f 100644 --- a/lib/llm/src/protocols/anthropic/types.rs +++ b/lib/llm/src/protocols/anthropic/types.rs @@ -140,6 +140,7 @@ impl TryFrom for NvCreateChatCompletionRequest { None }, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }) } @@ -804,7 +805,6 @@ mod tests { reasoning_content: None, }, finish_reason: Some(dynamo_protocols::types::FinishReason::Stop), - stop_reason: None, logprobs: None, }], created: 1726000000, diff --git a/lib/llm/src/protocols/common.rs b/lib/llm/src/protocols/common.rs index aa3f147d29de..257bcc170bd2 100644 --- a/lib/llm/src/protocols/common.rs +++ b/lib/llm/src/protocols/common.rs @@ -18,6 +18,7 @@ use derive_builder::Builder; use serde::{Deserialize, Serialize}; use super::TokenIdType; +use dynamo_protocols::types::StopReason; /// Maximum nesting depth allowed in guided_grammar EBNF strings. const MAX_GRAMMAR_NESTING_DEPTH: usize = 500; @@ -239,6 +240,10 @@ pub struct StopConditions { /// List of tokens that stop the generation when they are /// generated. The returned output will NOT contain the stop tokens. + pub stop_token_ids: Option>, + + /// List of hidden/system tokens that stop generation when they are + /// generated. The returned output will NOT contain the stop tokens. pub stop_token_ids_hidden: Option>, /// The minimum number of tokens to generate @@ -259,6 +264,7 @@ impl StopConditions { pub fn apply_ignore_eos(&mut self) { if self.ignore_eos.unwrap_or(false) { self.stop = None; + self.stop_token_ids = None; self.stop_token_ids_hidden = None; } } @@ -514,6 +520,10 @@ pub struct OutputOptions { /// the tokenizer. This is useful for inspecting the behavior of prompt /// templates that are applied during the backend preprocessing. pub formatted_prompt: Option, + + /// When true, logprob token fields are returned as "token_id:" + /// instead of decoded text. + pub return_tokens_as_token_ids: Option, } // Struct for log probability information @@ -604,6 +614,10 @@ pub struct Delta { pub finish_reason: Option, + /// The stop string or token that triggered the stop condition. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub stop_reason: Option, + // new token_ids pub token_ids: Option>, diff --git a/lib/llm/src/protocols/openai.rs b/lib/llm/src/protocols/openai.rs index 68e9fd80e77d..0234ccc6f59e 100644 --- a/lib/llm/src/protocols/openai.rs +++ b/lib/llm/src/protocols/openai.rs @@ -63,6 +63,10 @@ pub(crate) trait OpenAIStopConditionsProvider { fn get_stop(&self) -> Option>; + fn get_stop_token_ids(&self) -> Option> { + None + } + fn nvext(&self) -> Option<&nvext::NvExt>; /// Get ignore_eos from CommonExt if the type supports it. @@ -91,6 +95,10 @@ pub(crate) trait OpenAIOutputOptionsProvider { fn get_skip_special_tokens(&self) -> Option; fn get_formatted_prompt(&self) -> Option; + + fn get_return_tokens_as_token_ids(&self) -> Option { + None + } } impl SamplingOptionsProvider for T { @@ -176,6 +184,7 @@ impl StopConditionsProvider for T { let max_tokens = self.get_max_tokens(); let min_tokens = self.get_min_tokens(); let stop = self.get_stop(); + let stop_token_ids = self.get_stop_token_ids(); let max_thinking_tokens = self.get_max_thinking_tokens(); if let Some(stop) = &stop @@ -183,6 +192,11 @@ impl StopConditionsProvider for T { { anyhow::bail!("stop conditions must be less than 4") } + if let Some(stop_token_ids) = &stop_token_ids + && stop_token_ids.len() > 4 + { + anyhow::bail!("stop token IDs must be less than 4") + } // Use the trait method to get ignore_eos, which handles precedence let ignore_eos = self.get_ignore_eos(); @@ -191,6 +205,7 @@ impl StopConditionsProvider for T { max_tokens, min_tokens, stop, + stop_token_ids, stop_token_ids_hidden: None, ignore_eos, max_thinking_tokens, @@ -204,12 +219,14 @@ impl OutputOptionsProvider for T { let prompt_logprobs = self.get_prompt_logprobs(); let skip_special_tokens = self.get_skip_special_tokens(); let formatted_prompt = self.get_formatted_prompt(); + let return_tokens_as_token_ids = self.get_return_tokens_as_token_ids(); Ok(common::OutputOptions { logprobs, prompt_logprobs, skip_special_tokens, formatted_prompt, + return_tokens_as_token_ids, }) } } @@ -231,14 +248,23 @@ pub(crate) fn convert_backend_top_logprobs( selected_token: &str, selected_token_id: TokenIdType, selected_logprob: f32, + return_tokens_as_token_ids: bool, ) -> Vec { let mut found_selected = false; let mut result: Vec = top_lps .iter() .map(|top_lp| { - let tok = top_lp.token.clone().unwrap_or_default(); + let tok = if return_tokens_as_token_ids { + format!("token_id:{}", top_lp.token_id) + } else { + top_lp.token.clone().unwrap_or_default() + }; found_selected = found_selected || top_lp.token_id == selected_token_id; - let bytes = top_lp.bytes.clone().or_else(|| token_to_utf8_bytes(&tok)); + let bytes = if return_tokens_as_token_ids { + token_to_utf8_bytes(&tok) + } else { + top_lp.bytes.clone().or_else(|| token_to_utf8_bytes(&tok)) + }; dynamo_protocols::types::TopLogprobs { token: tok, logprob: top_lp.logprob as f32, @@ -248,10 +274,15 @@ pub(crate) fn convert_backend_top_logprobs( .collect(); if !found_selected { + let token = if return_tokens_as_token_ids { + format!("token_id:{}", selected_token_id) + } else { + selected_token.to_string() + }; result.push(dynamo_protocols::types::TopLogprobs { - token: selected_token.to_string(), + bytes: token_to_utf8_bytes(&token), + token, logprob: selected_logprob, - bytes: token_to_utf8_bytes(selected_token), }); } result diff --git a/lib/llm/src/protocols/openai/chat_completions.rs b/lib/llm/src/protocols/openai/chat_completions.rs index 8a77038d5834..c88cfed8d9b4 100644 --- a/lib/llm/src/protocols/openai/chat_completions.rs +++ b/lib/llm/src/protocols/openai/chat_completions.rs @@ -59,6 +59,11 @@ pub struct NvCreateChatCompletionRequest { #[serde(default, skip_serializing_if = "Option::is_none")] pub media_io_kwargs: Option, + /// When true, logprob token fields are returned as "token_id:" instead + /// of decoded text. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub return_tokens_as_token_ids: Option, + /// Catch-all for unsupported fields - checked during validation #[serde(flatten, default, skip_serializing)] pub unsupported_fields: std::collections::HashMap, @@ -285,10 +290,11 @@ impl OpenAIStopConditionsProvider for NvCreateChatCompletionRequest { /// * `Some(Vec)` if stop conditions are set. /// * `None` if no stop conditions are defined. fn get_stop(&self) -> Option> { - self.inner.stop.as_ref().map(|stop| match stop { - dynamo_protocols::types::Stop::String(s) => vec![s.clone()], - dynamo_protocols::types::Stop::StringArray(arr) => arr.clone(), - }) + self.inner.stop.as_ref().and_then(|stop| stop.strings()) + } + + fn get_stop_token_ids(&self) -> Option> { + self.inner.stop.as_ref().and_then(|stop| stop.token_ids()) } /// Returns a reference to the optional `NvExt` extension, if available. @@ -330,6 +336,10 @@ impl OpenAIOutputOptionsProvider for NvCreateChatCompletionRequest { fn get_formatted_prompt(&self) -> Option { None } + + fn get_return_tokens_as_token_ids(&self) -> Option { + self.return_tokens_as_token_ids + } } /// Implements `ValidateRequest` for `NvCreateChatCompletionRequest`, @@ -381,7 +391,8 @@ impl ValidateRequest for NvCreateChatCompletionRequest { #[cfg(test)] mod tests { use super::*; - use crate::protocols::common::OutputOptionsProvider; + use crate::engines::ValidateRequest; + use crate::protocols::common::{OutputOptionsProvider, StopConditionsProvider}; use serde_json::json; #[test] @@ -426,4 +437,87 @@ mod tests { assert_eq!(output_options.skip_special_tokens, Some(skip_value)); } } + + #[test] + fn test_stop_contract() { + let one_stop = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}], + "stop": " The" + }); + let request: NvCreateChatCompletionRequest = + serde_json::from_value(one_stop).expect("Failed to deserialize request"); + assert_eq!(request.get_stop(), Some(vec![" The".to_string()])); + assert_eq!(request.get_stop_token_ids(), None); + + let many_stops = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}], + "stop": ["A", "B"] + }); + let request: NvCreateChatCompletionRequest = + serde_json::from_value(many_stops).expect("Failed to deserialize request"); + assert_eq!( + request.get_stop(), + Some(vec!["A".to_string(), "B".to_string()]) + ); + assert_eq!(request.get_stop_token_ids(), None); + + let token_id_stops = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}], + "stop": [32, 34] + }); + let request: NvCreateChatCompletionRequest = + serde_json::from_value(token_id_stops).expect("Failed to deserialize request"); + assert_eq!(request.get_stop(), None); + assert_eq!(request.get_stop_token_ids(), Some(vec![32, 34])); + + let stop_conditions = request + .extract_stop_conditions() + .expect("extract stop conditions"); + assert_eq!(stop_conditions.stop, None); + assert_eq!(stop_conditions.stop_token_ids, Some(vec![32, 34])); + + let token_id_display_string_stop = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}], + "stop": "token_id:576" + }); + let request: NvCreateChatCompletionRequest = + serde_json::from_value(token_id_display_string_stop) + .expect("Failed to deserialize request"); + assert_eq!(request.get_stop(), Some(vec!["token_id:576".to_string()])); + assert_eq!(request.get_stop_token_ids(), None); + + let token_id_display_string_array_stop = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}], + "stop": ["token_id:576"] + }); + let request: NvCreateChatCompletionRequest = + serde_json::from_value(token_id_display_string_array_stop) + .expect("Failed to deserialize request"); + assert_eq!(request.get_stop(), Some(vec!["token_id:576".to_string()])); + assert_eq!(request.get_stop_token_ids(), None); + + let scalar_token_id_stop = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}], + "stop": 576 + }); + let result: Result = + serde_json::from_value(scalar_token_id_stop); + assert!(result.is_err()); + + let unsupported_stop_token_ids = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}], + "stop_token_ids": [576] + }); + let request: NvCreateChatCompletionRequest = + serde_json::from_value(unsupported_stop_token_ids) + .expect("Failed to deserialize request"); + assert!(ValidateRequest::validate(&request).is_err()); + } } diff --git a/lib/llm/src/protocols/openai/chat_completions/aggregator.rs b/lib/llm/src/protocols/openai/chat_completions/aggregator.rs index 5def709448d8..0849b557fe21 100644 --- a/lib/llm/src/protocols/openai/chat_completions/aggregator.rs +++ b/lib/llm/src/protocols/openai/chat_completions/aggregator.rs @@ -11,10 +11,10 @@ use crate::protocols::{ Annotated, codec::{Message, SseCodecError}, convert_sse_stream, - openai::ParsingOptions, + openai::{ParsingOptions, nvext::merge_response_nvext}, }; -use dynamo_protocols::types::{ChatCompletionMessageContent, StopReason}; +use dynamo_protocols::types::ChatCompletionMessageContent; use dynamo_runtime::engine::DataStream; /// Aggregates a stream of [`NvCreateChatCompletionStreamResponse`]s into a single @@ -52,8 +52,6 @@ struct DeltaChoice { role: Option, /// The reason the completion was finished (if applicable). finish_reason: Option, - /// The stop string or token that triggered the stop condition. - stop_reason: Option, /// Optional log probabilities for the chat choice. logprobs: Option, // Tool-call chunks accumulated in the order they arrived from the stream, @@ -238,10 +236,7 @@ impl DeltaAggregator { aggregator.system_fingerprint = Some(system_fingerprint); } - // Aggregate nvext field (take the last non-None value) - if delta.nvext.is_some() { - aggregator.nvext = delta.nvext; - } + merge_response_nvext(&mut aggregator.nvext, delta.nvext); // Aggregate choices incrementally. for choice in delta.inner.choices { @@ -254,7 +249,6 @@ impl DeltaAggregator { text: "".to_string(), role: choice.delta.role, finish_reason: None, - stop_reason: None, logprobs: None, tool_call_chunks: BTreeMap::new(), tool_calls: None, @@ -307,11 +301,6 @@ impl DeltaAggregator { state_choice.finish_reason = Some(finish_reason); } - // Update stop reason if provided. - if let Some(stop_reason) = choice.stop_reason { - state_choice.stop_reason = Some(stop_reason); - } - // Update logprobs if let Some(logprobs) = &choice.logprobs { let state_lps = state_choice.logprobs.get_or_insert( @@ -466,7 +455,6 @@ impl From for dynamo_protocols::types::ChatChoice { }, index: delta.index, finish_reason, - stop_reason: delta.stop_reason, logprobs: delta.logprobs, } } @@ -581,7 +569,6 @@ mod tests { index, delta, finish_reason, - stop_reason: None, logprobs, }; @@ -631,7 +618,6 @@ mod tests { index, delta, finish_reason, - stop_reason: None, logprobs: None, }; let data = NvCreateChatCompletionStreamResponse { @@ -1041,6 +1027,43 @@ mod tests { assert_eq!(choice.message.role, dynamo_protocols::types::Role::User); } + #[tokio::test] + async fn test_multiple_deltas_merge_nvext_fields() { + let mut annotated_delta1 = create_test_delta( + 0, + "Hello", + Some(dynamo_protocols::types::Role::Assistant), + None, + None, + None, + ); + annotated_delta1.data.as_mut().expect("delta data").nvext = + Some(serde_json::json!({ "engine_data": { "trace_id": "abc" } })); + let mut annotated_delta2 = create_test_delta( + 0, + " world", + None, + Some(dynamo_protocols::types::FinishReason::Stop), + None, + None, + ); + annotated_delta2.data.as_mut().expect("delta data").nvext = + Some(serde_json::json!({ "stop_reason": 128001 })); + + let stream = Box::pin(stream::iter(vec![annotated_delta1, annotated_delta2])); + let response = DeltaAggregator::apply(stream, ParsingOptions::default()) + .await + .expect("aggregate stream"); + + assert_eq!( + response.nvext, + Some(serde_json::json!({ + "engine_data": { "trace_id": "abc" }, + "stop_reason": 128001, + })) + ); + } + #[allow(deprecated)] #[tokio::test] async fn test_multiple_choices() { @@ -1068,7 +1091,6 @@ mod tests { reasoning_content: None, }, finish_reason: Some(dynamo_protocols::types::FinishReason::Stop), - stop_reason: None, logprobs: None, }, dynamo_protocols::types::ChatChoiceStream { @@ -1084,7 +1106,6 @@ mod tests { reasoning_content: None, }, finish_reason: Some(dynamo_protocols::types::FinishReason::Stop), - stop_reason: None, logprobs: None, }, ], @@ -1557,7 +1578,6 @@ mod tests { text: String::new(), role: Some(dynamo_protocols::types::Role::Assistant), finish_reason: Some(dynamo_protocols::types::FinishReason::Stop), - stop_reason: None, logprobs: None, tool_call_chunks: BTreeMap::new(), tool_calls: None, diff --git a/lib/llm/src/protocols/openai/chat_completions/delta.rs b/lib/llm/src/protocols/openai/chat_completions/delta.rs index 4a3dfae2d405..6f42b2f54556 100644 --- a/lib/llm/src/protocols/openai/chat_completions/delta.rs +++ b/lib/llm/src/protocols/openai/chat_completions/delta.rs @@ -69,6 +69,7 @@ impl NvCreateChatCompletionRequest { enable_logprobs: self.inner.logprobs.unwrap_or(false) || self.inner.top_logprobs.unwrap_or(0) > 0, response_fields, + return_tokens_as_token_ids: self.return_tokens_as_token_ids.unwrap_or(false), runtime_config: ModelRuntimeConfig::default(), }; @@ -87,6 +88,8 @@ pub struct DeltaGeneratorOptions { pub enable_logprobs: bool, /// Determines which nvext response fields may be emitted for this request. pub response_fields: NvExtResponseFieldSelection, + /// When true, logprob token fields use "token_id:" format instead of decoded text. + pub return_tokens_as_token_ids: bool, pub runtime_config: ModelRuntimeConfig, } @@ -198,16 +201,23 @@ impl DeltaGenerator { .map(|(_, lp)| lp as f32) .collect::>(); + let return_as_ids = self.options.return_tokens_as_token_ids; let content = top_logprobs.map(|top_logprobs| { toks.iter() .zip(tok_lps) .zip(top_logprobs) .map(|(((t, tid), lp), top_lps)| { - let converted = convert_backend_top_logprobs(&top_lps, t, *tid, lp); + let token_str = if return_as_ids { + format!("token_id:{}", tid) + } else { + t.clone() + }; + let converted = + convert_backend_top_logprobs(&top_lps, t, *tid, lp, return_as_ids); dynamo_protocols::types::ChatCompletionTokenLogprob { - token: t.clone(), + token: token_str.clone(), logprob: lp, - bytes: token_to_utf8_bytes(t), + bytes: token_to_utf8_bytes(&token_str), top_logprobs: converted, } }) @@ -227,18 +237,15 @@ impl DeltaGenerator { /// * `text` - The text content for the response. /// * `finish_reason` - The reason why the response finished (e.g., stop, length, etc.). /// * `logprobs` - Optional log probabilities of the generated tokens. - /// * `stop_reason` - Optional stop string or token that triggered the stop. /// /// # Returns /// * An [`dynamo_protocols::types::CreateChatCompletionStreamResponse`] instance representing the choice. - #[allow(deprecated)] pub fn create_choice( &mut self, index: u32, text: Option, finish_reason: Option, logprobs: Option, - stop_reason: Option, ) -> NvCreateChatCompletionStreamResponse { let delta = dynamo_protocols::types::ChatCompletionStreamResponseDelta { content: text.map(dynamo_protocols::types::ChatCompletionMessageContent::Text), @@ -257,7 +264,6 @@ impl DeltaGenerator { index, delta, finish_reason, - stop_reason, logprobs, }; @@ -391,16 +397,11 @@ impl crate::protocols::openai::DeltaGeneratorExt None, }; + let stop_reason = delta.stop_reason.clone(); // Create the streaming response. let index = delta.index.unwrap_or(0); - let mut stream_response = self.create_choice( - index, - delta.text, - finish_reason, - logprobs, - delta.stop_reason, - ); + let mut stream_response = self.create_choice(index, delta.text, finish_reason, logprobs); // Record finish for timing/ITL accounting even when timing is not returned to the client. // Kept at call site because it's a side effect on the tracker — not a gating decision. @@ -419,6 +420,7 @@ impl crate::protocols::openai::DeltaGeneratorExt>, finish_reason: Option, - stop_reason: Option, logprobs: Option, ) -> ChatChoiceStream { #[allow(deprecated)] @@ -140,7 +139,6 @@ fn create_choice_stream( reasoning_content: None, }, finish_reason, - stop_reason, logprobs, } } @@ -228,7 +226,6 @@ impl ChoiceJailState { &prefix, None, choice.finish_reason, - None, choice.logprobs.clone(), ); emissions.push(ChoiceEmission::PassThrough(prefix_choice)); @@ -279,7 +276,6 @@ impl ChoiceJailState { trailing_part, None, choice.finish_reason, - None, choice.logprobs.clone(), ); emissions.push(ChoiceEmission::Trailing(trailing_choice)); @@ -310,7 +306,6 @@ impl ChoiceJailState { &prefix, None, choice.finish_reason, - None, choice.logprobs.clone(), ); emissions.push(ChoiceEmission::PassThrough(prefix_choice)); @@ -352,7 +347,6 @@ impl ChoiceJailState { &content, None, choice.finish_reason, - None, choice.logprobs.clone(), ); emissions.push(ChoiceEmission::PassThrough(pass_through_choice)); @@ -415,7 +409,6 @@ impl ChoiceJailState { &trailing_owned, None, choice.finish_reason, - None, choice.logprobs.clone(), ); emissions.push(ChoiceEmission::Trailing(trailing_choice)); @@ -438,7 +431,6 @@ impl ChoiceJailState { &self.accumulated_content, None, self.stream_finish_reason, // For the accumulated content, assign the original stream finish reason, otherwise it will get lost - None, self.accumulated_logprobs.clone(), ); @@ -666,7 +658,6 @@ impl JailedStream { index: choice.index, delta: choice.delta.clone(), finish_reason: choice.finish_reason, - stop_reason: choice.stop_reason.clone(), logprobs: choice.logprobs.clone(), }; all_emissions.push(ChoiceEmission::PassThrough(pass_through_choice)); @@ -980,7 +971,6 @@ impl JailedStream { normal_text.as_deref().unwrap_or(""), None, base_choice.finish_reason, - base_choice.stop_reason.clone(), base_choice.logprobs.clone(), ); } @@ -1005,7 +995,6 @@ impl JailedStream { normal_text.as_deref().unwrap_or(""), Some(tool_call_chunks), None, - None, base_choice.logprobs.clone(), ) } @@ -1033,7 +1022,6 @@ impl JailedStream { content, None, base_choice.finish_reason, - base_choice.stop_reason.clone(), base_choice.logprobs.clone(), ) } @@ -1052,7 +1040,6 @@ impl JailedStream { "", None, base_choice.finish_reason, - base_choice.stop_reason.clone(), base_choice.logprobs.clone(), ) } @@ -1171,7 +1158,6 @@ impl JailedStream { "", Some(tool_call_chunks), base_choice.finish_reason, - None, base_choice.logprobs.clone(), ) } else if filter_dropped_all { @@ -1183,7 +1169,6 @@ impl JailedStream { "", None, base_choice.finish_reason, - base_choice.stop_reason.clone(), base_choice.logprobs.clone(), ) } else { @@ -1194,7 +1179,6 @@ impl JailedStream { accumulated_content, None, base_choice.finish_reason, - base_choice.stop_reason.clone(), base_choice.logprobs.clone(), ) } @@ -1568,7 +1552,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }; @@ -1637,6 +1620,7 @@ mod tests { } /// Helper: build a single-choice stream chunk with text content and logprobs + #[allow(deprecated)] fn text_chunk_with_logprobs(text: &str) -> Annotated { let logprobs = ChatChoiceLogprobs { content: Some( @@ -1668,7 +1652,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: Some(logprobs), }; diff --git a/lib/llm/src/protocols/openai/completions.rs b/lib/llm/src/protocols/openai/completions.rs index 0537277522e9..e60890214bb9 100644 --- a/lib/llm/src/protocols/openai/completions.rs +++ b/lib/llm/src/protocols/openai/completions.rs @@ -40,6 +40,11 @@ pub struct NvCreateCompletionRequest { #[serde(skip_serializing_if = "Option::is_none")] pub metadata: Option, + /// When true, logprob token fields are returned as "token_id:" + /// instead of the decoded text. + #[serde(skip_serializing_if = "Option::is_none")] + pub return_tokens_as_token_ids: Option, + /// Catch-all for unsupported fields - checked during validation #[serde(flatten, default, skip_serializing)] pub unsupported_fields: std::collections::HashMap, @@ -243,11 +248,11 @@ impl OpenAIStopConditionsProvider for NvCreateCompletionRequest { } fn get_stop(&self) -> Option> { - use dynamo_protocols::types::Stop; - self.inner.stop.as_ref().map(|s| match s { - Stop::String(s) => vec![s.clone()], - Stop::StringArray(arr) => arr.clone(), - }) + self.inner.stop.as_ref().and_then(|stop| stop.strings()) + } + + fn get_stop_token_ids(&self) -> Option> { + self.inner.stop.as_ref().and_then(|stop| stop.token_ids()) } fn nvext(&self) -> Option<&NvExt> { @@ -416,6 +421,10 @@ impl OpenAIOutputOptionsProvider for NvCreateCompletionRequest { fn get_formatted_prompt(&self) -> Option { None } + + fn get_return_tokens_as_token_ids(&self) -> Option { + self.return_tokens_as_token_ids + } } /// Implements `ValidateRequest` for `NvCreateCompletionRequest`, @@ -660,6 +669,7 @@ mod tests { let request: NvCreateCompletionRequest = serde_json::from_value(null_stop).expect("Failed to deserialize request"); assert_eq!(request.get_stop(), None); + assert_eq!(request.get_stop_token_ids(), None); let one_stop = json!({ "model": "test-model", @@ -669,6 +679,7 @@ mod tests { let request: NvCreateCompletionRequest = serde_json::from_value(one_stop).expect("Failed to deserialize request"); assert_eq!(request.get_stop(), Some(vec!["foo".to_string()])); + assert_eq!(request.get_stop_token_ids(), None); let many_stops = json!({ "model": "test-model", @@ -681,5 +692,54 @@ mod tests { request.get_stop(), Some(vec!["foo".to_string(), "bar".to_string()]) ); + assert_eq!(request.get_stop_token_ids(), None); + + let token_id_stop = json!({ + "model": "test-model", + "prompt": [1, 2, 3], + "stop": [32, 34] + }); + let request: NvCreateCompletionRequest = + serde_json::from_value(token_id_stop).expect("Failed to deserialize request"); + assert_eq!(request.get_stop(), None); + assert_eq!(request.get_stop_token_ids(), Some(vec![32, 34])); + + let stop_conditions = request + .extract_stop_conditions() + .expect("extract stop conditions"); + assert_eq!(stop_conditions.stop, None); + assert_eq!(stop_conditions.stop_token_ids, Some(vec![32, 34])); + assert_eq!(stop_conditions.stop_token_ids_hidden, None); + + let token_id_display_string_scalar_stop = json!({ + "model": "test-model", + "prompt": [1, 2, 3], + "stop": "token_id:576" + }); + let request: NvCreateCompletionRequest = + serde_json::from_value(token_id_display_string_scalar_stop) + .expect("Failed to deserialize request"); + assert_eq!(request.get_stop(), Some(vec!["token_id:576".to_string()])); + assert_eq!(request.get_stop_token_ids(), None); + + let token_id_display_string_stop = json!({ + "model": "test-model", + "prompt": [1, 2, 3], + "stop": ["token_id:576"] + }); + let request: NvCreateCompletionRequest = + serde_json::from_value(token_id_display_string_stop) + .expect("Failed to deserialize request"); + assert_eq!(request.get_stop(), Some(vec!["token_id:576".to_string()])); + assert_eq!(request.get_stop_token_ids(), None); + + let unsupported_stop_token_ids = json!({ + "model": "test-model", + "prompt": [1, 2, 3], + "stop_token_ids": [576] + }); + let request: NvCreateCompletionRequest = serde_json::from_value(unsupported_stop_token_ids) + .expect("Failed to deserialize request"); + assert!(ValidateRequest::validate(&request).is_err()); } } diff --git a/lib/llm/src/protocols/openai/completions/aggregator.rs b/lib/llm/src/protocols/openai/completions/aggregator.rs index 6bea09f59adb..90dfe93f9733 100644 --- a/lib/llm/src/protocols/openai/completions/aggregator.rs +++ b/lib/llm/src/protocols/openai/completions/aggregator.rs @@ -12,7 +12,7 @@ use crate::protocols::{ codec::{Message, SseCodecError}, common::FinishReason, convert_sse_stream, - openai::ParsingOptions, + openai::{ParsingOptions, nvext::merge_response_nvext}, }; /// Aggregates a stream of [`CompletionResponse`]s into a single [`CompletionResponse`]. @@ -85,10 +85,7 @@ impl DeltaAggregator { if let Some(system_fingerprint) = delta.inner.system_fingerprint { aggregator.system_fingerprint = Some(system_fingerprint); } - // Aggregate nvext field (take the last non-None value) - if delta.nvext.is_some() { - aggregator.nvext = delta.nvext; - } + merge_response_nvext(&mut aggregator.nvext, delta.nvext); // handle the choices for choice in delta.inner.choices { @@ -334,8 +331,10 @@ mod tests { // One will have a MessageRole and no FinishReason, // the other will have a FinishReason and no MessageRole let annotated_delta1 = create_test_delta(0, "Hello,", None, Some(-0.1)); - let annotated_delta2 = + let mut annotated_delta2 = create_test_delta(0, " world!", Some("stop".to_string()), Some(-0.2)); + annotated_delta2.data.as_mut().expect("delta data").nvext = + Some(serde_json::json!({ "stop_reason": 128001 })); // Create a stream let annotated_deltas = vec![annotated_delta1, annotated_delta2]; @@ -357,6 +356,10 @@ mod tests { choice.finish_reason, Some(dynamo_protocols::types::CompletionFinishReason::Stop) ); + assert_eq!( + response.nvext, + Some(serde_json::json!({ "stop_reason": 128001 })) + ); assert_eq!(choice.logprobs.as_ref().unwrap().tokens.len(), 2); assert_eq!( choice.logprobs.as_ref().unwrap().token_logprobs, @@ -364,6 +367,29 @@ mod tests { ); } + #[tokio::test] + async fn test_multiple_deltas_merge_nvext_fields() { + let mut annotated_delta1 = create_test_delta(0, "Hello,", None, None); + annotated_delta1.data.as_mut().expect("delta data").nvext = + Some(serde_json::json!({ "engine_data": { "trace_id": "abc" } })); + let mut annotated_delta2 = create_test_delta(0, " world!", Some("stop".to_string()), None); + annotated_delta2.data.as_mut().expect("delta data").nvext = + Some(serde_json::json!({ "stop_reason": 128001 })); + + let stream = Box::pin(stream::iter(vec![annotated_delta1, annotated_delta2])); + let response = DeltaAggregator::apply(stream, ParsingOptions::default()) + .await + .expect("aggregate stream"); + + assert_eq!( + response.nvext, + Some(serde_json::json!({ + "engine_data": { "trace_id": "abc" }, + "stop_reason": 128001, + })) + ); + } + #[tokio::test] async fn test_multiple_choices() { // Create a delta with multiple choices diff --git a/lib/llm/src/protocols/openai/completions/delta.rs b/lib/llm/src/protocols/openai/completions/delta.rs index 3f039399e1cd..38cc8d66b642 100644 --- a/lib/llm/src/protocols/openai/completions/delta.rs +++ b/lib/llm/src/protocols/openai/completions/delta.rs @@ -60,8 +60,9 @@ impl NvCreateCompletionRequest { .as_ref() .map(|opts| opts.continuous_usage_stats) .unwrap_or(false), - enable_logprobs: self.inner.logprobs.unwrap_or(0) > 0, + enable_logprobs: self.inner.logprobs.is_some(), response_fields, + return_tokens_as_token_ids: self.return_tokens_as_token_ids.unwrap_or(false), }; DeltaGenerator::new(self.inner.model.clone(), options, request_id) @@ -74,6 +75,8 @@ pub struct DeltaGeneratorOptions { pub continuous_usage_stats: bool, pub enable_logprobs: bool, pub response_fields: NvExtResponseFieldSelection, + /// When true, logprob token fields use "token_id:" format instead of decoded text. + pub return_tokens_as_token_ids: bool, } pub struct DeltaGenerator { @@ -158,19 +161,32 @@ impl DeltaGenerator { .map(|(_, lp)| lp as f32) .collect::>(); + let return_as_ids = self.options.return_tokens_as_token_ids; let top_lps = top_logprobs.map_or(vec![], |top_logprobs| { toks.iter() .zip(tok_lps.iter()) .zip(top_logprobs.iter()) .map(|(((t, tid), lp), top_lps)| { - let converted = convert_backend_top_logprobs(top_lps, t, *tid, *lp); + let converted = + convert_backend_top_logprobs(top_lps, t, *tid, *lp, return_as_ids); serde_json::to_value(converted).unwrap() }) .collect() }); + let tokens_out: Vec = toks + .iter() + .map(|(t, tid)| { + if return_as_ids { + format!("token_id:{}", tid) + } else { + t.clone() + } + }) + .collect(); + Some(dynamo_protocols::types::Logprobs { - tokens: toks.iter().map(|(t, _)| t.clone()).collect(), + tokens: tokens_out, token_logprobs: tok_lps.into_iter().map(Some).collect(), text_offset: vec![], top_logprobs: top_lps, @@ -291,6 +307,7 @@ impl crate::protocols::openai::DeltaGeneratorExt for ); let finish_reason = delta.finish_reason.map(Into::into); + let stop_reason = delta.stop_reason.clone(); // create choice let index = delta.index.unwrap_or(0); @@ -313,6 +330,7 @@ impl crate::protocols::openai::DeltaGeneratorExt for delta.disaggregated_params.as_ref(), finish_reason.is_some(), delta.engine_data, + stop_reason, ) && let Ok(nvext_json) = serde_json::to_value(&nvext_response) { response.nvext = Some(nvext_json); @@ -378,6 +396,7 @@ mod tests { common: Default::default(), nvext: None, metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } @@ -427,6 +446,7 @@ mod tests { .unwrap(), ), metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } @@ -466,6 +486,102 @@ mod tests { assert!(response.nvext.is_none()); } + #[test] + fn test_stop_reason_is_suppressed_without_nvext_extra_field() { + let request = create_test_request(); + let mut generator = request.response_generator("req-stop-reason".to_string()); + let mut output = final_backend_output(); + output.stop_reason = Some(dynamo_protocols::types::StopReason::String( + "END".to_string(), + )); + + let response = generator + .choice_from_postprocessor(output) + .expect("choice generation"); + + let response_json = serde_json::to_value(&response).expect("serialize response"); + assert!(response_json["choices"][0].get("stop_reason").is_none()); + assert!(response_json.get("nvext").is_none()); + } + + #[test] + fn test_stop_reason_emits_in_nvext_when_requested() { + let request = create_test_request_with_extra_fields(vec!["stop_reason".to_string()]); + let mut generator = request.response_generator("req-stop-reason-nvext".to_string()); + let mut output = final_backend_output(); + output.stop_reason = Some(dynamo_protocols::types::StopReason::String( + "END".to_string(), + )); + + let response = generator + .choice_from_postprocessor(output) + .expect("choice generation"); + + let response_json = serde_json::to_value(&response).expect("serialize response"); + assert!(response_json["choices"][0].get("stop_reason").is_none()); + assert_eq!(response_json["nvext"]["stop_reason"], "END"); + } + + #[test] + fn test_logprobs_zero_emits_chosen_token_logprob() { + let mut request = create_test_request(); + request.inner.logprobs = Some(0); + let mut generator = request.response_generator("req-logprobs-zero".to_string()); + let mut output = final_backend_output(); + output.log_probs = Some(vec![-0.5]); + + let response = generator + .choice_from_postprocessor(output) + .expect("choice generation"); + let logprobs = response.inner.choices[0] + .logprobs + .as_ref() + .expect("logprobs"); + + assert_eq!(logprobs.tokens, vec!["hello"]); + assert_eq!(logprobs.token_logprobs, vec![Some(-0.5)]); + assert!(logprobs.top_logprobs.is_empty()); + } + + #[test] + fn test_return_token_ids_formats_selected_top_logprob_fallback() { + let mut request = create_test_request(); + request.inner.logprobs = Some(1); + request.return_tokens_as_token_ids = Some(true); + let generator = request.response_generator("req-token-id-logprobs".to_string()); + + let logprobs = generator + .create_logprobs( + vec![Some("hello".to_string())], + vec![123], + Some(vec![-0.5]), + Some(vec![vec![common::llm_backend::TopLogprob { + rank: 1, + token_id: 999, + token: Some("other".to_string()), + logprob: -1.0, + bytes: None, + }]]), + ) + .expect("logprobs"); + + assert_eq!(logprobs.tokens, vec!["token_id:123"]); + let top_logprobs = logprobs.top_logprobs[0] + .as_array() + .expect("top_logprobs array"); + let other = top_logprobs + .iter() + .find(|item| item["token"] == "token_id:999") + .expect("top token_id formatting"); + assert_eq!(other["bytes"], serde_json::json!(b"token_id:999")); + let selected = top_logprobs + .iter() + .find(|item| item["token"] == "token_id:123") + .expect("selected token fallback"); + assert_eq!(selected["token"], "token_id:123"); + assert_eq!(selected["bytes"], serde_json::json!(b"token_id:123")); + } + #[test] fn test_timing_extra_field_emits_timing_on_final_chunk() { use crate::protocols::openai::nvext::NvExt; diff --git a/lib/llm/src/protocols/openai/nvext.rs b/lib/llm/src/protocols/openai/nvext.rs index 3a454990dcc2..7c561dd4a59d 100644 --- a/lib/llm/src/protocols/openai/nvext.rs +++ b/lib/llm/src/protocols/openai/nvext.rs @@ -3,6 +3,7 @@ use axum::http::HeaderMap; use derive_builder::Builder; +use dynamo_protocols::types::StopReason; use serde::{Deserialize, Serialize}; use utoipa::ToSchema; use validator::{Validate, ValidationError}; @@ -120,6 +121,32 @@ pub struct NvExtResponse { /// Dynamo does not inspect this; it is forwarded as-is to the client. #[serde(skip_serializing_if = "Option::is_none")] pub engine_data: Option, + + /// Backend-specific matched stop condition. This is not part of the + /// OpenAI response schema, so it is only returned under nvext when requested. + /// + /// This is response-level for Dynamo's current single-choice serving paths. + /// If `n > 1` is supported here, this needs an indexed/per-choice shape. + #[serde(skip_serializing_if = "Option::is_none")] + pub stop_reason: Option, +} + +pub(crate) fn merge_response_nvext( + target: &mut Option, + incoming: Option, +) { + let Some(incoming) = incoming else { + return; + }; + + match (target.as_mut(), incoming) { + (Some(serde_json::Value::Object(target_obj)), serde_json::Value::Object(incoming_obj)) => { + target_obj.extend(incoming_obj); + } + (_, incoming) => { + *target = Some(incoming); + } + } } /// Response nvext fields requested for a given request. @@ -137,6 +164,7 @@ pub struct NvExtResponseFieldSelection { pub token_ids: bool, pub routed_experts: bool, pub engine_data: bool, + pub stop_reason: bool, } impl NvExtResponseFieldSelection { @@ -153,6 +181,7 @@ impl NvExtResponseFieldSelection { "timing" => selection.timing = true, "routed_experts" => selection.routed_experts = true, "engine_data" => selection.engine_data = true, + "stop_reason" => selection.stop_reason = true, _ => {} } } @@ -185,12 +214,14 @@ impl NvExtResponseFieldSelection { /// `disaggregated_params` (cloned as-is, no validation). /// - `timing` requires the selection flag, `finish_reason_present == true`, **and** a tracker. /// - `engine_data` requires the selection flag **and** a non-`None` `engine_data_from_backend`. + /// - `stop_reason` requires the selection flag **and** a non-`None` `stop_reason_from_backend`. pub fn build_response_nvext( &self, tracker: Option<&std::sync::Arc>, disaggregated_params: Option<&serde_json::Value>, finish_reason_present: bool, engine_data_from_backend: Option, + stop_reason_from_backend: Option, ) -> Option { let worker_id = if self.worker_id { tracker.and_then(|t| t.get_worker_info()) @@ -226,11 +257,18 @@ impl NvExtResponseFieldSelection { None }; + let stop_reason = if self.stop_reason { + stop_reason_from_backend.and_then(|reason| serde_json::to_value(reason).ok()) + } else { + None + }; + if worker_id.is_none() && token_ids.is_none() && routed_experts.is_none() && timing.is_none() && engine_data.is_none() + && stop_reason.is_none() { return None; } @@ -241,6 +279,7 @@ impl NvExtResponseFieldSelection { token_ids, routed_experts, engine_data, + stop_reason, }) } } @@ -292,7 +331,7 @@ pub struct NvExt { /// Extra fields to be included in the response's nvext /// This is a list of field names that should be populated in the response /// Supported fields include "worker_id", "timing", "routed_experts", "engine_data", - /// which map to fields in NvExtResponse. + /// "stop_reason", which map to fields in NvExtResponse. #[serde(default, skip_serializing_if = "Option::is_none")] #[builder(default, setter(strip_option))] pub extra_fields: Option>, @@ -709,6 +748,22 @@ mod tests { ); } + #[test] + fn test_nvext_response_field_selection_stop_reason_only() { + let nvext = NvExt::builder() + .extra_fields(vec!["stop_reason".to_string()]) + .build() + .unwrap(); + + assert_eq!( + NvExtResponseFieldSelection::from_nvext(Some(&nvext)), + NvExtResponseFieldSelection { + stop_reason: true, + ..Default::default() + } + ); + } + // Helpers for build_response_nvext tests ----------------------------- fn sel_all_false() -> NvExtResponseFieldSelection { @@ -736,11 +791,13 @@ mod tests { fn test_build_response_nvext_all_false_returns_none() { let sel = sel_all_false(); assert!( - sel.build_response_nvext(None, None, false, None).is_none(), + sel.build_response_nvext(None, None, false, None, None) + .is_none(), "no fields selected → None" ); assert!( - sel.build_response_nvext(None, None, true, None).is_none(), + sel.build_response_nvext(None, None, true, None, None) + .is_none(), "finish_reason alone does not force emission" ); } @@ -755,7 +812,7 @@ mod tests { // finish_reason=false: worker_id still emitted (only timing is finish-gated). let out = sel - .build_response_nvext(Some(&tracker), None, false, None) + .build_response_nvext(Some(&tracker), None, false, None, None) .expect("worker_id should emit regardless of finish_reason"); assert!(out.worker_id.is_some()); @@ -774,7 +831,7 @@ mod tests { // timing alone + finish_reason=false → nothing to emit, returns None. assert!( - sel.build_response_nvext(Some(&tracker), None, false, None) + sel.build_response_nvext(Some(&tracker), None, false, None, None) .is_none(), "timing is gated on finish_reason_present" ); @@ -789,7 +846,7 @@ mod tests { let tracker = tracker_with_prefill_worker(); let out = sel - .build_response_nvext(Some(&tracker), None, true, None) + .build_response_nvext(Some(&tracker), None, true, None, None) .expect("timing should emit on finish"); assert!(out.timing.is_some()); @@ -805,7 +862,10 @@ mod tests { ..Default::default() }; // finish=true but no tracker → timing not populated → None. - assert!(sel.build_response_nvext(None, None, true, None).is_none()); + assert!( + sel.build_response_nvext(None, None, true, None, None) + .is_none() + ); } #[test] @@ -817,7 +877,7 @@ mod tests { let params = disagg_params_full(); let out = sel - .build_response_nvext(None, Some(¶ms), false, None) + .build_response_nvext(None, Some(¶ms), false, None, None) .expect("token_ids should emit when present"); assert_eq!(out.token_ids, Some(vec![11u32, 22, 33])); @@ -836,7 +896,7 @@ mod tests { let params = serde_json::json!({ "token_ids": "not-an-array" }); assert!( - sel.build_response_nvext(None, Some(¶ms), false, None) + sel.build_response_nvext(None, Some(¶ms), false, None, None) .is_none(), "malformed token_ids silently suppressed; nothing else selected → None" ); @@ -851,7 +911,7 @@ mod tests { let params = disagg_params_full(); let out = sel - .build_response_nvext(None, Some(¶ms), false, None) + .build_response_nvext(None, Some(¶ms), false, None, None) .expect("routed_experts should emit when present"); assert_eq!( @@ -860,6 +920,43 @@ mod tests { ); } + #[test] + fn test_build_response_nvext_stop_reason_when_requested() { + let sel = NvExtResponseFieldSelection { + stop_reason: true, + ..Default::default() + }; + + let out = sel + .build_response_nvext( + None, + None, + true, + None, + Some(StopReason::String("END".to_string())), + ) + .expect("stop_reason should emit when requested and present"); + + assert_eq!(out.stop_reason, Some(serde_json::json!("END"))); + assert!(out.worker_id.is_none()); + assert!(out.timing.is_none()); + assert!(out.token_ids.is_none()); + assert!(out.routed_experts.is_none()); + } + + #[test] + fn test_build_response_nvext_stop_reason_suppressed_when_absent() { + let sel = NvExtResponseFieldSelection { + stop_reason: true, + ..Default::default() + }; + + assert!( + sel.build_response_nvext(None, None, true, None, None) + .is_none() + ); + } + #[test] fn test_build_response_nvext_combined_emission() { let sel = NvExtResponseFieldSelection { @@ -868,12 +965,13 @@ mod tests { token_ids: true, routed_experts: true, engine_data: false, + stop_reason: false, }; let tracker = tracker_with_prefill_worker(); let params = disagg_params_full(); let out = sel - .build_response_nvext(Some(&tracker), Some(¶ms), true, None) + .build_response_nvext(Some(&tracker), Some(¶ms), true, None, None) .expect("all fields selected and available → Some"); assert!(out.worker_id.is_some()); @@ -904,6 +1002,7 @@ mod tests { token_ids: false, // only enabled via query_instance_id routed_experts: true, engine_data: false, + stop_reason: false, } ); } diff --git a/lib/llm/src/protocols/openai/responses/mod.rs b/lib/llm/src/protocols/openai/responses/mod.rs index 2dcc530573c6..ae1c868af84a 100644 --- a/lib/llm/src/protocols/openai/responses/mod.rs +++ b/lib/llm/src/protocols/openai/responses/mod.rs @@ -731,6 +731,7 @@ impl TryFrom for NvCreateChatCompletionRequest { nvext: resp.nvext, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }) } @@ -2102,7 +2103,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }], created: now, @@ -2163,7 +2163,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }], created: now, @@ -2571,7 +2570,6 @@ thinking audio: None, }, finish_reason: Some(FinishReason::Stop), - stop_reason: None, logprobs: None, }], created: 0, diff --git a/lib/llm/src/protocols/openai/responses/stream_converter.rs b/lib/llm/src/protocols/openai/responses/stream_converter.rs index ff631bff01f7..d4a3a4948939 100644 --- a/lib/llm/src/protocols/openai/responses/stream_converter.rs +++ b/lib/llm/src/protocols/openai/responses/stream_converter.rs @@ -726,7 +726,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }], created: 0, @@ -756,7 +755,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }], created: 0, diff --git a/lib/llm/src/protocols/openai/validate.rs b/lib/llm/src/protocols/openai/validate.rs index 237e84bc75be..559dd109ac1b 100644 --- a/lib/llm/src/protocols/openai/validate.rs +++ b/lib/llm/src/protocols/openai/validate.rs @@ -380,6 +380,18 @@ pub fn validate_stop(stop: &Option) -> Result<(), } } } + dynamo_protocols::types::Stop::TokenIdArray(token_ids) => { + if token_ids.is_empty() { + anyhow::bail!("Stop token IDs array cannot be empty"); + } + if token_ids.len() > MAX_STOP_SEQUENCES { + anyhow::bail!( + "Maximum of {} stop token IDs allowed, got {}", + MAX_STOP_SEQUENCES, + token_ids.len() + ); + } + } } } Ok(()) diff --git a/lib/llm/src/protocols/unified.rs b/lib/llm/src/protocols/unified.rs index 6ce62744e7f3..c748678126e5 100644 --- a/lib/llm/src/protocols/unified.rs +++ b/lib/llm/src/protocols/unified.rs @@ -404,10 +404,19 @@ impl OpenAIStopConditionsProvider for UnifiedRequest { } fn get_stop(&self) -> Option> { - self.inner.inner.stop.as_ref().map(|stop| match stop { - dynamo_protocols::types::Stop::String(s) => vec![s.clone()], - dynamo_protocols::types::Stop::StringArray(arr) => arr.clone(), - }) + self.inner + .inner + .stop + .as_ref() + .and_then(|stop| stop.strings()) + } + + fn get_stop_token_ids(&self) -> Option> { + self.inner + .inner + .stop + .as_ref() + .and_then(|stop| stop.token_ids()) } fn nvext(&self) -> Option<&NvExt> { @@ -534,6 +543,7 @@ mod tests { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; diff --git a/lib/llm/tests/aggregators.rs b/lib/llm/tests/aggregators.rs index 5ce5e86eb95f..367522b89118 100644 --- a/lib/llm/tests/aggregators.rs +++ b/lib/llm/tests/aggregators.rs @@ -172,7 +172,6 @@ fn make_stream_delta( reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }] } else { diff --git a/lib/llm/tests/http-service.rs b/lib/llm/tests/http-service.rs index 71d0c15b47bf..9bdc40d7ea6a 100644 --- a/lib/llm/tests/http-service.rs +++ b/lib/llm/tests/http-service.rs @@ -85,7 +85,7 @@ impl let stream = stream! { tokio::time::sleep(std::time::Duration::from_millis(max_tokens)).await; for i in 0..10 { - let output = generator.create_choice(i, Some(format!("choice {i}")), None, None, None); + let output = generator.create_choice(i, Some(format!("choice {i}")), None, None); yield Annotated::from_data(output); } diff --git a/lib/llm/tests/http_metrics.rs b/lib/llm/tests/http_metrics.rs index dd54a9011ffc..2d295e89359b 100644 --- a/lib/llm/tests/http_metrics.rs +++ b/lib/llm/tests/http_metrics.rs @@ -55,7 +55,7 @@ impl // output_sequence_tokens is properly recorded (the histogram only // records when osl > 0, which requires the annotation to be present). for i in 0..5 { - let output = generator.create_choice(i, Some(format!("Mock response {i}")), None, None, None); + let output = generator.create_choice(i, Some(format!("Mock response {i}")), None, None); let mut annotated = Annotated::from_data(output); let metrics = LLMMetricAnnotation { input_tokens: 5, diff --git a/lib/llm/tests/logprob_analysis_integration.rs b/lib/llm/tests/logprob_analysis_integration.rs index 1adcf655ae06..f684991fbf5d 100644 --- a/lib/llm/tests/logprob_analysis_integration.rs +++ b/lib/llm/tests/logprob_analysis_integration.rs @@ -388,7 +388,6 @@ fn create_response_with_linear_probs( reasoning_content: None, }, finish_reason: Some(FinishReason::Stop), - stop_reason: None, logprobs: Some(ChatChoiceLogprobs { content: Some(token_logprobs), refusal: None, @@ -469,7 +468,6 @@ fn create_multi_choice_response( reasoning_content: None, }, finish_reason: Some(FinishReason::Stop), - stop_reason: None, logprobs: Some(ChatChoiceLogprobs { content: Some(token_logprobs), refusal: None, diff --git a/lib/llm/tests/openai_completions.rs b/lib/llm/tests/openai_completions.rs index 2d772b6af6e6..fd23d916ec44 100644 --- a/lib/llm/tests/openai_completions.rs +++ b/lib/llm/tests/openai_completions.rs @@ -29,6 +29,7 @@ impl CompletionSample { common: Default::default(), nvext: None, metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; diff --git a/lib/llm/tests/parallel_tool_call_integration.rs b/lib/llm/tests/parallel_tool_call_integration.rs index 2827239d4754..4eeb89026f74 100644 --- a/lib/llm/tests/parallel_tool_call_integration.rs +++ b/lib/llm/tests/parallel_tool_call_integration.rs @@ -92,6 +92,7 @@ fn create_mock_chat_completion_request() -> NvCreateChatCompletionRequest { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } diff --git a/lib/llm/tests/postprocessor_parsing_stream.rs b/lib/llm/tests/postprocessor_parsing_stream.rs index 2bb48cd4cbfc..9a5ef8fb3305 100644 --- a/lib/llm/tests/postprocessor_parsing_stream.rs +++ b/lib/llm/tests/postprocessor_parsing_stream.rs @@ -289,7 +289,6 @@ fn mock_content_chunk(content: &str) -> NvCreateChatCompletionStreamResponse { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }; NvCreateChatCompletionStreamResponse { @@ -324,7 +323,6 @@ fn mock_final_chunk() -> NvCreateChatCompletionStreamResponse { reasoning_content: None, }, finish_reason: Some(FinishReason::Stop), - stop_reason: None, logprobs: None, }; NvCreateChatCompletionStreamResponse { diff --git a/lib/llm/tests/preprocessor.rs b/lib/llm/tests/preprocessor.rs index 9de0ad62e842..144115f527a2 100644 --- a/lib/llm/tests/preprocessor.rs +++ b/lib/llm/tests/preprocessor.rs @@ -260,6 +260,7 @@ impl Request { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } @@ -700,6 +701,7 @@ mod context_length_validation { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } diff --git a/lib/llm/tests/test_common_ext.rs b/lib/llm/tests/test_common_ext.rs index 8e49c7377b09..b2366001eab1 100644 --- a/lib/llm/tests/test_common_ext.rs +++ b/lib/llm/tests/test_common_ext.rs @@ -69,6 +69,7 @@ fn test_sampling_parameters_include_stop_str_in_output_extraction() { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; @@ -299,6 +300,7 @@ fn test_serialization_preserves_structure() { }), chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; @@ -351,6 +353,7 @@ fn test_sampling_parameters_extraction() { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; diff --git a/lib/llm/tests/test_jail.rs b/lib/llm/tests/test_jail.rs index 7a8da4aa4dcd..9c79eead37ab 100644 --- a/lib/llm/tests/test_jail.rs +++ b/lib/llm/tests/test_jail.rs @@ -45,7 +45,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }; @@ -88,7 +87,6 @@ mod tests { reasoning_content: None, }, finish_reason: Some(FinishReason::Stop), - stop_reason: None, logprobs: None, }; @@ -135,7 +133,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }; @@ -181,7 +178,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, } }) @@ -229,7 +225,6 @@ mod tests { reasoning_content: None, }, finish_reason: Some(FinishReason::Stop), - stop_reason: None, logprobs: None, } }) @@ -2405,7 +2400,6 @@ mod parallel_jail_tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, } }) diff --git a/lib/llm/tests/test_reasoning_parser.rs b/lib/llm/tests/test_reasoning_parser.rs index 49113b8a660a..c5d0202816b8 100644 --- a/lib/llm/tests/test_reasoning_parser.rs +++ b/lib/llm/tests/test_reasoning_parser.rs @@ -34,7 +34,6 @@ fn create_mock_response_chunk( reasoning_content, }, finish_reason: None, - stop_reason: None, logprobs: None, }; diff --git a/lib/llm/tests/test_stop_behavior.rs b/lib/llm/tests/test_stop_behavior.rs index 1a45638efd65..e61e2612eec3 100644 --- a/lib/llm/tests/test_stop_behavior.rs +++ b/lib/llm/tests/test_stop_behavior.rs @@ -136,3 +136,22 @@ fn stop_token_priority_over_sequence() { Some(StopTrigger::HiddenStopTokenDetected(id)) if id == STOP )); } + +#[test] +fn user_stop_token_reports_distinct_trigger() { + let tokenizer: Arc = Arc::new(TestTokenizer); + let decode_stream = tokenizers::DecodeStream::new(tokenizer, &[], false); + let stop_conditions = StopConditions { + stop_token_ids: Some(vec![STOP]), + stop_token_ids_hidden: Some(vec![EOS]), + ..Default::default() + }; + let mut decoder = Decoder::new(decode_stream, stop_conditions, false, None); + let result = decoder.process_token_ids(&[HI, STOP]).unwrap(); + + assert_eq!(result.text.as_deref(), Some("hi")); + assert!(matches!( + result.stop_trigger, + Some(StopTrigger::UserStopTokenDetected(id)) if id == STOP + )); +} diff --git a/lib/llm/tests/test_streaming_tool_parsers.rs b/lib/llm/tests/test_streaming_tool_parsers.rs index d8b95b9d96bd..2fb30575ce8e 100644 --- a/lib/llm/tests/test_streaming_tool_parsers.rs +++ b/lib/llm/tests/test_streaming_tool_parsers.rs @@ -1290,7 +1290,6 @@ mod tests { reasoning_content: None, }, finish_reason, - stop_reason: None, logprobs: None, }; Annotated { diff --git a/lib/llm/tests/test_streaming_usage.rs b/lib/llm/tests/test_streaming_usage.rs index 0a6fd3178bf6..eb91c305da68 100644 --- a/lib/llm/tests/test_streaming_usage.rs +++ b/lib/llm/tests/test_streaming_usage.rs @@ -194,6 +194,7 @@ fn create_chat_request( nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } @@ -501,6 +502,7 @@ fn create_cmpl_request(include_usage: Option, stream: bool) -> NvCreateCom common: Default::default(), nvext: None, metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } @@ -528,6 +530,7 @@ fn create_nonstreaming_chat_request() -> NvCreateChatCompletionRequest { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } diff --git a/lib/llm/tests/tool_choice.rs b/lib/llm/tests/tool_choice.rs index 350498fbaaa6..0a720f985a36 100644 --- a/lib/llm/tests/tool_choice.rs +++ b/lib/llm/tests/tool_choice.rs @@ -40,6 +40,7 @@ fn create_test_request() -> NvCreateChatCompletionRequest { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } @@ -490,7 +491,6 @@ fn make_text_chunk( } else { None }, - stop_reason: None, logprobs: None, }], created: 1234567890, diff --git a/lib/llm/tests/tool_choice_finish_reasons.rs b/lib/llm/tests/tool_choice_finish_reasons.rs index d3d190c3953c..9f66338d24b7 100644 --- a/lib/llm/tests/tool_choice_finish_reasons.rs +++ b/lib/llm/tests/tool_choice_finish_reasons.rs @@ -33,6 +33,7 @@ fn create_test_request() -> NvCreateChatCompletionRequest { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } diff --git a/lib/protocols/src/types/chat.rs b/lib/protocols/src/types/chat.rs index c9cef996b9df..75026c9ff22d 100644 --- a/lib/protocols/src/types/chat.rs +++ b/lib/protocols/src/types/chat.rs @@ -77,8 +77,72 @@ pub use async_openai::types::chat::{ WebSearchUserLocationType, }; -// Upstream renamed Stop -> StopConfiguration; re-export under old name for compat -pub use async_openai::types::chat::StopConfiguration as Stop; +/// OpenAI stop configuration, with Dynamo's token-id stop extension. +/// +/// The standard OpenAI shape accepts a string or string array. Dynamo also +/// accepts an integer array, e.g. `"stop": [576]`, to express token-id stop +/// conditions for tokenized in/out workflows. Strings like `"token_id:576"` +/// remain ordinary string stops; the `token_id:` format is only an output +/// display format for logprobs. +#[derive(Debug, Deserialize, Serialize, Clone, PartialEq)] +#[serde(untagged)] +pub enum Stop { + String(String), + StringArray(Vec), + TokenIdArray(Vec), +} + +impl Stop { + pub fn strings(&self) -> Option> { + match self { + Stop::String(s) => Some(vec![s.clone()]), + Stop::StringArray(arr) => Some(arr.clone()), + Stop::TokenIdArray(_) => None, + } + } + + pub fn token_ids(&self) -> Option> { + match self { + Stop::TokenIdArray(arr) => Some(arr.clone()), + Stop::String(_) | Stop::StringArray(_) => None, + } + } +} + +impl From for Stop { + fn from(value: String) -> Self { + Stop::String(value) + } +} + +impl From<&str> for Stop { + fn from(value: &str) -> Self { + Stop::String(value.to_string()) + } +} + +impl From> for Stop { + fn from(value: Vec) -> Self { + Stop::StringArray(value) + } +} + +impl From> for Stop { + fn from(value: Vec) -> Self { + Stop::TokenIdArray(value) + } +} + +impl From for Stop { + fn from(value: async_openai::types::chat::StopConfiguration) -> Self { + match value { + async_openai::types::chat::StopConfiguration::String(value) => Stop::String(value), + async_openai::types::chat::StopConfiguration::StringArray(value) => { + Stop::StringArray(value) + } + } + } +} // Upstream renamed FinishReason (streaming) -- re-export pub use async_openai::types::chat::FinishReason; @@ -281,11 +345,13 @@ pub struct ChatCompletionTool { /// Inference backends (vLLM, SGLang) report which stop condition triggered: /// - `String`: a matched user-provided stop sequence /// - `Int`: a matched stop token ID +/// - `IntArray`: matched stop token IDs reported as a sequence #[derive(Debug, Deserialize, Serialize, Clone, PartialEq)] #[serde(untagged)] pub enum StopReason { String(String), Int(i64), + IntArray(Vec), } /// Reasoning content from a previous assistant turn. @@ -694,9 +760,6 @@ pub struct ChatChoice { pub index: u32, pub message: ChatCompletionResponseMessage, pub finish_reason: Option, - /// Matched stop condition from the backend. - #[serde(skip_serializing_if = "Option::is_none")] - pub stop_reason: Option, pub logprobs: Option, } @@ -745,19 +808,12 @@ pub struct ChatCompletionStreamResponseDeltaFunctionCall { pub arguments: Option, } -/// Streaming chat choice with stop reason support. -/// -/// Extends upstream `ChatChoiceStream` with: -/// - `stop_reason`: the matched stop sequence (string) or stop token ID (integer) -/// reported by inference backends +/// Streaming chat choice. #[derive(Debug, Deserialize, Clone, PartialEq, Serialize)] pub struct ChatChoiceStream { pub index: u32, pub delta: ChatCompletionStreamResponseDelta, pub finish_reason: Option, - /// Matched stop condition from the backend. - #[serde(skip_serializing_if = "Option::is_none")] - pub stop_reason: Option, pub logprobs: Option, } @@ -778,6 +834,56 @@ pub struct CreateChatCompletionStreamResponse { mod tests { use super::*; + #[test] + fn stop_accepts_token_id_array() { + let stop: Stop = serde_json::from_value(serde_json::json!([32, 34])).unwrap(); + + assert_eq!(stop, Stop::TokenIdArray(vec![32, 34])); + } + + #[test] + fn stop_accepts_string_and_string_array() { + let stop: Stop = serde_json::from_value(serde_json::json!(" The")).unwrap(); + + assert_eq!(stop, Stop::String(" The".to_string())); + + let stop: Stop = serde_json::from_value(serde_json::json!(["A", "B"])).unwrap(); + + assert_eq!( + stop, + Stop::StringArray(vec!["A".to_string(), "B".to_string()]) + ); + } + + #[test] + fn stop_token_id_display_string_remains_string_stop() { + let stop: Stop = serde_json::from_value(serde_json::json!("token_id:576")).unwrap(); + + assert_eq!(stop, Stop::String("token_id:576".to_string())); + + let stop: Stop = serde_json::from_value(serde_json::json!(["token_id:576"])).unwrap(); + + assert_eq!(stop, Stop::StringArray(vec!["token_id:576".to_string()])); + } + + #[test] + fn stop_rejects_single_token_id() { + let result = serde_json::from_value::(serde_json::json!(576)); + + assert!(result.is_err()); + } + + #[test] + fn stop_converts_from_upstream_stop_configuration() { + let upstream = + async_openai::types::chat::StopConfiguration::StringArray(vec!["END".to_string()]); + + assert_eq!( + Stop::from(upstream), + Stop::StringArray(vec!["END".to_string()]) + ); + } + #[test] fn tool_call_defaults_type_on_deserialize() { let tool_call: ChatCompletionMessageToolCall = serde_json::from_value(serde_json::json!({ diff --git a/lib/protocols/src/types/completion.rs b/lib/protocols/src/types/completion.rs index d86a99f0a40c..4c1ecd1ff584 100644 --- a/lib/protocols/src/types/completion.rs +++ b/lib/protocols/src/types/completion.rs @@ -13,7 +13,7 @@ use serde::{Deserialize, Serialize}; use crate::error::OpenAIError; -use super::{ChatCompletionStreamOptions, Choice, CompletionUsage, Prompt, Stop}; +use super::{ChatCompletionStreamOptions, Prompt, Stop}; // Re-export response type from upstream (identical) pub use async_openai::types::completions::CreateCompletionResponse; @@ -170,4 +170,83 @@ mod tests { assert!(err_msg.contains("string")); assert!(err_msg.contains("echo parameter")); } + + #[test] + fn completion_choice_serializes_openai_shape() { + use crate::types::{Choice, CompletionFinishReason}; + + let choice = Choice { + text: "hello".to_string(), + index: 0, + logprobs: None, + finish_reason: Some(CompletionFinishReason::Stop), + }; + + let value = serde_json::to_value(choice).expect("serialize choice"); + + assert_eq!(value["finish_reason"], "stop"); + assert_eq!(value["text"], "hello"); + } + + #[test] + fn stop_accepts_token_id_array() { + let json = r#"{"model": "test_model", "prompt": [1, 2, 3], "stop": [32, 34]}"#; + let request: CreateCompletionRequest = serde_json::from_str(json).unwrap(); + + assert_eq!(request.stop, Some(Stop::TokenIdArray(vec![32, 34]))); + } + + #[test] + fn stop_accepts_string_and_string_array() { + let one_stop = r#"{"model": "test_model", "prompt": "hello", "stop": " The"}"#; + let request: CreateCompletionRequest = serde_json::from_str(one_stop).unwrap(); + + assert_eq!(request.stop, Some(Stop::String(" The".to_string()))); + + let many_stops = r#"{"model": "test_model", "prompt": "hello", "stop": ["A", "B"]}"#; + let request: CreateCompletionRequest = serde_json::from_str(many_stops).unwrap(); + + assert_eq!( + request.stop, + Some(Stop::StringArray(vec!["A".to_string(), "B".to_string()])) + ); + } + + #[test] + fn stop_token_id_display_string_remains_string_stop() { + let json = r#"{"model": "test_model", "prompt": [1, 2, 3], "stop": "token_id:576"}"#; + let request: CreateCompletionRequest = serde_json::from_str(json).unwrap(); + + assert_eq!(request.stop, Some(Stop::String("token_id:576".to_string()))); + + let json = r#"{"model": "test_model", "prompt": [1, 2, 3], "stop": ["token_id:576"]}"#; + let request: CreateCompletionRequest = serde_json::from_str(json).unwrap(); + + assert_eq!( + request.stop, + Some(Stop::StringArray(vec!["token_id:576".to_string()])) + ); + } + + #[test] + fn builder_accepts_upstream_stop_configuration() { + let upstream_stop = async_openai::types::chat::StopConfiguration::String("END".to_string()); + + let request = CreateCompletionRequestArgs::default() + .model("test_model") + .prompt(Prompt::String("hello".to_string())) + .stop(upstream_stop) + .build() + .unwrap(); + + assert_eq!(request.stop, Some(Stop::String("END".to_string()))); + } + + #[test] + fn stop_rejects_single_token_id() { + let json = r#"{"model": "test_model", "prompt": [1, 2, 3], "stop": 576}"#; + let result: Result = serde_json::from_str(json); + + assert!(result.is_err()); + } }