chore: use token ids for reasoning parsing instead of strings directly

lib/llm/src/engines.rs

@@ -183,7 +183,7 @@ impl
         incoming_request: SingleIn<NvCreateChatCompletionRequest>,
     ) -> Result<ManyOut<Annotated<NvCreateChatCompletionStreamResponse>>, Error> {
         let (request, context) = incoming_request.transfer(());
-        let mut deltas = request.response_generator();
+        let mut deltas = request.response_generator(None, None);
         let ctx = context.context();
         let req = request.inner.messages.into_iter().next_back().unwrap();
 
@@ -204,12 +204,12 @@ 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);
+                let response = deltas.create_choice(0, Some(c.to_string()), None, None, None);
                 yield Annotated{ id: Some(id.to_string()), data: Some(response), event: None, comment: None };
                 id += 1;
             }
 
-            let response = deltas.create_choice(0, None, Some(dynamo_async_openai::types::FinishReason::Stop), None);
+            let response = deltas.create_choice(0, None, None, Some(dynamo_async_openai::types::FinishReason::Stop), None);
             yield Annotated { id: Some(id.to_string()), data: Some(response), event: None, comment: None };
         };
 

lib/llm/src/preprocessor.rs

@@ -94,6 +94,7 @@ pub struct OpenAIPreprocessor {
     formatter: Arc<dyn OAIPromptFormatter>,
     tokenizer: Arc<dyn Tokenizer>,
     model_info: Arc<dyn ModelInfo>,
+    vocab: HashMap<String, u32>,
 }
 
 impl OpenAIPreprocessor {
@@ -113,6 +114,7 @@ impl OpenAIPreprocessor {
                 );
             }
         };
+        let vocab = tokenizer.get_vocab(true);
         let tokenizer = Arc::new(tokenizer);
 
         let Some(model_info) = mdc.model_info else {
@@ -127,6 +129,7 @@ impl OpenAIPreprocessor {
             tokenizer,
             model_info,
             mdcsum,
+            vocab,
         }))
     }
 
@@ -499,7 +502,8 @@ impl
         let (request, context) = request.into_parts();
 
         // create a response generator
-        let response_generator = request.response_generator();
+        let response_generator =
+            request.response_generator(Some(&self.tokenizer), Some(&self.vocab));
         let mut response_generator = Box::new(response_generator);
 
         // convert the chat completion request to a common completion request

lib/llm/src/protocols/openai/chat_completions.rs

@@ -28,7 +28,7 @@ use super::{
 };
 
 pub mod aggregator;
-mod delta;
+pub mod delta;
 
 pub use aggregator::DeltaAggregator;
 pub use delta::DeltaGenerator;

lib/llm/src/protocols/openai/chat_completions/delta.rs

@@ -1,6 +1,9 @@
 // SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+use crate::tokenizers::traits::Tokenizer;
+use std::{collections::HashMap, sync::Arc};
+
 use dynamo_parsers::{ParserResult, ReasoningParser, ReasoningParserType, ReasoningParserWrapper};
 
 use super::{NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse};
