Skip to main content

reasoning_parser/parsers/
cohere_cmd.rs

1//! Cohere Command model reasoning parser.
2//!
3//! Parses thinking blocks from `<|START_THINKING|>...<|END_THINKING|>`.
4//! Supports CMD3 and CMD4 format.
5
6use crate::{
7    parsers::BaseReasoningParser,
8    traits::{ParseError, ParserConfig, ParserResult, ReasoningParser, DEFAULT_MAX_BUFFER_SIZE},
9};
10
11/// Cohere Command model reasoning parser.
12///
13/// Handles `<|START_THINKING|>` and `<|END_THINKING|>` tokens.
14/// Unlike DeepSeek-R1, Cohere requires explicit start token (always_in_reasoning=false).
15pub struct CohereCmdParser {
16    base: BaseReasoningParser,
17}
18
19impl CohereCmdParser {
20    /// Create a new Cohere Command parser.
21    pub fn new() -> Self {
22        let config = ParserConfig {
23            think_start_token: "<|START_THINKING|>".to_string(),
24            think_end_token: "<|END_THINKING|>".to_string(),
25            stream_reasoning: true,
26            max_buffer_size: DEFAULT_MAX_BUFFER_SIZE,
27            always_in_reasoning: false,
28        };
29
30        Self {
31            base: BaseReasoningParser::new(config).with_model_type("cohere_cmd".to_string()),
32        }
33    }
34}
35
36impl Default for CohereCmdParser {
37    fn default() -> Self {
38        Self::new()
39    }
40}
41
42impl ReasoningParser for CohereCmdParser {
43    fn detect_and_parse_reasoning(&mut self, text: &str) -> Result<ParserResult, ParseError> {
44        self.base.detect_and_parse_reasoning(text)
45    }
46
47    fn parse_reasoning_streaming_incremental(
48        &mut self,
49        text: &str,
50    ) -> Result<ParserResult, ParseError> {
51        self.base.parse_reasoning_streaming_incremental(text)
52    }
53
54    fn reset(&mut self) {
55        self.base.reset();
56    }
57
58    fn model_type(&self) -> &str {
59        self.base.model_type()
60    }
61
62    fn is_in_reasoning(&self) -> bool {
63        self.base.is_in_reasoning()
64    }
65
66    fn mark_reasoning_started(&mut self) {
67        self.base.mark_reasoning_started();
68    }
69
70    fn mark_think_start_stripped(&mut self) {
71        self.base.mark_think_start_stripped();
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78
79    #[test]
80    fn test_cohere_cmd_no_reasoning() {
81        let mut parser = CohereCmdParser::new();
82
83        // Without thinking tags, text should be returned as normal
84        let result = parser
85            .detect_and_parse_reasoning("This is a normal response.")
86            .unwrap();
87        assert_eq!(result.normal_text, "This is a normal response.");
88        assert_eq!(result.reasoning_text, "");
89    }
90
91    #[test]
92    fn test_cohere_cmd_with_thinking() {
93        let mut parser = CohereCmdParser::new();
94
95        let result = parser
96            .detect_and_parse_reasoning(
97                "<|START_THINKING|>Let me analyze this step by step.<|END_THINKING|>The answer is 42.",
98            )
99            .unwrap();
100        assert_eq!(result.normal_text, "The answer is 42.");
101        assert_eq!(result.reasoning_text, "Let me analyze this step by step.");
102    }
103
104    #[test]
105    fn test_cohere_cmd_truncated_thinking() {
106        let mut parser = CohereCmdParser::new();
107
108        // Thinking block without end token (truncated)
109        let result = parser
110            .detect_and_parse_reasoning("<|START_THINKING|>Analyzing the problem...")
111            .unwrap();
112        assert_eq!(result.normal_text, "");
113        assert_eq!(result.reasoning_text, "Analyzing the problem...");
114    }
115
116    #[test]
117    fn test_cohere_cmd_streaming() {
118        let mut parser = CohereCmdParser::new();
119
120        // First chunk - start of thinking
121        let result1 = parser
122            .parse_reasoning_streaming_incremental("<|START_THINKING|>Step 1: ")
123            .unwrap();
124        assert_eq!(result1.reasoning_text, "Step 1: ");
125        assert_eq!(result1.normal_text, "");
126
127        // Second chunk - continue thinking
128        let result2 = parser
129            .parse_reasoning_streaming_incremental("Analyze inputs. ")
130            .unwrap();
131        assert_eq!(result2.reasoning_text, "Analyze inputs. ");
132        assert_eq!(result2.normal_text, "");
133
134        // Third chunk - end thinking and normal text
135        let result3 = parser
136            .parse_reasoning_streaming_incremental(
137                "Step 2: Check.<|END_THINKING|>Here's the answer.",
138            )
139            .unwrap();
140        assert_eq!(result3.reasoning_text, "Step 2: Check.");
141        assert_eq!(result3.normal_text, "Here's the answer.");
142    }
143
144    #[test]
145    fn test_cohere_cmd_streaming_partial_token() {
146        let mut parser = CohereCmdParser::new();
147
148        // Partial start token
149        let result1 = parser
150            .parse_reasoning_streaming_incremental("<|START_")
151            .unwrap();
152        assert_eq!(result1.reasoning_text, "");
153        assert_eq!(result1.normal_text, "");
154
155        // Complete the start token
156        let result2 = parser
157            .parse_reasoning_streaming_incremental("THINKING|>reasoning")
158            .unwrap();
159        assert_eq!(result2.reasoning_text, "reasoning");
160        assert_eq!(result2.normal_text, "");
161    }
162
163    #[test]
164    fn test_cohere_cmd_reset() {
165        let mut parser = CohereCmdParser::new();
166
167        // Process some text
168        parser
169            .parse_reasoning_streaming_incremental("<|START_THINKING|>thinking<|END_THINKING|>done")
170            .unwrap();
171
172        // Reset
173        parser.reset();
174        assert!(!parser.is_in_reasoning());
175
176        // Should work fresh again
177        let result = parser
178            .detect_and_parse_reasoning("<|START_THINKING|>new<|END_THINKING|>text")
179            .unwrap();
180        assert_eq!(result.reasoning_text, "new");
181        assert_eq!(result.normal_text, "text");
182    }
183
184    #[test]
185    fn test_model_type() {
186        let parser = CohereCmdParser::new();
187        assert_eq!(parser.model_type(), "cohere_cmd");
188    }
189
190    #[test]
191    fn test_cohere_full_response_format() {
192        let mut parser = CohereCmdParser::new();
193
194        // Simulate full Cohere response with thinking and response markers
195        let input = r"<|START_THINKING|>
196Let me analyze this step by step.
1971. First, I'll consider the question.
1982. Then, I'll formulate a response.
199<|END_THINKING|>
200<|START_RESPONSE|>The answer is 42.<|END_RESPONSE|>";
201
202        let result = parser.detect_and_parse_reasoning(input).unwrap();
203        assert!(result.reasoning_text.contains("step by step"));
204        // Note: Response markers are passed through - cleaned by tool_parser or response processor
205        assert!(result.normal_text.contains("START_RESPONSE"));
206    }
207
208    #[test]
209    fn test_cohere_cmd_empty_thinking() {
210        let mut parser = CohereCmdParser::new();
211
212        let result = parser
213            .detect_and_parse_reasoning("<|START_THINKING|><|END_THINKING|>The answer.")
214            .unwrap();
215        assert_eq!(result.normal_text, "The answer.");
216        assert_eq!(result.reasoning_text, "");
217    }
218
219    #[test]
220    fn test_cohere_cmd_unicode_in_thinking() {
221        let mut parser = CohereCmdParser::new();
222
223        let result = parser
224            .detect_and_parse_reasoning(
225                "<|START_THINKING|>分析这个问题 🤔 emoji test<|END_THINKING|>答案是42。",
226            )
227            .unwrap();
228        assert_eq!(result.reasoning_text, "分析这个问题 🤔 emoji test");
229        assert_eq!(result.normal_text, "答案是42。");
230    }
231
232    #[test]
233    fn test_cohere_cmd_angle_brackets_in_thinking() {
234        let mut parser = CohereCmdParser::new();
235
236        // Angle brackets that don't match tokens should pass through
237        let result = parser
238            .detect_and_parse_reasoning(
239                "<|START_THINKING|>check if x < 10 and y > 5<|END_THINKING|>result",
240            )
241            .unwrap();
242        assert_eq!(result.reasoning_text, "check if x < 10 and y > 5");
243        assert_eq!(result.normal_text, "result");
244    }
245
246    #[test]
247    fn test_cohere_cmd_whitespace_only_thinking() {
248        let mut parser = CohereCmdParser::new();
249
250        let result = parser
251            .detect_and_parse_reasoning("<|START_THINKING|>   \n\t  <|END_THINKING|>answer")
252            .unwrap();
253        // BaseReasoningParser trims reasoning text
254        assert_eq!(result.reasoning_text, "");
255        assert_eq!(result.normal_text, "answer");
256    }
257}