diff --git a/components/src/dynamo/frontend/sglang_prepost.py b/components/src/dynamo/frontend/sglang_prepost.py index 45e42c087d4c..5f34314d5617 100644 --- a/components/src/dynamo/frontend/sglang_prepost.py +++ b/components/src/dynamo/frontend/sglang_prepost.py @@ -1279,6 +1279,8 @@ def __init__( self._pending_stop_text = "" self._locally_finished = False self._local_stop_reason: str | None = None + # Decode context, decoded IDs, stop text and hidden stop-token count. + self._local_stop_boundary: tuple[list[int], list[int], str, int] | None = None # Keep a small, known-complete prompt suffix as decode context. Generated # tokens are promoted to context only after they decode without a @@ -1385,6 +1387,25 @@ def _incremental_decode( self._pending_decode_ids = [] return delta_text + def _count_tokens_through( + self, context_ids: list[int], token_ids: list[int], text: str + ) -> int: + """Count leading ``token_ids`` whose incremental decode covers ``text``. + + Once a decoded prefix covers ``text``, every longer one does, so bisect + instead of decoding every prefix. + """ + context_text = self._decode_ids(context_ids) + low, high = 0, len(token_ids) + while low < high: + count = (low + high) // 2 + decoded = self._decode_ids(context_ids + token_ids[:count]) + if decoded[len(context_text) :].startswith(text): + high = count + else: + low = count + 1 + return low + def _build_openai_logprobs( self, log_probs: list[float], @@ -1527,6 +1548,18 @@ def locally_finished(self) -> bool: def local_stop_reason(self) -> str | None: return self._local_stop_reason + def post_stop_token_count(self) -> int: + """Count the latest generated IDs that follow a local stop string. + + The token completing the stop is kept. The count spans IDs withheld by + incremental decoding and is computed on demand by re-decoding them. + """ + if self._local_stop_boundary is None: + return 0 + context_ids, decoded_ids, stop_text, hidden_count = self._local_stop_boundary + kept = self._count_tokens_through(context_ids, decoded_ids, stop_text) + return len(decoded_ids) - kept + hidden_count + @property def has_pending_stop_text(self) -> bool: return bool(self._pending_stop_text) @@ -1571,8 +1604,10 @@ def _filter_stop_string_delta( text: str, finish_reason: str | None, stop_reason: Any, - ) -> tuple[str, bool]: - text = self._pending_stop_text + text + ) -> tuple[str, int | None]: + """Return visible text and, after a local stop, the stop's end in ``text``.""" + held_text = self._pending_stop_text + text = held_text + text self._pending_stop_text = "" match = self._find_stop_string(text, stop_reason) @@ -1587,16 +1622,17 @@ def _filter_stop_string_delta( del self._pending_logprobs_content[-suppressed_count:] self._locally_finished = True self._local_stop_reason = matched_stop_string - return text[:retained_end], True + stop_end = match_index + len(matched_stop_string) - len(held_text) + return text[:retained_end], stop_end if finish_reason or not text or not self._stop_strings: - return text, False + return text, None pending_len = trailing_stop_prefix_len(text, self._stop_strings) if pending_len: self._pending_stop_text = text[-pending_len:] - return text[:-pending_len], False - return text, False + return text[:-pending_len], None + return text, None def _parse_reasoning_delta( self, delta_text: str, finish_reason: str | None @@ -1669,33 +1705,45 @@ def process_output(self, engine_response: dict[str, Any]) -> dict[str, Any] | No stop_terminated = engine_response.get( "stop_terminated", finish_reason == "stop" ) + generated_ids = token_ids if stop_terminated: - raw_token_count = len(token_ids) token_ids = self._strip_matched_stop_token_ids(list(token_ids), stop_reason) - retained_token_count = len(token_ids) - if log_probs is not None and len(log_probs) == raw_token_count: - log_probs = log_probs[:retained_token_count] - if top_logprobs is not None and len(top_logprobs) == raw_token_count: - top_logprobs = top_logprobs[:retained_token_count] - delta_text = ( + context_ids = self._decode_context_ids + held_ids = list(self._pending_decode_ids) + decoded_text = ( self._incremental_decode(token_ids, flush=finish_reason is not None) if token_ids or finish_reason is not None else "" ) - openai_logprobs = None + # A hidden stop token is generated output: score it, but do not display it. + stop_logprobs: list[dict[str, Any]] = [] if log_probs is not None: + scored_ids = ( + generated_ids if len(log_probs) == len(generated_ids) else token_ids + ) openai_logprobs = self._build_openai_logprobs( - log_probs, top_logprobs, token_ids + log_probs, top_logprobs, scored_ids ) if openai_logprobs is not None: - self._pending_logprobs_content.extend(openai_logprobs["content"]) + content = openai_logprobs["content"] + self._pending_logprobs_content.extend(content[: len(token_ids)]) + stop_logprobs = content[len(token_ids) :] self._logprob_context_ids = (self._logprob_context_ids + token_ids)[-4:] - delta_text, locally_finished = self._filter_stop_string_delta( - delta_text, finish_reason, stop_reason + delta_text, stop_end = self._filter_stop_string_delta( + decoded_text, finish_reason, stop_reason ) - if locally_finished: + if stop_end is None: + # Queued after stop-string suppression, which matches display text. + self._pending_logprobs_content.extend(stop_logprobs) + else: finish_reason = "stop" + self._local_stop_boundary = ( + context_ids, + held_ids + token_ids, + decoded_text[:stop_end], + len(generated_ids) - len(token_ids), + ) if self._fast_plain_text: if delta_text: diff --git a/components/src/dynamo/frontend/sglang_processor.py b/components/src/dynamo/frontend/sglang_processor.py index 290a23588bf3..fee11389ecc7 100644 --- a/components/src/dynamo/frontend/sglang_processor.py +++ b/components/src/dynamo/frontend/sglang_processor.py @@ -832,6 +832,10 @@ async def _generate_and_stream( stream_interval = self.stream_interval prompt_logprobs = PromptLogprobsAccumulator(request) logprobs_count = _request_logprobs_count(request) + include_completion_ids = nvext_extra_field_requested( + request, "completion_token_ids" + ) + include_prompt_ids = nvext_extra_field_requested(request, "prompt_token_ids") try: dynamo_stream = await self.routed_engine.generate( @@ -845,6 +849,8 @@ async def _generate_and_stream( pending_token_ids: list[int] = [] pending_log_probs: list[float] | None = None pending_top_logprobs: list[list[dict[str, Any]]] | None = None + # Requested generated IDs wait here while parsing emits no chunk. + completion_ids: list[int] = [] latest_usage: dict[str, Any] = {} pending_terminal: dict[str, Any] | None = None first_chunk = True @@ -868,6 +874,7 @@ def flush_pending( nonlocal pending_token_ids nonlocal pending_log_probs nonlocal pending_top_logprobs + nonlocal completion_ids nonlocal latest_usage nonlocal first_chunk nonlocal post_proc_total_ms @@ -889,6 +896,12 @@ def flush_pending( t_pp0 = time.monotonic() choice = post.process_output(mapped_response) + if include_completion_ids: + completion_ids.extend(pending_token_ids) + if post.locally_finished: + # IDs that decoding withheld before the stop are still queued. + dropped = post.post_stop_token_count() + del completion_ids[len(completion_ids) - dropped :] if post.locally_finished and not latest_usage: latest_usage = { @@ -935,6 +948,11 @@ def flush_pending( request, "engine_data" ): response_nvext["engine_data"] = engine_data + if include_completion_ids: + response_nvext["completion_token_ids"] = completion_ids + completion_ids = [] + if include_prompt_ids and (finish_reason or post.locally_finished): + response_nvext["prompt_token_ids"] = tokens if response_nvext: dynamo_out["nvext"] = response_nvext 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 5ca2f681974d..f611da2f4a32 100644 --- a/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py +++ b/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py @@ -4757,7 +4757,9 @@ async def collect(): assert chunks[1]["nvext"]["stop_reason"] == "END" assert "prompt_logprobs" not in chunks[1]["choices"][0] - def _run_logprob_stream(self, items, request=None): + def _run_logprob_stream( + self, items, request=None, eos_token_ids=None, prompt_ids=() + ): if request is None: request = {"model": "test-model", "logprobs": True} processor = SglangProcessor( @@ -4772,19 +4774,61 @@ def _run_logprob_stream(self, items, request=None): tokenizer=self.ByteTokenizer(), tool_call_parser=None, reasoning_parser=None, + eos_token_ids=eos_token_ids, ) async def collect(): return [ item["data"] async for item in processor._generate_and_stream( - "req-logprobs", request, {}, [], post + "req-logprobs", request, {}, list(prompt_ids), post ) if "data" in item ] return asyncio.run(collect()) + def test_hidden_eos_keeps_score_and_requested_ids(self): + """All 48 generated tokens keep a score and ID; only EOS text stays hidden.""" + eos = 0 + generated = [*b"x" * 47, eos] + logprobs = [-float(position) for position in range(1, len(generated) + 1)] + items = [ + {"token_ids": [token_id], "log_probs": [logprob]} + for token_id, logprob in zip(generated, logprobs) + ] + items[-1]["finish_reason"] = "stop" + + chunks = self._run_logprob_stream( + items, + request={ + "model": "test-model", + "logprobs": True, + "nvext": {"extra_fields": ["completion_token_ids", "prompt_token_ids"]}, + }, + eos_token_ids=[eos], + prompt_ids=[1, 2, 3], + ) + + choices = [chunk["choices"][0] for chunk in chunks] + content = "".join(choice["delta"].get("content", "") for choice in choices) + assert content == "x" * 47 + assert choices[-1]["finish_reason"] == "stop" + scores = [ + entry["logprob"] + for choice in choices + for entry in choice["logprobs"]["content"] + ] + assert scores == logprobs + completion_ids = [ + token_id + for chunk in chunks + for token_id in chunk["nvext"]["completion_token_ids"] + ] + assert completion_ids == generated + prompt_ids = [chunk["nvext"].get("prompt_token_ids") for chunk in chunks] + assert prompt_ids == [None] * (len(chunks) - 1) + [[1, 2, 3]] + @pytest.mark.parametrize( "prompt_count,logprobs,top_logprobs,expected_top_count", [ diff --git a/docs/fern/pages/use-cases/tool-calling-and-reasoning/troubleshooting-tool-calls.md b/docs/fern/pages/use-cases/tool-calling-and-reasoning/troubleshooting-tool-calls.md index 455b2f2aae38..1743a00bf3a4 100644 --- a/docs/fern/pages/use-cases/tool-calling-and-reasoning/troubleshooting-tool-calls.md +++ b/docs/fern/pages/use-cases/tool-calling-and-reasoning/troubleshooting-tool-calls.md @@ -1,5 +1,5 @@ --- -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 title: Troubleshooting Tool Calls subtitle: Capture raw model output with logprobs so issues can be localized @@ -19,7 +19,13 @@ what the response will look like, so you can capture and share useful diagnostic info. > [!IMPORTANT] -> Recipe applies to non-streaming requests against Dynamo's OpenAI `/v1/chat/completions` endpoint. For multi-channel reasoning models (`harmony`, `kimi_k2`, `kimi_k25`, `gemma4`), the recipe recovers only the assistant-content channel; the reasoning channel is not surfaced in `logprobs.content`. +> Recipe applies to non-streaming requests against Dynamo's OpenAI `/v1/chat/completions` endpoint. +> On the V2 parser route (see [Parser Generation Routing](#parser-generation-routing)), +> `logprobs.content` lists every token entry the backend returns, in generation order, including +> reasoning, tool-call, and control tokens that the parsed message does not show. Streamed responses +> carry each backend chunk's entries once, on one chunk for that choice, so the entries need not line +> up with the text in `reasoning_content`, `content`, and `tool_calls`. This note does not describe +> the V1 route, where late entries can arrive in a final empty-`delta` chunk after `finish_reason`. ## The request @@ -88,11 +94,22 @@ engine's raw token stream: } ``` -Each entry in `logprobs.content` is one generated token with its exact UTF-8 -`bytes`. Concatenating those bytes in order reconstructs the raw model -output, before any tool-call parser touched it. That is the key piece for -triage: it tells us what the model actually produced, separately from what -the parser made of it. +Each entry in `logprobs.content` is one generated token and keeps the score +the backend returned for it. The `token` text and `bytes` are decoded +representations: an end-of-sequence token can appear as its literal text and +bytes, such as `<|im_end|>` with `--dyn-chat-processor sglang`, or as empty +text and no bytes, and `return_tokens_as_token_ids` can replace the text with a +`token_id:` label. Without that option, concatenating the `bytes` in +order approximates the model output before any tool-call parser touched it. +It can include control-token literals such as a trailing `<|im_end|>`, omits +tokens that decode to nothing, and, because each token can be decoded on its +own, can lose spacing that the tokenizer adds only in +context. That is the key piece for triage: it tells us what the model +actually produced, separately from what the parser made of it. For an exact +replay, also request the generated token IDs with +`"nvext": {"extra_fields": ["completion_token_ids"]}` where the response +supports them (see +[NVIDIA Request Extensions](../../developer-guide/additional-resources/nvidia-request-extensions-nvext.md)). ## What to include when reporting an issue @@ -108,8 +125,9 @@ Share these four things in the bug report or issue thread: value if set. With those four pieces, the Dynamo team can usually localize the bug -without standing up your model. The team will reconstruct the raw stream -from the `bytes` arrays and compare it against `message.content` and +without standing up your model. The team will rebuild the raw stream from +the per-token entries, using exact token IDs over decoded `bytes` where the +response provides them, and compare it against `message.content` and `message.tool_calls` to decide whether the issue is in the model output, the parser configuration, or the parser logic. diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index 8b7c02d8029c..329e7c878dd4 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -6597,7 +6597,9 @@ impl OpenAIPreprocessor { /// /// The parser can buffer and rewrite several input chunks before it emits an /// output. Completion token IDs therefore describe the ordered buffered - /// group, not the rewritten text in one output delta. + /// group, not the rewritten text in one output delta. Chosen-token logprobs + /// follow the same rule per choice: each choice's next emitted delta carries + /// every score received since its previous one. pub fn apply_tool_calling_jail( tool_call_parser: Option, tool_choice: Option, @@ -6609,6 +6611,7 @@ impl OpenAIPreprocessor { where S: Stream> + Send + 'static, { + use crate::protocols::openai::chat_completions::unified_parser::empty_choice; use dynamo_parsers::tool_calling::jail::{ Annotated as JailAnnotated, apply_tool_calling_jail, apply_tool_calling_jail_with_guided_streaming, @@ -6633,12 +6636,18 @@ impl OpenAIPreprocessor { // `nvext` uses the unary aggregator's merge rules: completion token IDs // are appended, while the latest supplied value wins for every other // top-level field. `engine_data` is replaced as one complete value. + // + // Choice logprobs are copied here per choice index. The jail still sees + // them, so its text decisions are unchanged, but it splits them by + // character count and can drop or repeat them around buffered markers + // and terminal chunks. Its copies are therefore replaced on the way out. #[derive(Default)] struct PendingDynamoMetadata { metrics_template: Option, chunk_tokens: usize, nvext: Option, prompt_logprobs: Option>, + logprobs: std::collections::BTreeMap, response_template: Option, } let pending = Arc::new(Mutex::new(PendingDynamoMetadata::default())); @@ -6693,7 +6702,10 @@ impl OpenAIPreprocessor { // dynamo `Annotated` -> jail `Annotated` (buffer Dynamo metadata) let jail_input = stream.map(move |mut a| { let has_metadata = a.data.as_ref().is_some_and(|nv| { - nv.llm_metrics.is_some() || nv.nvext.is_some() || nv.prompt_logprobs.is_some() + nv.llm_metrics.is_some() + || nv.nvext.is_some() + || nv.prompt_logprobs.is_some() + || nv.inner.choices.iter().any(|choice| choice.logprobs.is_some()) }); if has_metadata { let mut p = pending_in @@ -6722,6 +6734,28 @@ impl OpenAIPreprocessor { if let Some(prompt_logprobs) = nv.prompt_logprobs.take() { p.prompt_logprobs = Some(prompt_logprobs); } + for choice in &nv.inner.choices { + let Some(logprobs) = &choice.logprobs else { + continue; + }; + let empty = dynamo_protocols::types::ChatChoiceLogprobs { + content: None, + refusal: None, + }; + let queued = p.logprobs.entry(choice.index).or_insert(empty); + if let Some(content) = &logprobs.content { + queued + .content + .get_or_insert(Vec::new()) + .extend_from_slice(content); + } + if let Some(refusal) = &logprobs.refusal { + queued + .refusal + .get_or_insert(Vec::new()) + .extend_from_slice(refusal); + } + } } } if is_glm47 && let Some(data) = &a.data { @@ -6816,7 +6850,7 @@ impl OpenAIPreprocessor { }; let pending_out = Arc::clone(&pending); let pending_eof = Arc::clone(&pending); - let jailed_output = jailed.flat_map(move |a| { + let jailed_output = jailed.flat_map(move |mut a| { debug_assert!( a.error.is_none(), "dynamo-parsers must not construct errors" @@ -6827,10 +6861,15 @@ impl OpenAIPreprocessor { let has_choices = a.data.as_ref().is_some_and(|data| !data.choices.is_empty()); let is_payload_usage = a.event.as_deref() == Some(ANNOTATION_PAYLOAD_USAGE); let mut prompt_logprobs = None; - let (llm_metrics, nvext) = a.data.as_ref().map_or((None, None), |_| { + let (llm_metrics, nvext) = a.data.as_mut().map_or((None, None), |data| { let mut p = pending_out .lock() .expect("jail Dynamo metadata buffer poisoned"); + // The jail yields everything for one input before it reads the + // next, so a queue holds only scores up to this output. + for choice in &mut data.choices { + choice.logprobs = p.logprobs.remove(&choice.index); + } let chunk_tokens = p.chunk_tokens; p.chunk_tokens = 0; let metrics = p.metrics_template.take().map(|mut metrics| { @@ -6959,6 +6998,7 @@ impl OpenAIPreprocessor { p.chunk_tokens = 0; p.nvext = None; p.prompt_logprobs = None; + p.logprobs.clear(); p.response_template = None; } yield error; @@ -6977,22 +7017,39 @@ impl OpenAIPreprocessor { }); let nvext = p.nvext.take(); let prompt_logprobs = p.prompt_logprobs.take(); - if llm_metrics.is_none() && nvext.is_none() && prompt_logprobs.is_none() { + // Scores for text the jail consumed after a choice's last delta. + // Each carrier adds only the scores, never content or a finish. + let choices: Vec<_> = std::mem::take(&mut p.logprobs) + .into_iter() + .map(|(index, logprobs)| { + let mut choice = empty_choice(index); + choice.logprobs = Some(logprobs); + choice + }) + .collect(); + if llm_metrics.is_none() + && nvext.is_none() + && prompt_logprobs.is_none() + && choices.is_empty() + { None } else { - p.response_template.take().map(|inner| Annotated { - data: Some(NvCreateChatCompletionStreamResponse { - message_boundary: None, - inner, - nvext, - prompt_logprobs, - llm_metrics, - tool_call_completion: Vec::new(), - }), - id: None, - event: None, - comment: None, - error: None, + p.response_template.take().map(|mut inner| { + inner.choices = choices; + Annotated { + data: Some(NvCreateChatCompletionStreamResponse { + message_boundary: None, + inner, + nvext, + prompt_logprobs, + llm_metrics, + tool_call_completion: Vec::new(), + }), + id: None, + event: None, + comment: None, + error: None, + } }) } }; @@ -12420,6 +12477,53 @@ responses_tools: None, ); } + /// `required` starts the jail active. With no text to flush, it swallows a contentless + /// score, so only the end-of-stream choice can deliver it after the finish. + #[tokio::test] + async fn legacy_jail_delivers_swallowed_eos_score_after_finish() { + let score = |token: &str| { + Some( + serde_json::from_value(serde_json::json!({ + "content": [{"token": token, "logprob": -0.5, "bytes": [], "top_logprobs": []}], + "refusal": null + })) + .unwrap(), + ) + }; + let mut finish = terminal_chat_stream_chunk(); + let choice = &mut finish.data.as_mut().unwrap().inner.choices[0]; + choice.delta.reasoning_content = Some("thought".to_string()); + choice.logprobs = score("thought"); + let mut eos = chat_stream_chunk(0, None); + let choice = &mut eos.data.as_mut().unwrap().inner.choices[0]; + choice.delta.content = None; + choice.logprobs = score("eos"); + + let output = OpenAIPreprocessor::apply_tool_calling_jail( + Some("hermes".to_string()), + Some(ChatCompletionToolChoiceOption::Required), + None, + false, + false, + stream::iter(vec![finish, eos]), + ) + .collect::>() + .await; + + let choices: Vec<_> = output + .iter() + .flat_map(|response| response.data.iter()) + .flat_map(|data| data.inner.choices.iter()) + .collect(); + assert_eq!(choices.len(), 2); + assert_eq!(choices[0].finish_reason, Some(FinishReason::Stop)); + assert_eq!(choices[0].logprobs, score("thought")); + let mut score_only = + crate::protocols::openai::chat_completions::unified_parser::empty_choice(0); + score_only.logprobs = score("eos"); + assert_eq!(choices[1], &score_only); + } + fn kimi_k3_reasoning_chunk(reasoning: &str) -> Annotated { let mut chunk = chat_stream_chunk(0, Some(Role::Assistant)); let choice = &mut chunk.data.as_mut().unwrap().inner.choices[0]; diff --git a/lib/llm/src/protocols/openai/chat_completions/unified_parser.rs b/lib/llm/src/protocols/openai/chat_completions/unified_parser.rs index 1b4ff5e12ab2..76f00af44a54 100644 --- a/lib/llm/src/protocols/openai/chat_completions/unified_parser.rs +++ b/lib/llm/src/protocols/openai/chat_completions/unified_parser.rs @@ -45,9 +45,10 @@ use dynamo_parsers_v2::{ UnifiedParserStartingState, UnifiedToolOutputMode, create_unified_parser_for_family, }; use dynamo_protocols::types::{ - ChatChoiceStream, ChatCompletionMessageContent, ChatCompletionMessageToolCall, - ChatCompletionMessageToolCallChunk, ChatCompletionStreamResponseDelta, FinishReason, - FunctionCall, FunctionCallStream, FunctionType, + ChatChoiceLogprobs, ChatChoiceStream, ChatCompletionMessageContent, + ChatCompletionMessageToolCall, ChatCompletionMessageToolCallChunk, + ChatCompletionStreamResponseDelta, FinishReason, FunctionCall, FunctionCallStream, + FunctionType, }; use dynamo_runtime::protocols::annotated::Annotated; use futures::{Stream, StreamExt}; @@ -1438,9 +1439,8 @@ pub(crate) fn fanout_llm_metrics_position(choices: &[ChatChoiceStream]) -> Optio /// An empty streaming choice for `index`, used as the base every emitted chunk is /// filled in from. /// -/// `logprobs` is dropped on purpose: once parsing rewrites a choice, the emitted text -/// no longer lines up token-for-token with the backend's raw stream, so per-token -/// logprobs would be attached to the wrong characters. +/// It carries no `logprobs`: `choices_for` moves a source chunk's raw-token logprobs +/// onto exactly one of the choices parsed from it. pub(crate) fn empty_choice(index: u32) -> ChatChoiceStream { #[allow(deprecated)] ChatChoiceStream { @@ -1481,6 +1481,9 @@ struct QwenXmlStreamState { pending_opener: String, /// Must follow recovered opener text on mismatch and precede completed call events. deferred_reasoning: Vec, + /// Logprobs of chunks generated after the latest deferred reasoning, held until it + /// is released. + deferred_logprobs: Option, /// Entries in event/output lists, not XML character offsets: these preserve /// deferred reasoning's original response metadata through projection. deferred_reasoning_event_indices: HashSet, @@ -2143,11 +2146,23 @@ impl ChoiceState { } } - fn queue_qwen_opener_reasoning(&mut self, reasoning: String) { + /// Scores held since the previous queued reasoning were generated before this + /// one, so they lead its source envelope's `logprobs`. + fn queue_qwen_opener_reasoning( + &mut self, + reasoning: String, + source_logprobs: &mut Option, + ) { let NativeGrammarState::QwenXml(qwen) = &mut self.native_grammar_state else { return; }; qwen.deferred_reasoning.push(reasoning); + if let Some(mut held) = qwen.deferred_logprobs.take() { + if let Some(source) = source_logprobs.take() { + extend_logprobs(&mut held, source); + } + *source_logprobs = Some(held); + } } fn prepend_pending_qwen_reasoning(&mut self, events: &mut Vec) { @@ -2367,7 +2382,9 @@ impl ChoiceState { /// /// `role` / `refusal` ride on the first choice owned by the current source /// response; queued reasoning keeps its earlier source envelope. The terminating - /// `finish_reason` stays after every emitted delta. + /// `finish_reason` stays after every emitted delta. The source chunk's `logprobs` + /// move once to its last choice, waiting for any queued reasoning to go first: + /// they score generated tokens, not parsed text. pub(crate) fn choices_for( &mut self, original: &ChatChoiceStream, @@ -2520,6 +2537,33 @@ impl ChoiceState { choices.push(choice); } + // Queued reasoning arrived before this chunk, so this chunk's logprobs must + // not ride ahead of it: hold them while it is queued, then follow it. + let mut logprobs = original.logprobs.take(); + if let NativeGrammarState::QwenXml(qwen) = &mut self.native_grammar_state { + if let Some(mut held) = qwen.deferred_logprobs.take() { + if let Some(source) = logprobs.take() { + extend_logprobs(&mut held, source); + } + logprobs = Some(held); + } + if !qwen.deferred_reasoning.is_empty() { + qwen.deferred_logprobs = logprobs.take(); + } + } + if let Some(logprobs) = logprobs { + let last_is_queued = choices + .len() + .checked_sub(1) + .is_some_and(|last| self.is_deferred_qwen_xml_reasoning_choice(last)); + if choices.is_empty() || last_is_queued { + choices.push(empty_choice(index)); + } + if let Some(choice) = choices.last_mut() { + choice.logprobs = Some(logprobs); + } + } + choices } @@ -2934,6 +2978,15 @@ enum FinishedStreamChunk { Source(Annotated>), } +fn extend_logprobs(held: &mut ChatChoiceLogprobs, later: ChatChoiceLogprobs) { + if let Some(content) = later.content { + held.content.get_or_insert_default().extend(content); + } + if let Some(refusal) = later.refusal { + held.refusal.get_or_insert_default().extend(refusal); + } +} + fn take_qwen_reasoning_response( record: &mut ChoiceRecord, ) -> Annotated> { @@ -3600,7 +3653,10 @@ where states .get_mut(&original.index) .expect("preserved parser state must still exist") - .queue_qwen_opener_reasoning(reasoning.clone()); + .queue_qwen_opener_reasoning( + reasoning.clone(), + &mut original.logprobs, + ); let mut source_data = chat.clone(); source_data.inner.choices.clear(); let keep_source_metadata = !qwen_source_metadata_deferred; @@ -10977,6 +11033,63 @@ pub(crate) mod tests { } } } + + /// Scores stay in generation order while raw text and reasoning alternate behind a + /// split Qwen opener. + #[tokio::test] + async fn qwen_queued_reasoning_keeps_source_logprob_order() { + fn scored( + mut response: Annotated, + token: &str, + ) -> Annotated { + response.data.as_mut().unwrap().inner.choices[0].logprobs = Some( + serde_json::from_value(serde_json::json!({ + "content": [{"token": token, "logprob": -0.5, "bytes": [], "top_logprobs": []}], + "refusal": null + })) + .unwrap(), + ); + response + } + + fn reasoning(text: &str, token: &str) -> Annotated { + let mut source = scored(chunk("", false), token); + let choice = &mut source.data.as_mut().unwrap().inner.choices[0]; + choice.delta.content = None; + choice.delta.reasoning_content = Some(text.to_string()); + source + } + + let completed_call = concat!( + "ll>Paris", + "" + ); + let responses = apply_stream( + stream::iter(vec![ + scored(chunk(">() + .await; + + let tokens: Vec<_> = collect_choices(&responses) + .into_iter() + .filter_map(|choice| choice.logprobs.as_ref()?.content.as_ref()) + .flatten() + .map(|entry| entry.token.as_str()) + .collect(); + assert_eq!(tokens, ["A", "B", "C", "D", "E"]); + } + #[tokio::test] #[ignore = "Parser 0.7.17 missing-parameter-close recovery delays text across external reasoning; run explicitly to measure ordering"] async fn qwen_recovered_parameter_close_restores_opener_ordering_at_every_split() { diff --git a/lib/llm/src/protocols/openai/responses/stream_converter.rs b/lib/llm/src/protocols/openai/responses/stream_converter.rs index d20287108e62..bf4285da85d8 100644 --- a/lib/llm/src/protocols/openai/responses/stream_converter.rs +++ b/lib/llm/src/protocols/openai/responses/stream_converter.rs @@ -1841,6 +1841,38 @@ mod tests { ); } + /// The V1 tool jail can send late scores on an empty-delta choice after `finish_reason`. + /// Responses carries no chat logprobs, so that choice must not change the output. + #[test] + fn post_finish_score_only_choice_emits_no_events() { + let mut score_only = finish_chunk(FinishReason::Stop); + let choice = &mut score_only.inner.choices[0]; + *choice = crate::protocols::openai::chat_completions::unified_parser::empty_choice(0); + choice.logprobs = Some( + serde_json::from_value(serde_json::json!({ + "content": [{"token": "eos", "logprob": -0.5, "bytes": [], "top_logprobs": []}], + "refusal": null + })) + .unwrap(), + ); + + let mut conv = ResponseStreamConverter::new("test-model".into(), default_params()); + conv.process_chunk(&text_chunk("hi")); + conv.process_chunk(&finish_chunk(FinishReason::Stop)); + + assert!(conv.process_chunk(&score_only).is_empty()); + assert_eq!( + event_types(&conv.emit_end_events()), + [ + "response.output_text.done", + "response.content_part.done", + "response.output_item.done", + "response.completed", + ] + ); + assert_eq!(conv.completed_output().len(), 1); + } + /// Parseable arguments remain open until an explicit tool-call finish reason. #[test] fn test_complete_tool_call_closes_on_finish_reason() {