@@ -15,14 +18,18 @@ impl NvCreateChatCompletionRequest {
     ///
     /// # Returns
     /// * [`DeltaGenerator`] configured with model name and response options.
-    pub fn response_generator(&self) -> DeltaGenerator {
+    pub fn response_generator(
+        &self,
+        tokenizer: Option<&Arc<(dyn Tokenizer + 'static)>>,
+        vocab: Option<&HashMap<String, u32>>,
+    ) -> DeltaGenerator {
         let options = DeltaGeneratorOptions {
             enable_usage: true,
             enable_logprobs: self.inner.logprobs.unwrap_or(false)
                 || self.inner.top_logprobs.unwrap_or(0) > 0,
         };
 
-        DeltaGenerator::new(self.inner.model.clone(), options)
+        DeltaGenerator::new(self.inner.model.clone(), options, tokenizer, vocab)
     }
 }
 
@@ -36,7 +43,6 @@ pub struct DeltaGeneratorOptions {
 }
 
 /// Generates incremental chat completion responses in a streaming fashion.
-#[derive(Debug)]
 pub struct DeltaGenerator {
     /// Unique identifier for the chat completion session.
     id: String,
@@ -58,7 +64,9 @@ pub struct DeltaGenerator {
 
     /// Reasoning Parser object
     /// This is used to parse reasoning content in the response.
-    reasoning_parser: ReasoningParserWrapper,
+    reasoning_parser: Option<ReasoningParserWrapper>,
+
+    tokenizer: Option<Arc<(dyn Tokenizer + 'static)>>,
 }
 
 impl DeltaGenerator {
@@ -70,7 +78,12 @@ impl DeltaGenerator {
     ///
     /// # Returns
     /// * A new instance of [`DeltaGenerator`].
-    pub fn new(model: String, options: DeltaGeneratorOptions) -> Self {
+    pub fn new(
+        model: String,
+        options: DeltaGeneratorOptions,
+        tokenizer: Option<&Arc<(dyn Tokenizer + 'static)>>,
+        vocab: Option<&HashMap<String, u32>>,
+    ) -> Self {
         let now = std::time::SystemTime::now()
             .duration_since(std::time::UNIX_EPOCH)
             .unwrap()
@@ -94,7 +107,7 @@ impl DeltaGenerator {
         let reasoning_parser_type = ReasoningParserType::Basic;
 
         // Reasoning parser wrapper
-        let reasoning_parser = reasoning_parser_type.get_reasoning_parser();
+        let reasoning_parser = vocab.map(|v| reasoning_parser_type.get_reasoning_parser(v));
 
         Self {
             id: format!("chatcmpl-{}", uuid::Uuid::new_v4()),
@@ -107,6 +120,7 @@ impl DeltaGenerator {
             msg_counter: 0,
             options,
             reasoning_parser,
+            tokenizer: tokenizer.cloned(),
         }
     }
 
@@ -183,13 +197,16 @@ impl DeltaGenerator {
         })
     }
 
-    fn create_reasoning_content(&mut self, text: Option<String>) -> Option<ParserResult> {
-        let text = text?;
-        let parser_result = self
-            .reasoning_parser
-            .parse_reasoning_streaming_incremental(&text);
-
-        Some(parser_result)
+    fn create_reasoning_content(&mut self, token_ids: Vec<u32>) -> Option<ParserResult> {
+        if self.tokenizer.is_none() || self.reasoning_parser.is_none() {
+            return None;
+        }
+        Some(
+            self.reasoning_parser
+                .as_mut()
+                .unwrap()
+                .parse_reasoning_streaming_incremental(&token_ids),
+        )
     }
 
     /// Creates a choice within a chat completion response.
@@ -207,17 +224,12 @@ impl DeltaGenerator {
         &mut self,
         index: u32,
         text: Option<String>,
+        reasoning_content: Option<String>,
         finish_reason: Option<dynamo_async_openai::types::FinishReason>,
         logprobs: Option<dynamo_async_openai::types::ChatChoiceLogprobs>,
     ) -> NvCreateChatCompletionStreamResponse {
-        let reasoning_parser_result = self.create_reasoning_content(text).unwrap_or_default();
-
-        let (normal_text, reasoning_content) = (
-            reasoning_parser_result.get_some_normal_text(),
-            reasoning_parser_result.get_some_reasoning(),
-        );
         let delta = dynamo_async_openai::types::ChatCompletionStreamResponseDelta {
-            content: normal_text,
+            content: text,
             function_call: None,
             tool_calls: None,
             role: if self.msg_counter == 0 {
@@ -292,7 +304,7 @@ impl crate::protocols::openai::DeltaGeneratorExt<NvCreateChatCompletionStreamRes
 
         let logprobs = self.create_logprobs(
             delta.tokens,
-            delta.token_ids,
+            delta.token_ids.clone(),
             delta.log_probs,
             delta.top_logprobs,
         );
@@ -319,8 +331,44 @@ impl crate::protocols::openai::DeltaGeneratorExt<NvCreateChatCompletionStreamRes
         };
 
         // Create the streaming response.
+        let reasoning_parser_result = self.create_reasoning_content(delta.token_ids);
+
+        let (normal_text, reasoning_content) = if let Some(parser_result) = reasoning_parser_result
+        {
+            let none_if_empty = |vec: String| {
+                if vec.is_empty() {
+                    None
+                } else {
+                    Some(vec)
+                }
+            };
+            (
+                none_if_empty(
+                    self.tokenizer
+                        .as_ref()
+                        .unwrap()
+                        .decode(&parser_result.normal_token_ids, false)
+                        .unwrap(),
+                ),
+                none_if_empty(
+                    self.tokenizer
+                        .as_ref()
+                        .unwrap()
+                        .decode(&parser_result.reasoning_token_ids, false)
+                        .unwrap(),
+                ),
+            )
+        } else {
+            (delta.text, None)
+        };
         let index = 0;
-        let stream_response = self.create_choice(index, delta.text, finish_reason, logprobs);
+        let stream_response = self.create_choice(
+            index,
+            normal_text,
+            reasoning_content,
+            finish_reason,
+            logprobs,
+        );
 
         Ok(stream_response)
     }

lib/llm/src/tokenizers/hf.rs

@@ -13,6 +13,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashMap;
+
 use tokenizers::tokenizer::Tokenizer as HfTokenizer;
 
 use super::{
@@ -35,6 +37,10 @@ impl HuggingFaceTokenizer {
     pub fn from_tokenizer(tokenizer: HfTokenizer) -> Self {
         HuggingFaceTokenizer { tokenizer }
     }
+
+    pub fn get_vocab(&self, with_added_tokens: bool) -> HashMap<String, u32> {
+        self.tokenizer.get_vocab(with_added_tokens)
+    }
 }
 
 impl Encoder for HuggingFaceTokenizer {

lib/llm/tests/http-service.rs

@@ -95,12 +95,12 @@ impl
         let max_tokens = request.inner.max_tokens.unwrap_or(0) as u64;
 
         // let generator = NvCreateChatCompletionStreamResponse::generator(request.model.clone());
-        let mut generator = request.response_generator();
+        let mut generator = request.response_generator(None, None);
 
         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);
+                let output = generator.create_choice(i,Some(format!("choice {i}")), None, None, None);
 
                 yield Annotated::from_data(output);
             }

lib/parsers/Cargo.toml

@@ -32,4 +32,5 @@ serde_json = { workspace = true }
 tracing = { workspace = true }
 uuid = { workspace = true }
 
-regex = "1"
+regex = "1"
+openai-harmony = { git = "https://github.com/openai/harmony", version = "0.0.4" }