From 700d8a3245f1133c37e035039e511a1e5dce1a5d Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Mon, 3 Feb 2025 22:35:01 +0800
Subject: feat: strip reasoning contents (#1141)

---
 src/utils/mod.rs | 6 ++++++
 1 file changed, 6 insertions(+)

(limited to 'src/utils/mod.rs')

diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index 9dde751..049f361 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -26,11 +26,13 @@ use anyhow::{Context, Result};
 use fancy_regex::Regex;
 use fuzzy_matcher::{skim::SkimMatcherV2, FuzzyMatcher};
 use is_terminal::IsTerminal;
+use std::borrow::Cow;
 use std::{env, path::PathBuf, process};
 use unicode_segmentation::UnicodeSegmentation;
 
 lazy_static::lazy_static! {
     pub static ref CODE_BLOCK_RE: Regex = Regex::new(r"(?ms)```\w*(.*)```").unwrap();
+    pub static ref THINK_TAG_RE: Regex = Regex::new(r"(?s)^\s*<think>.*?</think>(\s*|$)").unwrap();
     pub static ref IS_STDOUT_TERMINAL: bool = std::io::stdout().is_terminal();
     pub static ref NO_COLOR: bool = env::var("NO_COLOR").ok().and_then(|v| parse_bool(&v)).unwrap_or_default() || !*IS_STDOUT_TERMINAL;
 }
@@ -59,6 +61,10 @@ pub fn parse_bool(value: &str) -> Option<bool> {
     }
 }
 
+pub fn strip_think_tag(text: &str) -> Cow<str> {
+    THINK_TAG_RE.replace_all(text, "")
+}
+
 pub fn estimate_token_length(text: &str) -> usize {
     let words: Vec<&str> = text.unicode_words().collect();
     let mut output: f32 = 0.0;
-- 
cgit v1.2.3