babel/
utils.rs

1use anyhow::Result;
2use crossterm::{
3    execute,
4    style::{Color, Print, ResetColor, SetForegroundColor},
5};
6use serde_json::Value;
7use std::io::stdout;
8
9// Add this helper function before the StreamResponse implementation
10pub fn strip_markdown_code_blocks(s: &str) -> String {
11    let s = s.trim();
12
13    // Check if string starts with ```json and ends with ```
14    if s.starts_with("```json") && s.ends_with("```") {
15        // Extract the content between the markers
16        let without_start = s.strip_prefix("```json").unwrap_or(s);
17        let content = without_start.strip_suffix("```").unwrap_or(without_start);
18        return content.trim().to_string();
19    }
20    // s.replace("```json", "").replace("```", "")
21    // If not a full code block, just return the original string
22    s.to_string()
23}
24
25// Define our rendering states with associated colors
26#[allow(dead_code)]
27#[derive(Debug, Clone, PartialEq)]
28enum MarkdownState {
29    Normal,
30    Heading(usize), // Level 1-6
31    Bold,
32    Italic,
33    BoldItalic,
34    CodeBlock(String), // Language
35    InlineCode,
36    UnorderedList(usize),      // Nesting level
37    OrderedList(usize, usize), // Nesting level, current number
38    Link,
39    LinkUrl,
40    Blockquote(usize), // Nesting level
41}
42
43pub struct MarkdownStreamRenderer {
44    buffer: String,
45    in_response: bool,
46    depth: i32,
47    extracted: String,
48    state_stack: Vec<MarkdownState>,
49    current_line: String,
50}
51
52impl MarkdownStreamRenderer {
53    pub fn new() -> Self {
54        Self {
55            buffer: String::new(),
56            in_response: false,
57            depth: 0,
58            extracted: String::new(),
59            state_stack: vec![MarkdownState::Normal],
60            current_line: String::new(),
61        }
62    }
63
64    pub fn process_chunk(&mut self, chunk: &str) -> String {
65        self.buffer.push_str(chunk);
66        let mut output = String::new();
67
68        // Try to find response field
69        if !self.in_response && self.buffer.contains(r#""response":"#) {
70            if let Some(pos) = self.buffer.find(r#""response":"#) {
71                // Safely get the part after the response field
72                let safe_start = pos + r#""response":"#.len();
73                let chars: Vec<char> = self.buffer.chars().collect();
74                if safe_start < self.buffer.len() {
75                    // Use character iterator to ensure we slice at character boundaries
76                    let new_buffer: String = chars.into_iter().skip(safe_start).collect();
77                    self.buffer = new_buffer;
78                } else {
79                    self.buffer.clear();
80                }
81                self.in_response = true;
82                self.depth = 0;
83            }
84        }
85
86        // If we've found the response field, start extracting its content
87        if self.in_response {
88            let chars: Vec<char> = self.buffer.chars().collect();
89            let mut i = 0;
90
91            while i < chars.len() {
92                let c = chars[i];
93
94                // Check for quote characters
95                let is_escape_quote = i > 0 && chars[i - 1] == '\\' && c == '"';
96
97                if c == '"' && !is_escape_quote {
98                    if self.depth == 0 {
99                        // Found start quote - start capturing but don't include the quote itself
100                        self.depth = 1;
101                        i += 1; // Skip the opening quote
102                        continue;
103                    } else {
104                        // Found end quote - stop capturing
105                        self.depth = 0;
106                        self.in_response = false;
107                        break;
108                    }
109                }
110
111                if self.depth == 1 {
112                    if c == '\\' && i + 1 < chars.len() {
113                        // Normal case: both backslash and escaped char are in the same chunk
114                        let next_char = chars[i + 1];
115                        if next_char == 'n' {
116                            output.push('\n');
117                        } else if next_char == '"' {
118                            output.push('"');
119                        } else if next_char == '\\' {
120                            output.push('\\');
121                        } else {
122                            output.push(c);
123                            i -= 1; // Unknown escape, don't skip next character
124                        }
125                        i += 1; // Skip the next character in the escape sequence
126                    } else if self.extracted.chars().last() == Some('\\') {
127                        // Split case: backslash was at end of previous chunk
128                        if c == 'n' {
129                            output.push('\n');
130                            // Remove the backslash from extracted to prevent double-handling
131                            self.extracted.pop();
132                        } else if c == '"' {
133                            output.push('"');
134                            self.extracted.pop();
135                        } else if c == '\\' {
136                            output.push('\\');
137                            self.extracted.pop();
138                        } else {
139                            // Not an escape sequence we recognize, just add the original backslash and this char
140                            output.push(c);
141                        }
142                    } else {
143                        output.push(c);
144                    }
145                }
146
147                i += 1;
148            }
149
150            // Update buffer, using character indices instead of byte indices
151            if !self.in_response {
152                self.buffer.clear();
153            } else if i < chars.len() {
154                self.buffer = chars.into_iter().skip(i).collect();
155            } else {
156                self.buffer.clear();
157            }
158        }
159        self.extracted.push_str(&output);
160        let output = if output.chars().last() == Some('\\') {
161            output.pop();
162            output
163        } else {
164            output
165        };
166        // Render markdown if we have content
167        if !output.is_empty() {
168            let _ = self.render_increment(&output);
169            return String::new(); // Return empty since we're handling rendering
170        }
171
172        output
173    }
174
175    fn current_state(&self) -> &MarkdownState {
176        self.state_stack.last().unwrap_or(&MarkdownState::Normal)
177    }
178
179    fn push_state(&mut self, state: MarkdownState) {
180        self.state_stack.push(state);
181    }
182
183    fn pop_state(&mut self) -> Option<MarkdownState> {
184        self.state_stack.pop()
185    }
186
187    fn render_increment(&mut self, text: &str) -> Result<()> {
188        // Don't clear previous render - just render the new text incrementally
189        // Process and render the new text
190        self.process_text(text)?;
191
192        // We don't need to track render height since we're not clearing anything
193
194        Ok(())
195    }
196
197    fn process_text(&mut self, text: &str) -> Result<()> {
198        let mut chars = text.chars().peekable();
199
200        while let Some(c) = chars.next() {
201            // Add character to current line buffer
202            self.current_line.push(c);
203
204            // Apply styling based on current character and state
205            match c {
206                '#' => {
207                    // Check for heading at start of line
208                    if self.current_line.trim() == "#" {
209                        // First # of a potential heading
210                        let mut level = 1;
211
212                        // Look ahead to count consecutive # characters
213                        let mut lookahead = chars.clone();
214                        while lookahead.next_if_eq(&'#').is_some() {
215                            level += 1;
216                        }
217
218                        // Set heading color based on level
219                        let color = match level {
220                            1 => Color::Magenta,
221                            2 => Color::DarkMagenta,
222                            3 => Color::Cyan,
223                            _ => Color::White,
224                        };
225
226                        // Print the # with appropriate color
227                        execute!(stdout(), SetForegroundColor(color), Print("#"))?;
228
229                        // Push heading state
230                        self.push_state(MarkdownState::Heading(level));
231                    } else {
232                        // Just a regular # character, not at start of line
233                        self.print_with_current_style("#")?;
234                    }
235                }
236                '*' => {
237                    // Handle asterisks for bold/italic or list items
238                    if self.current_line.trim() == "*" && chars.peek() == Some(&' ') {
239                        // Unordered list item
240                        execute!(stdout(), SetForegroundColor(Color::Green), Print("*"))?;
241                        self.push_state(MarkdownState::UnorderedList(0));
242                    } else if self.current_line.ends_with("**") {
243                        // Bold marker
244                        execute!(stdout(), SetForegroundColor(Color::Yellow), Print("**"))?;
245                        self.current_line.pop(); // Remove the last * we just printed
246                        self.current_line.pop(); // Remove the second-to-last *
247
248                        // Toggle bold state
249                        match self.current_state() {
250                            MarkdownState::Bold => {
251                                let _ = self.pop_state();
252                                execute!(stdout(), ResetColor)?;
253                            }
254                            _ => self.push_state(MarkdownState::Bold),
255                        }
256                    } else if self.current_line.ends_with("*") && !self.current_line.ends_with("**")
257                    {
258                        // Italic marker
259                        execute!(stdout(), SetForegroundColor(Color::Blue), Print("*"))?;
260                        self.current_line.pop(); // Remove the * we just printed
261
262                        // Toggle italic state
263                        match self.current_state() {
264                            MarkdownState::Italic => {
265                                let _ = self.pop_state();
266                                execute!(stdout(), ResetColor)?;
267                            }
268                            _ => self.push_state(MarkdownState::Italic),
269                        }
270                    } else {
271                        // Just a regular asterisk
272                        self.print_with_current_style("*")?;
273                    }
274                }
275                '`' => {
276                    // Handle backticks for code
277                    if self.current_line.ends_with("```") {
278                        // Code block marker
279                        execute!(stdout(), SetForegroundColor(Color::Yellow), Print("```"))?;
280                        self.current_line.pop(); // Remove the last `
281                        self.current_line.pop(); // Remove the second `
282                        self.current_line.pop(); // Remove the third `
283
284                        // Toggle code block state
285                        match self.current_state() {
286                            MarkdownState::CodeBlock(_) => {
287                                let _ = self.pop_state();
288                                execute!(stdout(), ResetColor)?;
289                            }
290                            _ => self.push_state(MarkdownState::CodeBlock(String::new())),
291                        }
292                    } else if self.current_line.ends_with("`") {
293                        // Inline code marker
294                        execute!(stdout(), SetForegroundColor(Color::Yellow), Print("`"))?;
295                        self.current_line.pop(); // Remove the ` we just printed
296
297                        // Toggle inline code state
298                        match self.current_state() {
299                            MarkdownState::InlineCode => {
300                                let _ = self.pop_state();
301                                execute!(stdout(), ResetColor)?;
302                            }
303                            _ => self.push_state(MarkdownState::InlineCode),
304                        }
305                    } else {
306                        // Just a regular backtick
307                        self.print_with_current_style("`")?;
308                    }
309                }
310                '[' => {
311                    // Link opening bracket
312                    execute!(stdout(), SetForegroundColor(Color::Blue), Print("["))?;
313                    self.push_state(MarkdownState::Link);
314                }
315                ']' => {
316                    // Link closing bracket
317                    execute!(stdout(), SetForegroundColor(Color::Blue), Print("]"))?;
318
319                    // Check if we're in a link state
320                    if matches!(self.current_state(), MarkdownState::Link) {
321                        self.pop_state();
322
323                        // Check for opening parenthesis for URL
324                        if chars.peek() == Some(&'(') {
325                            self.push_state(MarkdownState::LinkUrl);
326                        }
327                    }
328                }
329                '(' => {
330                    if matches!(self.current_state(), MarkdownState::LinkUrl) {
331                        // URL opening parenthesis
332                        execute!(stdout(), SetForegroundColor(Color::DarkBlue), Print("("))?;
333                    } else {
334                        // Regular parenthesis
335                        self.print_with_current_style("(")?;
336                    }
337                }
338                ')' => {
339                    if matches!(self.current_state(), MarkdownState::LinkUrl) {
340                        // URL closing parenthesis
341                        execute!(stdout(), SetForegroundColor(Color::DarkBlue), Print(")"))?;
342                        self.pop_state();
343                    } else {
344                        // Regular parenthesis
345                        self.print_with_current_style(")")?;
346                    }
347                }
348                '1'..='9' => {
349                    // Check for ordered list at start of line
350                    if self.current_line.trim().len() == 1 && chars.peek() == Some(&'.') {
351                        // Potential ordered list item
352                        execute!(stdout(), SetForegroundColor(Color::Green), Print(c))?;
353                    } else {
354                        // Regular digit
355                        self.print_with_current_style(c.to_string().as_str())?;
356                    }
357                }
358                '.' => {
359                    if self.current_line.trim().len() >= 1
360                        && self
361                            .current_line
362                            .trim()
363                            .chars()
364                            .next()
365                            .unwrap()
366                            .is_digit(10)
367                        && self.current_line.trim().ends_with('.')
368                        && chars.peek() == Some(&' ')
369                    {
370                        // Ordered list dot
371                        execute!(stdout(), SetForegroundColor(Color::Green), Print("."))?;
372                        self.push_state(MarkdownState::OrderedList(0, 0));
373                    } else {
374                        // Regular dot
375                        self.print_with_current_style(".")?;
376                    }
377                }
378                '>' => {
379                    // Blockquote
380                    if self.current_line.trim() == ">" {
381                        execute!(stdout(), SetForegroundColor(Color::Cyan), Print(">"))?;
382                        self.push_state(MarkdownState::Blockquote(1));
383                    } else {
384                        // Regular > character
385                        self.print_with_current_style(">")?;
386                    }
387                }
388                '\n' => {
389                    // End of line
390                    execute!(stdout(), Print("\n"))?;
391                    self.current_line.clear();
392
393                    // Reset line-specific states
394                    self.reset_line_states()?;
395                }
396                _ => {
397                    // Regular character
398                    self.print_with_current_style(c.to_string().as_str())?;
399                }
400            }
401        }
402
403        Ok(())
404    }
405
406    // Helper method to print text with current style
407    fn print_with_current_style(&self, text: &str) -> Result<()> {
408        match self.current_state() {
409            MarkdownState::Normal => {
410                execute!(stdout(), ResetColor, Print(text))?;
411            }
412            MarkdownState::Heading(level) => {
413                let color = match level {
414                    1 => Color::Magenta,
415                    2 => Color::DarkMagenta,
416                    3 => Color::Cyan,
417                    _ => Color::White,
418                };
419                execute!(stdout(), SetForegroundColor(color), Print(text))?;
420            }
421            MarkdownState::Bold => {
422                execute!(stdout(), SetForegroundColor(Color::Yellow), Print(text))?;
423            }
424            MarkdownState::Italic => {
425                execute!(stdout(), SetForegroundColor(Color::Blue), Print(text))?;
426            }
427            MarkdownState::BoldItalic => {
428                execute!(stdout(), SetForegroundColor(Color::Magenta), Print(text))?;
429            }
430            MarkdownState::CodeBlock(_) => {
431                execute!(stdout(), SetForegroundColor(Color::Yellow), Print(text))?;
432            }
433            MarkdownState::InlineCode => {
434                execute!(stdout(), SetForegroundColor(Color::Yellow), Print(text))?;
435            }
436            MarkdownState::Link => {
437                execute!(stdout(), SetForegroundColor(Color::Blue), Print(text))?;
438            }
439            MarkdownState::LinkUrl => {
440                execute!(stdout(), SetForegroundColor(Color::DarkBlue), Print(text))?;
441            }
442            MarkdownState::UnorderedList(_) | MarkdownState::OrderedList(_, _) => {
443                execute!(stdout(), SetForegroundColor(Color::Green), Print(text))?;
444            }
445            MarkdownState::Blockquote(_) => {
446                execute!(stdout(), SetForegroundColor(Color::Cyan), Print(text))?;
447            }
448        }
449
450        Ok(())
451    }
452
453    fn reset_line_states(&mut self) -> Result<()> {
454        // Reset states that shouldn't persist across lines
455        match self.current_state() {
456            MarkdownState::Heading(_) => {
457                self.pop_state();
458                execute!(stdout(), ResetColor)?;
459            }
460            MarkdownState::UnorderedList(_) => {
461                self.pop_state();
462            }
463            MarkdownState::OrderedList(_, _) => {
464                self.pop_state();
465            }
466            MarkdownState::Blockquote(_) => {
467                self.pop_state();
468            }
469            _ => {}
470        }
471
472        Ok(())
473    }
474}
475
476// Helper function to extract response from JSON
477pub fn extract_response(partial: &str) -> Option<&str> {
478    // Find the start of the response field
479    let response_marker = r#""response": ""#;
480    let start_pos = partial.find(response_marker)?;
481    let start_pos = start_pos + response_marker.len();
482
483    // Get substring starting from the response value
484    let response_substring = &partial[start_pos..];
485
486    // Track the JSON string boundaries more carefully
487    let mut depth = 0;
488    let mut is_escaped = false;
489    let mut end_pos = 0;
490
491    for (i, c) in response_substring.char_indices() {
492        if is_escaped {
493            is_escaped = false;
494            continue;
495        }
496
497        if c == '\\' {
498            is_escaped = true;
499            continue;
500        }
501
502        // Only consider a quote as the end if we're at the top level
503        if c == '"' && depth == 0 {
504            end_pos = i;
505            break;
506        }
507
508        // Track balanced pairs for proper parsing
509        match c {
510            '[' => depth += 1,
511            ']' => {
512                if depth > 0 {
513                    depth -= 1
514                }
515            }
516            '{' => depth += 1,
517            '}' => {
518                if depth > 0 {
519                    depth -= 1
520                }
521            }
522            _ => {}
523        }
524    }
525
526    // If we didn't find a closing quote, the JSON might be incomplete
527    if end_pos == 0 {
528        // Return what we have so far
529        Some(response_substring)
530    } else {
531        // Extract just the response value
532        let response = &response_substring[..end_pos];
533
534        // We'll handle unescaping in the renderer
535        Some(response)
536    }
537}
538
539// JSON parsing utility functions
540pub fn contains_end_tag(content: &str) -> bool {
541    // Strip markdown code blocks first
542    let clean_content = strip_markdown_code_blocks(content);
543
544    // Try JSON
545    match serde_json::from_str::<Value>(&clean_content) {
546        Ok(json) => {
547            if json
548                .get("finished")
549                .and_then(Value::as_bool)
550                .unwrap_or(false)
551            {
552                return true;
553            }
554        }
555        Err(_) => {}
556    }
557    return false;
558}
559
560pub fn contains_tool_call(content: &str) -> Option<(String, String)> {
561    // Strip markdown code blocks first
562    let clean_content = strip_markdown_code_blocks(content);
563
564    // Try JSON
565    match serde_json::from_str::<Value>(&clean_content) {
566        Ok(json) => {
567            if let Some(tool) = json.get("tool") {
568                if tool.is_null() {
569                    return None;
570                }
571
572                let tool_name = tool.get("name")?.as_str()?;
573                let tool_content = tool.get("content")?.as_str()?;
574
575                return Some((tool_name.to_string(), tool_content.to_string()));
576            }
577        }
578        Err(_) => {}
579    }
580
581    None
582}
583
584pub fn extract_tool_content(content: &str) -> Option<String> {
585    // Try JSON parsing
586    if let Some(tool_content) = extract_tool_content_json(content) {
587        return Some(tool_content);
588    }
589    None
590}
591
592fn extract_tool_content_json(content: &str) -> Option<String> {
593    // Strip markdown code blocks first
594    let clean_content = strip_markdown_code_blocks(content);
595
596    match serde_json::from_str::<Value>(&clean_content) {
597        Ok(json) => {
598            if let Some(tool) = json.get("tool") {
599                if tool.is_object() && tool.get("name")? == "cli" {
600                    return tool.get("content")?.as_str().map(String::from);
601                }
602            }
603            None
604        }
605        Err(_) => None,
606    }
607}
babel/utils.rs

babel/
utils.rs
