dynamo_parsers/reasoning/
granite_parser.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::ParserResult;
5use crate::ReasoningParser;
6
7#[derive(Debug, Clone, PartialEq, Eq, Hash)]
8pub struct GraniteReasoningParser {
9    think_start_tokens: Vec<String>,
10    think_end_tokens: Vec<String>,
11    buffer: String,
12    stripped_think_start: bool,
13    in_reasoning: bool,
14}
15
16impl GraniteReasoningParser {
17    pub fn new() -> Self {
18        Self {
19            think_start_tokens: ["Here's my thought process:", "Here is my thought process:"]
20                .iter()
21                .map(|s| s.to_string())
22                .collect(),
23            think_end_tokens: ["Here's my response:", "Here is my response:"]
24                .iter()
25                .map(|s| s.to_string())
26                .collect(),
27            buffer: String::new(),
28            stripped_think_start: false,
29            in_reasoning: false,
30        }
31    }
32}
33
34impl Default for GraniteReasoningParser {
35    fn default() -> Self {
36        Self::new()
37    }
38}
39
40impl ReasoningParser for GraniteReasoningParser {
41    fn detect_and_parse_reasoning(&mut self, text: &str, _: &[u32]) -> ParserResult {
42        let think_start_token = self
43            .think_start_tokens
44            .iter()
45            .find(|&token| text.contains(token))
46            .unwrap_or_else(|| self.think_start_tokens.first().unwrap());
47
48        let think_end_token = self
49            .think_end_tokens
50            .iter()
51            .find(|&token| text.contains(token))
52            .unwrap_or_else(|| self.think_end_tokens.first().unwrap());
53        // Implement parsing logic specific to Granite format
54        let in_reasoning = self.in_reasoning
55            || self
56                .think_start_tokens
57                .iter()
58                .any(|token| text.contains(token));
59        if !in_reasoning {
60            return ParserResult {
61                normal_text: text.to_string(),
62                reasoning_text: String::new(),
63            };
64        }
65
66        // The text is considered to be in a reasoning block.
67        let processed_text = text.replacen(think_start_token, "", 1).trim().to_string();
68
69        if !processed_text.contains(think_end_token) {
70            // Assume reasoning was truncated before `think_end_token`
71            return ParserResult {
72                normal_text: String::new(),
73                reasoning_text: processed_text,
74            };
75        }
76
77        // Extract reasoning content
78        let splits: Vec<&str> = processed_text.splitn(2, think_end_token).collect();
79        let reasoning_text = splits.first().unwrap_or(&"").to_string();
80        let normal_text = splits
81            .get(1)
82            .map(|s| s.trim().to_string())
83            .unwrap_or_default();
84
85        ParserResult {
86            normal_text,
87            reasoning_text,
88        }
89    }
90
91    fn parse_reasoning_streaming_incremental(&mut self, text: &str, _: &[u32]) -> ParserResult {
92        // Implement streaming parsing logic specific to Granite format
93
94        // Incrementally parse the streaming text
95        self.buffer.push_str(text);
96        let mut current_text = self.buffer.to_string();
97        // If the current text is a prefix of the think token, keep buffering
98
99        for think_start_token in &self.think_start_tokens {
100            if think_start_token.starts_with(&current_text)
101                && think_start_token.as_str() != current_text.as_str()
102            {
103                return ParserResult {
104                    normal_text: String::new(),
105                    reasoning_text: String::new(),
106                };
107            }
108        }
109        for think_end_token in &self.think_end_tokens {
110            if think_end_token.starts_with(&current_text)
111                && think_end_token.as_str() != current_text.as_str()
112            {
113                return ParserResult {
114                    normal_text: String::new(),
115                    reasoning_text: String::new(),
116                };
117            }
118        }
119
120        let think_start_token = self
121            .think_start_tokens
122            .iter()
123            .find(|&token| current_text.contains(token))
124            .unwrap_or_else(|| self.think_start_tokens.first().unwrap());
125
126        let think_end_token = self
127            .think_end_tokens
128            .iter()
129            .find(|&token| current_text.contains(token))
130            .unwrap_or_else(|| self.think_end_tokens.first().unwrap());
131
132        if !self.stripped_think_start && current_text.contains(think_start_token) {
133            current_text = current_text.replacen(think_start_token, "", 1);
134            self.buffer = current_text.to_string();
135            self.stripped_think_start = true;
136            self.in_reasoning = true;
137        }
138        // Handle end of reasoning block
139        let mut think_end_idx = current_text.len();
140        if self.in_reasoning {
141            think_end_idx = current_text
142                .find(think_end_token)
143                .unwrap_or(current_text.len());
144        }
145        if self.in_reasoning && think_end_idx < current_text.len() {
146            let reasoning_text = &current_text[..think_end_idx];
147            self.buffer.clear();
148            self.in_reasoning = false;
149            let start_idx = think_end_idx + think_end_token.len();
150            let normal_text = if start_idx < current_text.len() {
151                &current_text[start_idx..]
152            } else {
153                ""
154            };
155            return ParserResult {
156                normal_text: normal_text.to_string(),
157                reasoning_text: reasoning_text.to_string(),
158            };
159        }
160        // Continue with reasoning content
161        if self.in_reasoning {
162            // Stream the content immediately
163            let reasoning_text = current_text;
164            self.buffer.clear();
165            ParserResult {
166                normal_text: String::new(),
167                reasoning_text,
168            }
169        } else {
170            // If we're not in a reasoning block return as normal text
171            let normal_text = current_text;
172            self.buffer.clear();
173            ParserResult {
174                normal_text,
175                reasoning_text: String::new(),
176            }
177        }
178    }
179}
180
181#[cfg(test)]
182mod tests {
183    use super::*;
184
185    #[test]
186    fn test_basic_reasoning_detection() {
187        let mut parser = GraniteReasoningParser::new();
188        let text = "Here's my thought process: I need to think about this. Here's my response: The answer is 42.";
189        let result = parser.parse_reasoning_streaming_incremental(text, &[]);
190
191        assert_eq!(result.reasoning_text, " I need to think about this. ");
192        assert_eq!(result.normal_text, " The answer is 42.");
193    }
194
195    #[test]
196    fn test_alternative_start_token() {
197        let mut parser = GraniteReasoningParser::new();
198        let text = "Here is my thought process: Different thinking here. Here is my response: Final answer.";
199        let result = parser.parse_reasoning_streaming_incremental(text, &[]);
200
201        assert_eq!(result.reasoning_text, " Different thinking here. ");
202        assert_eq!(result.normal_text, " Final answer.");
203    }
204
205    #[test]
206    fn test_streaming_partial_tokens() {
207        let mut parser = GraniteReasoningParser::new();
208
209        // Test partial start token
210        let result1 = parser.parse_reasoning_streaming_incremental("Here's", &[]);
211        assert_eq!(result1.normal_text, "");
212        assert_eq!(result1.reasoning_text, "");
213
214        // Complete the start token and add reasoning
215        let result2 = parser
216            .parse_reasoning_streaming_incremental(" my thought process: This is reasoning", &[]);
217        assert_eq!(result2.reasoning_text, " This is reasoning");
218        assert_eq!(result2.normal_text, "");
219    }
220
221    #[test]
222    fn test_streaming_partial_end_tokens() {
223        let mut parser = GraniteReasoningParser::new();
224
225        // Start reasoning
226        parser
227            .parse_reasoning_streaming_incremental("Here's my thought process: Thinking... ", &[]);
228
229        parser.parse_reasoning_streaming_incremental("Here", &[]);
230
231        // Partial end token should buffer
232        let result = parser.parse_reasoning_streaming_incremental("'s my", &[]);
233        assert_eq!(result.normal_text, "");
234        assert_eq!(result.reasoning_text, "");
235
236        // Complete end token
237        let result2 = parser.parse_reasoning_streaming_incremental(" response: Done!", &[]);
238        assert_eq!(result2.reasoning_text, "");
239        assert_eq!(result2.normal_text, " Done!");
240    }
241
242    #[test]
243    fn test_no_reasoning_tokens() {
244        let mut parser = GraniteReasoningParser::new();
245        let text = "This is just normal text without any special tokens.";
246        let result = parser.parse_reasoning_streaming_incremental(text, &[]);
247
248        assert_eq!(result.normal_text, text);
249        assert_eq!(result.reasoning_text, "");
250    }
251
252    #[test]
253    fn test_only_start_token_no_end() {
254        let mut parser = GraniteReasoningParser::new();
255
256        let result1 = parser.parse_reasoning_streaming_incremental(
257            "Here's my thought process: This is reasoning content",
258            &[],
259        );
260        assert_eq!(result1.reasoning_text, " This is reasoning content");
261        assert_eq!(result1.normal_text, "");
262
263        // More reasoning content without end token
264        let result2 = parser.parse_reasoning_streaming_incremental(" and more thinking", &[]);
265        assert_eq!(result2.reasoning_text, " and more thinking");
266        assert_eq!(result2.normal_text, "");
267    }
268
269    #[test]
270    fn test_empty_reasoning_block() {
271        let mut parser = GraniteReasoningParser::new();
272        let text = "Here's my thought process:Here's my response: Direct answer.";
273        let result = parser.parse_reasoning_streaming_incremental(text, &[]);
274
275        assert_eq!(result.reasoning_text, "");
276        assert_eq!(result.normal_text, " Direct answer.");
277    }
278
279    #[test]
280    fn test_reasoning_with_whitespace() {
281        let mut parser = GraniteReasoningParser::new();
282        let text = "Here's my thought process:   \n  Indented reasoning  \n  Here's my response:   Final result  ";
283        let result = parser.parse_reasoning_streaming_incremental(text, &[]);
284
285        assert_eq!(result.reasoning_text, "   \n  Indented reasoning  \n  ");
286        assert_eq!(result.normal_text, "   Final result  ");
287    }
288
289    #[test]
290    fn test_case_sensitive_tokens() {
291        let mut parser = GraniteReasoningParser::new();
292        let text = "here's my thought process: lowercase. here's my response: answer.";
293        let result = parser.parse_reasoning_streaming_incremental(text, &[]);
294
295        // Should not detect lowercase tokens
296        assert_eq!(result.normal_text, text);
297        assert_eq!(result.reasoning_text, "");
298    }
299
300    #[test]
301    fn test_nested_or_repeated_tokens() {
302        let mut parser = GraniteReasoningParser::new();
303        let text = "Here's my thought process: I think Here's my thought process: is confusing. Here's my response: Done.";
304        let result = parser.parse_reasoning_streaming_incremental(text, &[]);
305
306        assert_eq!(
307            result.reasoning_text,
308            " I think Here's my thought process: is confusing. "
309        );
310        assert_eq!(result.normal_text, " Done.");
311    }
312
313    #[test]
314    fn test_detect_and_parse_reasoning_basic() {
315        let mut parser = GraniteReasoningParser::new();
316        let text = "Here's my thought process: I need to analyze this problem. Here's my response: The solution is clear.";
317        let result = parser.detect_and_parse_reasoning(text, &[]);
318
319        assert_eq!(result.reasoning_text, "I need to analyze this problem. ");
320        assert_eq!(result.normal_text, "The solution is clear.");
321    }
322
323    #[test]
324    fn test_detect_and_parse_reasoning_alternative_tokens() {
325        let mut parser = GraniteReasoningParser::new();
326        let text = "Here is my thought process: Different reasoning approach. Here is my response: Final conclusion.";
327        let result = parser.detect_and_parse_reasoning(text, &[]);
328
329        assert_eq!(result.reasoning_text, "Different reasoning approach. ");
330        assert_eq!(result.normal_text, "Final conclusion.");
331    }
332
333    #[test]
334    fn test_detect_and_parse_reasoning_no_tokens() {
335        let mut parser = GraniteReasoningParser::new();
336        let text = "This is just normal text without special markers.";
337        let result = parser.detect_and_parse_reasoning(text, &[]);
338
339        assert_eq!(result.normal_text, text);
340        assert_eq!(result.reasoning_text, "");
341    }
342
343    #[test]
344    fn test_detect_and_parse_reasoning_only_start_token() {
345        let mut parser = GraniteReasoningParser::new();
346        let text = "Here's my thought process: This reasoning has no end marker.";
347        let result = parser.detect_and_parse_reasoning(text, &[]);
348
349        assert_eq!(result.reasoning_text, "This reasoning has no end marker.");
350        assert_eq!(result.normal_text, "");
351    }
352
353    #[test]
354    fn test_detect_and_parse_reasoning_empty_sections() {
355        let mut parser = GraniteReasoningParser::new();
356        let text = "Here's my thought process:Here's my response:";
357        let result = parser.detect_and_parse_reasoning(text, &[]);
358
359        assert_eq!(result.reasoning_text, "");
360        assert_eq!(result.normal_text, "");
361    }
362
363    #[test]
364    fn test_detect_and_parse_reasoning_whitespace_handling() {
365        let mut parser = GraniteReasoningParser::new();
366        let text = "Here's my thought process:   \n\tSpaced reasoning\n   Here's my response:  \n  Spaced response\n";
367        let result = parser.detect_and_parse_reasoning(text, &[]);
368
369        assert_eq!(result.reasoning_text, "Spaced reasoning\n   ");
370        assert_eq!(result.normal_text, "Spaced response");
371    }
372
373    #[test]
374    fn test_detect_and_parse_reasoning_multiple_end_tokens() {
375        let mut parser = GraniteReasoningParser::new();
376        let text = "Here's my thought process: Thinking about Here's my response: in the middle. Here's my response: Real end.";
377        let result = parser.detect_and_parse_reasoning(text, &[]);
378
379        assert_eq!(result.reasoning_text, "Thinking about ");
380        assert_eq!(
381            result.normal_text,
382            "in the middle. Here's my response: Real end."
383        );
384    }
385
386    #[test]
387    fn test_detect_and_parse_reasoning_case_sensitivity() {
388        let mut parser = GraniteReasoningParser::new();
389        let text =
390            "here's my thought process: lowercase tokens. here's my response: should not work.";
391        let result = parser.detect_and_parse_reasoning(text, &[]);
392
393        assert_eq!(result.normal_text, text);
394        assert_eq!(result.reasoning_text, "");
395    }
396
397    #[test]
398    fn test_detect_and_parse_reasoning_mixed_tokens() {
399        let mut parser = GraniteReasoningParser::new();
400        let text = "Here's my thought process: First reasoning. Here is my response: Mixed token response.";
401        let result = parser.detect_and_parse_reasoning(text, &[]);
402
403        assert_eq!(result.reasoning_text, "First reasoning. ");
404        assert_eq!(result.normal_text, "Mixed token response.");
405    }
406
407    #[test]
408    fn test_detect_and_parse_reasoning_long_content() {
409        let mut parser = GraniteReasoningParser::new();
410        let text = "Here's my thought process: This is a very long reasoning section that spans multiple sentences. I need to consider various factors. The analysis requires careful thought. Here's my response: After all that thinking, here is the comprehensive answer with multiple parts and detailed explanation.";
411        let result = parser.detect_and_parse_reasoning(text, &[]);
412
413        assert_eq!(
414            result.reasoning_text,
415            "This is a very long reasoning section that spans multiple sentences. I need to consider various factors. The analysis requires careful thought. "
416        );
417        assert_eq!(
418            result.normal_text,
419            "After all that thinking, here is the comprehensive answer with multiple parts and detailed explanation."
420        );
421    }
422}
dynamo_parsers/reasoning/granite_parser.rs

dynamo_parsers/reasoning/
granite_parser.rs
