LLM支持优化

src/llm/thinking.rs

@@ -0,0 +1,152 @@
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+
+/// 统一的思考状态管理器，用于管理模型思考状态的显示与隐藏
+pub struct ThinkingStateManager {
+    is_thinking: AtomicBool,
+    on_start: Option<Box<dyn Fn() + Send + Sync>>,
+    on_end: Option<Box<dyn Fn() + Send + Sync>>,
+}
+
+impl ThinkingStateManager {
+    pub fn new() -> Self {
+        Self {
+            is_thinking: AtomicBool::new(false),
+            on_start: None,
+            on_end: None,
+        }
+    }
+
+    /// 设置思考开始回调
+    pub fn on_thinking_start<F: Fn() + Send + Sync + 'static>(mut self, callback: F) -> Self {
+        self.on_start = Some(Box::new(callback));
+        self
+    }
+
+    /// 设置思考结束回调
+    pub fn on_thinking_end<F: Fn() + Send + Sync + 'static>(mut self, callback: F) -> Self {
+        self.on_end = Some(Box::new(callback));
+        self
+    }
+
+    /// 开始思考状态
+    pub fn start_thinking(&self) {
+        if !self.is_thinking.load(Ordering::SeqCst) {
+            self.is_thinking.store(true, Ordering::SeqCst);
+            if let Some(ref cb) = self.on_start {
+                cb();
+            }
+        }
+    }
+
+    /// 结束思考状态
+    pub fn end_thinking(&self) {
+        if self.is_thinking.load(Ordering::SeqCst) {
+            self.is_thinking.store(false, Ordering::SeqCst);
+            if let Some(ref cb) = self.on_end {
+                cb();
+            }
+        }
+    }
+
+    /// 当前是否处于思考状态
+    pub fn is_thinking(&self) -> bool {
+        self.is_thinking.load(Ordering::SeqCst)
+    }
+}
+
+impl Default for ThinkingStateManager {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// 线程安全的思考状态管理器引用
+pub type SharedThinkingState = Arc<ThinkingStateManager>;
+
+/// 创建带有默认控制台输出的思考状态管理器
+/// 在思考开始时打印 "thinking..."，在思考结束时清除该标识
+pub fn create_console_thinking_state() -> SharedThinkingState {
+    Arc::new(
+        ThinkingStateManager::new()
+            .on_thinking_start(|| {
+                eprint!("\rthinking...");
+            })
+            .on_thinking_end(|| {
+                eprint!("\r          \r");
+            }),
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::sync::Mutex;
+
+    #[test]
+    fn test_thinking_state_transitions() {
+        let manager = ThinkingStateManager::new();
+        assert!(!manager.is_thinking());
+
+        manager.start_thinking();
+        assert!(manager.is_thinking());
+
+        manager.end_thinking();
+        assert!(!manager.is_thinking());
+    }
+
+    #[test]
+    fn test_thinking_idempotent_start() {
+        let manager = ThinkingStateManager::new();
+        manager.start_thinking();
+        manager.start_thinking(); // 重复调用不应触发回调两次
+        assert!(manager.is_thinking());
+    }
+
+    #[test]
+    fn test_thinking_idempotent_end() {
+        let manager = ThinkingStateManager::new();
+        manager.end_thinking(); // 未开始时结束不应触发问题
+        assert!(!manager.is_thinking());
+    }
+
+    #[test]
+    fn test_thinking_callbacks() {
+        let events: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
+        let events_clone = events.clone();
+
+        let manager = ThinkingStateManager::new()
+            .on_thinking_start(move || {
+                events_clone.lock().unwrap().push("start".to_string());
+            });
+
+        let events_clone2 = events.clone();
+        let manager = manager.on_thinking_end(move || {
+            events_clone2.lock().unwrap().push("end".to_string());
+        });
+
+        manager.start_thinking();
+        manager.end_thinking();
+
+        let recorded = events.lock().unwrap();
+        assert_eq!(recorded.len(), 2);
+        assert_eq!(recorded[0], "start");
+        assert_eq!(recorded[1], "end");
+    }
+
+    #[test]
+    fn test_create_console_thinking_state() {
+        let state = create_console_thinking_state();
+        assert!(!state.is_thinking());
+        state.start_thinking();
+        assert!(state.is_thinking());
+        state.end_thinking();
+        assert!(!state.is_thinking());
+    }
+
+    #[test]
+    fn test_default() {
+        let manager = ThinkingStateManager::default();
+        assert!(!manager.is_thinking());
+    }
+}