fix(stream): align logprobs with retained parser segments

deepseek-recipe/src/stream/mod.rs

@@ -59,9 +59,11 @@ impl std::error::Error for StreamError {}
 /// `ToolCallArgumentsDelta` chunks; output truncated inside a tool call ends
 /// that sequence early.
 ///
-/// `logprobs` holds the log probabilities of the tokens of the decoded text the
+/// `logprobs` holds original token probabilities only for complete tokens
-/// content comes from. Content split from one decoded text carries the log
+/// inside the emitted source-text segment. Tokens crossing a parser segment
-/// probabilities of the whole text.
+/// boundary are omitted, as are probabilities for replaced text (such as JSON
+/// padding). If vocabulary bytes do not match decoded text, probabilities for
+/// that decoded piece are omitted. No token probability is duplicated.
 #[derive(Debug, Clone, PartialEq)]
 pub enum OutputChunk {
     /// Output begins.

deepseek-recipe/src/stream/processor.rs

@@ -150,7 +150,7 @@ where
                                 .map(|decoded| OutputPiece {
                                     content: decoded.content,
                                     token_count: decoded.token_count,
-                                    logprobs: decoded.logprobs.map(Arc::from),
+                                    logprobs: decoded.logprobs,
                                     kind: if decoded.is_special_token {
                                         ChunkKind::SpecialToken
                                     } else {
@@ -237,15 +237,61 @@ where
 struct OutputPiece {
     content: String,
     token_count: usize,
-    logprobs: Option<Arc<[TokenLogprob]>>,
+    logprobs: Option<Vec<TokenLogprob>>,
     kind: ChunkKind,
 }
 
-/// Output text waiting for its parsing actions, with the log probabilities of
+/// Output text waiting for its parsing actions and the remaining token probabilities.
-/// the text it was split from.
 struct BufferedChunk {
     content: String,
-    logprobs: Option<Arc<[TokenLogprob]>>,
+    logprobs: Option<BufferedLogprobs>,
+}
+
+/// Consume each token probability once, using its original byte range.
+struct BufferedLogprobs {
+    tokens: VecDeque<TokenLogprob>,
+    /// Bytes already consumed from the front token. A split token is omitted.
+    offset: usize,
+}
+
+impl BufferedLogprobs {
+    fn new(content: &str, logprobs: Vec<TokenLogprob>) -> Option<Self> {
+        // A custom decoder can transform vocabulary bytes. Without matching
+        // bytes there is no reliable alignment; do not invent probabilities.
+        if logprobs.iter().any(|token| token.bytes.is_empty())
+            || !logprobs
+                .iter()
+                .flat_map(|token| token.bytes.iter().copied())
+                .eq(content.bytes())
+        {
+            return None;
+        }
+        Some(Self {
+            tokens: logprobs.into(),
+            offset: 0,
+        })
+    }
+
+    fn take(&mut self, mut byte_len: usize) -> Option<Arc<[TokenLogprob]>> {
+        let mut taken = Vec::new();
+        while let Some(token) = self.tokens.front() {
+            let remaining = token.bytes.len() - self.offset;
+            if remaining > byte_len {
+                self.offset += byte_len;
+                break;
+            }
+            let token = self.tokens.pop_front().expect("the front token exists");
+            if self.offset == 0 {
+                taken.push(token);
+            }
+            self.offset = 0;
+            byte_len -= remaining;
+            if byte_len == 0 {
+                break;
+            }
+        }
+        (!taken.is_empty()).then(|| taken.into())
+    }
 }
 
 /// Output text that the parser has not fully assigned yet, and the state
@@ -280,7 +326,8 @@ impl PendingOutput {
         }
     }
 
-    fn push(&mut self, content: String, logprobs: Option<Arc<[TokenLogprob]>>) {
+    fn push(&mut self, content: String, logprobs: Option<Vec<TokenLogprob>>) {
+        let logprobs = logprobs.and_then(|logprobs| BufferedLogprobs::new(&content, logprobs));
         self.chunks.push_back(BufferedChunk { content, logprobs });
     }
 
@@ -290,20 +337,24 @@ impl PendingOutput {
 
     /// Take the first `segment.byte_len` bytes of the front buffered chunk and
     /// convert them by the segment's action. `segment.byte_len` must not exceed
-    /// the length of the front chunk. Text split off from a chunk keeps the
+    /// the length of the front chunk. Only complete tokens inside the taken
-    /// chunk's log probabilities.
+    /// byte range keep their probabilities; a token crossing the split is omitted.
     fn take_piece(&mut self, segment: OutputActionSegment) -> Option<OutputChunk> {
         let BufferedChunk {
             content: mut piece,
-            logprobs,
+            mut logprobs,
         } = self.chunks.pop_front()?;
+        let piece_logprobs = logprobs
+            .as_mut()
+            .and_then(|logprobs| logprobs.take(segment.byte_len));
         if segment.byte_len < piece.len() {
             let remaining = piece.split_off(segment.byte_len);
             self.chunks.push_front(BufferedChunk {
                 content: remaining,
-                logprobs: logprobs.clone(),
+                logprobs,
             });
         }
+        let logprobs = piece_logprobs;
         let is_first_piece = self.previous_piece_action != segment.action;
         self.previous_piece_action = segment.action;
 
@@ -321,7 +372,7 @@ impl PendingOutput {
             }
             OutputAction::ToSpace => Some(OutputChunk::Content {
                 content: " ".to_string(),
-                logprobs,
+                logprobs: None,
             }),
             OutputAction::Skip | OutputAction::SkipInvalid { .. } => None,
             OutputAction::StopSequence => {

deepseek-recipe/tests/logprob_alignment.rs

@@ -0,0 +1,367 @@
+use deepseek_recipe::stream::{
+    ChunkGenerator, InferenceChunk, InferenceFinishReason, InferenceTokenLogprob, OutputChunk,
+    ParsingOptions, StreamProcessor, StreamTokenizer,
+};
+use tokio_stream::StreamExt;
+
+struct AsciiTokenizer;
+impl StreamTokenizer for AsciiTokenizer {
+    fn decode_ids(&self, ids: &[u32], _: bool) -> Result<String, String> {
+        Ok(ids.iter().map(|&id| char::from_u32(id).unwrap()).collect())
+    }
+    fn id_to_token(&self, id: u32) -> Option<String> {
+        char::from_u32(id).map(|ch| ch.to_string())
+    }
+}
+struct Passthrough;
+impl ChunkGenerator for Passthrough {
+    type Chunk = OutputChunk;
+    async fn generate(&mut self, chunk: OutputChunk) -> Vec<OutputChunk> {
+        vec![chunk]
+    }
+}
+async fn run(text: &str, batch: bool, stop: Vec<String>) -> (String, Vec<String>) {
+    let ids: Vec<u32> = text.chars().map(u32::from).collect();
+    let batches = if batch {
+        vec![ids]
+    } else {
+        ids.into_iter().map(|id| vec![id]).collect()
+    };
+    let mut input: Vec<InferenceChunk> = batches
+        .into_iter()
+        .map(|token_ids| {
+            let logprobs = token_ids
+                .iter()
+                .map(|_| InferenceTokenLogprob {
+                    logprob: -1.0,
+                    top_logprobs: vec![],
+                })
+                .collect();
+            InferenceChunk::Tokens {
+                token_ids,
+                logprobs: Some(logprobs),
+            }
+        })
+        .collect();
+    input.push(InferenceChunk::Finish {
+        finish_reason: InferenceFinishReason::Stop,
+    });
+    let processor = StreamProcessor::new(
+        Passthrough,
+        ParsingOptions {
+            stop_sequences: stop,
+            ..Default::default()
+        },
+    )
+    .with_tokenizer(AsciiTokenizer);
+    let stream = processor.into_protocol_stream(tokio_stream::iter(input));
+    let chunks: Vec<_> = std::pin::pin!(stream).collect().await;
+    let mut text = String::new();
+    let mut tokens = vec![];
+    for chunk in chunks {
+        if let OutputChunk::Content { content, logprobs } = chunk.unwrap() {
+            text.push_str(&content);
+            if let Some(entries) = logprobs {
+                tokens.extend(entries.iter().map(|entry| entry.token.clone()));
+            }
+        }
+    }
+    (text, tokens)
+}
+#[tokio::test]
+async fn stop_filtered_logprobs_do_not_depend_on_backend_batching() {
+    let batched = run("ABSTOPZ", true, vec!["STOP".into()]).await;
+    let single = run("ABSTOPZ", false, vec!["STOP".into()]).await;
+    assert_eq!(batched.0, "AB");
+    assert_eq!(batched, single);
+    assert_eq!(batched.1, vec!["A", "B"]);
+}
+#[tokio::test]
+async fn stripped_leading_newline_has_no_content_logprob() {
+    let batched = run("\nAB", true, vec![]).await;
+    let single = run("\nAB", false, vec![]).await;
+    assert_eq!(batched.0, "AB");
+    assert_eq!(batched, single);
+}
+
+#[tokio::test]
+async fn json_padding_and_labels_do_not_duplicate_logprobs() {
+    let text = "prefix{\"x\":1}\n```suffix";
+    let ids: Vec<u32> = text.chars().map(u32::from).collect();
+    for batch_size in 1..=ids.len() {
+        let input = ids
+            .chunks(batch_size)
+            .map(|ids| InferenceChunk::Tokens {
+                token_ids: ids.to_vec(),
+                logprobs: Some(
+                    ids.iter()
+                        .map(|_| InferenceTokenLogprob {
+                            logprob: -1.0,
+                            top_logprobs: vec![],
+                        })
+                        .collect(),
+                ),
+            })
+            .chain([InferenceChunk::Finish {
+                finish_reason: InferenceFinishReason::Stop,
+            }]);
+        let processor = StreamProcessor::new(
+            Passthrough,
+            ParsingOptions {
+                parse_json_output: true,
+                ..Default::default()
+            },
+        )
+        .with_tokenizer(AsciiTokenizer);
+        let stream = processor.into_protocol_stream(tokio_stream::iter(input));
+        let output: Vec<_> = std::pin::pin!(stream).collect().await;
+        let mut tokens = Vec::new();
+        for chunk in output {
+            if let OutputChunk::Content {
+                logprobs: Some(logprobs),
+                ..
+            } = chunk.unwrap()
+            {
+                tokens.extend(logprobs.iter().cloned());
+            }
+        }
+        assert_eq!(
+            tokens
+                .iter()
+                .map(|token| token.token.as_str())
+                .collect::<String>(),
+            "{\"x\":1}",
+            "batch size {batch_size}"
+        );
+    }
+}
+
+struct WholePieceTokenizer(&'static str);
+impl StreamTokenizer for WholePieceTokenizer {
+    fn decode_ids(&self, _: &[u32], _: bool) -> Result<String, String> {
+        Ok("ABSTOPZ".into())
+    }
+    fn id_to_token(&self, _: u32) -> Option<String> {
+        Some(self.0.into())
+    }
+}
+#[tokio::test]
+async fn a_token_crossing_the_stop_boundary_is_not_fabricated_or_repeated() {
+    let processor = StreamProcessor::new(
+        Passthrough,
+        ParsingOptions {
+            stop_sequences: vec!["STOP".into()],
+            ..Default::default()
+        },
+    )
+    .with_tokenizer(WholePieceTokenizer("ABSTOPZ"));
+    let input = tokio_stream::iter(vec![
+        InferenceChunk::Tokens {
+            token_ids: vec![1],
+            logprobs: Some(vec![InferenceTokenLogprob {
+                logprob: -0.25,
+                top_logprobs: vec![],
+            }]),
+        },
+        InferenceChunk::Finish {
+            finish_reason: InferenceFinishReason::Stop,
+        },
+    ]);
+    let output: Vec<_> = std::pin::pin!(processor.into_protocol_stream(input))
+        .collect()
+        .await;
+    let content: Vec<_> = output
+        .into_iter()
+        .map(Result::unwrap)
+        .filter_map(|chunk| {
+            if let OutputChunk::Content { content, logprobs } = chunk {
+                Some((content, logprobs))
+            } else {
+                None
+            }
+        })
+        .collect();
+    assert_eq!(content.len(), 1);
+    assert_eq!(content[0].0, "AB");
+    assert!(content[0].1.is_none());
+}
+
+#[tokio::test]
+async fn multibyte_token_logprobs_survive_stop_filtering_at_every_token_split() {
+    let tokenizer = tokenizers::Tokenizer::from_file(concat!(
+        env!("CARGO_MANIFEST_DIR"),
+        "/../static/tokenizers/v41/tokenizer.json"
+    ))
+    .unwrap();
+    let text = "你好🤭STOPsuffix";
+    let ids = tokenizer.encode(text, false).unwrap().get_ids().to_vec();
+    for split in 0..=ids.len() {
+        let input = [&ids[..split], &ids[split..]]
+            .into_iter()
+            .map(|ids| InferenceChunk::Tokens {
+                token_ids: ids.to_vec(),
+                logprobs: Some(
+                    ids.iter()
+                        .map(|_| InferenceTokenLogprob {
+                            logprob: -1.0,
+                            top_logprobs: vec![],
+                        })
+                        .collect(),
+                ),
+            })
+            .chain([InferenceChunk::Finish {
+                finish_reason: InferenceFinishReason::Stop,
+            }]);
+        let processor = StreamProcessor::new(
+            Passthrough,
+            ParsingOptions {
+                stop_sequences: vec!["STOP".into()],
+                ..Default::default()
+            },
+        )
+        .with_tokenizer(tokenizer.clone());
+        let output: Vec<_> =
+            std::pin::pin!(processor.into_protocol_stream(tokio_stream::iter(input)))
+                .collect()
+                .await;
+        let mut content = String::new();
+        let mut bytes = Vec::new();
+        for chunk in output {
+            if let OutputChunk::Content {
+                content: part,
+                logprobs,
+            } = chunk.unwrap()
+            {
+                content.push_str(&part);
+                if let Some(logprobs) = logprobs {
+                    bytes.extend(
+                        logprobs
+                            .iter()
+                            .flat_map(|token| token.bytes.iter().copied()),
+                    );
+                }
+            }
+        }
+        assert_eq!(content, "你好🤭", "split {split}");
+        assert_eq!(bytes, "你好🤭".as_bytes(), "split {split}");
+    }
+}
+
+#[tokio::test]
+async fn vocabulary_bytes_that_do_not_match_decoded_text_are_not_reported() {
+    let processor = StreamProcessor::new(Passthrough, ParsingOptions::default())
+        .with_tokenizer(WholePieceTokenizer("different bytes"));
+    let input = tokio_stream::iter(vec![
+        InferenceChunk::Tokens {
+            token_ids: vec![1],
+            logprobs: Some(vec![InferenceTokenLogprob {
+                logprob: -0.25,
+                top_logprobs: vec![],
+            }]),
+        },
+        InferenceChunk::Finish {
+            finish_reason: InferenceFinishReason::Stop,
+        },
+    ]);
+    let output: Vec<_> = std::pin::pin!(processor.into_protocol_stream(input))
+        .collect()
+        .await;
+    for chunk in output {
+        if let OutputChunk::Content { content, logprobs } = chunk.unwrap() {
+            assert_eq!(content, "ABSTOPZ");
+            assert!(logprobs.is_none());
+        }
+    }
+}
+
+async fn public_chunks<G: ChunkGenerator>(
+    generator: G,
+    options: ParsingOptions,
+) -> Vec<serde_json::Value>
+where
+    G::Chunk: serde::Serialize + Send,
+{
+    let ids: Vec<u32> = "ABSTOPZ".chars().map(u32::from).collect();
+    let input = tokio_stream::iter(vec![
+        InferenceChunk::Tokens {
+            logprobs: Some(
+                ids.iter()
+                    .map(|_| InferenceTokenLogprob {
+                        logprob: -1.0,
+                        top_logprobs: vec![],
+                    })
+                    .collect(),
+            ),
+            token_ids: ids,
+        },
+        InferenceChunk::Finish {
+            finish_reason: InferenceFinishReason::Stop,
+        },
+    ]);
+    let processor = StreamProcessor::new(generator, options).with_tokenizer(AsciiTokenizer);
+    std::pin::pin!(processor.into_protocol_stream(input))
+        .map(|chunk| serde_json::to_value(chunk.unwrap()).unwrap())
+        .collect()
+        .await
+}
+#[tokio::test]
+async fn chat_completion_wire_logprobs_match_retained_content() {
+    use deepseek_recipe::openai::ChatCompletionRequest;
+    use deepseek_recipe::request::{ConversionOptions, ProtocolRequest};
+    let request: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
+        "model":"DeepSeek-V4.1-Flash", "messages":[{"role":"user","content":"Hello"}],
+        "thinking":{"type":"disabled"}, "logprobs":true, "stop":"STOP"
+    }))
+    .unwrap();
+    let converted = request.convert(ConversionOptions::default()).unwrap();
+    let generator = ChatCompletionRequest::chunk_generator(
+        &converted,
+        "test-id".into(),
+        "DeepSeek-V4.1-Flash".into(),
+    );
+    let chunks = public_chunks(generator, converted.parsing_options).await;
+    let text: String = chunks
+        .iter()
+        .filter_map(|chunk| chunk["choices"][0]["delta"]["content"].as_str())
+        .collect();
+    let tokens: String = chunks
+        .iter()
+        .filter_map(|chunk| chunk["choices"][0]["logprobs"]["content"].as_array())
+        .flatten()
+        .map(|entry| entry["token"].as_str().unwrap())
+        .collect();
+    assert_eq!(text, "AB");
+    assert_eq!(tokens, "AB");
+}
+#[tokio::test]
+async fn responses_wire_logprobs_match_retained_content() {
+    use deepseek_recipe::openai::ResponsesRequest;
+    use deepseek_recipe::request::{ConversionOptions, ProtocolRequest};
+    let request: ResponsesRequest = serde_json::from_value(serde_json::json!({
+        "model":"DeepSeek-V4.1-Flash", "input":"Hello", "reasoning":{"effort":"none"}, "top_logprobs":0
+    })).unwrap();
+    let mut converted = request.convert(ConversionOptions::default()).unwrap();
+    // Responses has no request-level stop field; the caller can configure parsing.
+    converted.parsing_options.stop_sequences = vec!["STOP".into()];
+    let generator = ResponsesRequest::chunk_generator(
+        &converted,
+        "test-id".into(),
+        "DeepSeek-V4.1-Flash".into(),
+    );
+    let chunks = public_chunks(generator, converted.parsing_options).await;
+    let deltas: Vec<_> = chunks
+        .iter()
+        .filter(|chunk| chunk["type"] == "response.output_text.delta")
+        .collect();
+    let text: String = deltas
+        .iter()
+        .map(|chunk| chunk["delta"].as_str().unwrap())
+        .collect();
+    let tokens: String = deltas
+        .iter()
+        .flat_map(|chunk| chunk["logprobs"].as_array().unwrap())
+        .map(|entry| entry["token"].as_str().unwrap())
+        .collect();
+    assert_eq!(text, "AB");
+    assert_eq!(tokens, "AB");
+}

docs/streaming.md

@@ -8,6 +8,13 @@ Completions chunks. The application supplies the backend and HTTP transport.
 In thinking mode, the processor counts `reasoning_tokens` for usage only when a
 tokenizer is attached; see [Use with tokenizer](tokenizer.md#decoding).
 
+Token log probabilities follow the original bytes retained by the parser.
+Only complete tokens within an emitted source-text segment are reported; tokens
+crossing a parser segment boundary are omitted. Removed text and replacement
+spaces around JSON output carry no probabilities. If a custom tokenizer's
+vocabulary bytes do not match its decoded text, probabilities for that decoded
+piece are omitted. Completion usage still counts the processed backend tokens.
+
 ## Rust
 
 Add the dependencies to your application:

docs/streaming.zh.md

@@ -4,6 +4,8 @@
 
 `StreamProcessor` 把后端输出转换为 Messages、Chat Completions 或 Responses 事件。以下示例使用 mock 推理并打印 Chat Completions chunk。后端与 HTTP 传输由应用提供。思考模式下，只有为 processor 附加 tokenizer 后，它才会统计 usage 中的 `reasoning_tokens`，见[使用 tokenizer](tokenizer.zh.md#解码)。
 
+Token 的 log probability 按解析器保留的原始字节归属。只报告完整落在输出原文片段内的 token；跨解析片段边界的 token 不报告概率。被删除的文本和 JSON 输出周围生成的替代空格不带概率。若自定义 tokenizer 的词表字节与解码文本不一致，该解码片段不报告概率。completion usage 仍按已处理的后端 token 计数。
+
 ## Rust
 
 为你的应用添加依赖：