llm_worker/timeline/
text_block_collector.rs

1//! TextBlockCollector - テキストブロック収集用ハンドラ
2//!
3//! TimelineのTextBlockHandler として登録され、
4//! ストリーム中のテキストブロックを収集する。
5
6use crate::handler::{Handler, TextBlockEvent, TextBlockKind};
7use std::sync::{Arc, Mutex};
8
9/// TextBlockから収集したテキスト情報を保持
10#[derive(Debug, Default)]
11pub struct TextCollectorState {
12    /// 蓄積中のテキスト
13    buffer: String,
14}
15
16/// TextBlockCollector - テキストブロックハンドラ
17///
18/// Timelineに登録してTextBlockイベントを受信し、
19/// 完了したテキストブロックを収集する。
20#[derive(Clone)]
21pub struct TextBlockCollector {
22    /// 収集されたテキストブロック
23    collected: Arc<Mutex<Vec<String>>>,
24}
25
26impl TextBlockCollector {
27    /// 新しいTextBlockCollectorを作成
28    pub fn new() -> Self {
29        Self {
30            collected: Arc::new(Mutex::new(Vec::new())),
31        }
32    }
33
34    /// 収集されたテキストを取得してクリア
35    pub fn take_collected(&self) -> Vec<String> {
36        let mut guard = self.collected.lock().unwrap();
37        std::mem::take(&mut *guard)
38    }
39
40    /// 収集されたテキストの参照を取得
41    pub fn collected(&self) -> Vec<String> {
42        self.collected.lock().unwrap().clone()
43    }
44
45    /// 収集されたテキストがあるかどうか
46    pub fn has_content(&self) -> bool {
47        !self.collected.lock().unwrap().is_empty()
48    }
49
50    /// 収集をクリア
51    pub fn clear(&self) {
52        self.collected.lock().unwrap().clear();
53    }
54}
55
56impl Default for TextBlockCollector {
57    fn default() -> Self {
58        Self::new()
59    }
60}
61
62impl Handler<TextBlockKind> for TextBlockCollector {
63    type Scope = TextCollectorState;
64
65    fn on_event(&mut self, scope: &mut Self::Scope, event: &TextBlockEvent) {
66        match event {
67            TextBlockEvent::Start(_) => {
68                scope.buffer.clear();
69            }
70            TextBlockEvent::Delta(text) => {
71                scope.buffer.push_str(text);
72            }
73            TextBlockEvent::Stop(_) => {
74                // ブロック完了時にテキストを確定
75                if !scope.buffer.is_empty() {
76                    let text = std::mem::take(&mut scope.buffer);
77                    self.collected.lock().unwrap().push(text);
78                }
79            }
80        }
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87    use crate::timeline::Timeline;
88    use crate::timeline::event::Event;
89
90    /// TextBlockCollectorが単一のテキストブロックを正しく収集することを確認
91    #[test]
92    fn test_collect_single_text_block() {
93        let collector = TextBlockCollector::new();
94        let mut timeline = Timeline::new();
95        timeline.on_text_block(collector.clone());
96
97        // テキストブロックのイベントシーケンスをディスパッチ
98        timeline.dispatch(&Event::text_block_start(0));
99        timeline.dispatch(&Event::text_delta(0, "Hello, "));
100        timeline.dispatch(&Event::text_delta(0, "World!"));
101        timeline.dispatch(&Event::text_block_stop(0, None));
102
103        // 収集されたテキストを確認
104        let texts = collector.take_collected();
105        assert_eq!(texts.len(), 1);
106        assert_eq!(texts[0], "Hello, World!");
107    }
108
109    /// TextBlockCollectorが複数のテキストブロックを正しく収集することを確認
110    #[test]
111    fn test_collect_multiple_text_blocks() {
112        let collector = TextBlockCollector::new();
113        let mut timeline = Timeline::new();
114        timeline.on_text_block(collector.clone());
115
116        // 1つ目のテキストブロック
117        timeline.dispatch(&Event::text_block_start(0));
118        timeline.dispatch(&Event::text_delta(0, "First"));
119        timeline.dispatch(&Event::text_block_stop(0, None));
120
121        // 2つ目のテキストブロック
122        timeline.dispatch(&Event::text_block_start(1));
123        timeline.dispatch(&Event::text_delta(1, "Second"));
124        timeline.dispatch(&Event::text_block_stop(1, None));
125
126        let texts = collector.take_collected();
127        assert_eq!(texts.len(), 2);
128        assert_eq!(texts[0], "First");
129        assert_eq!(texts[1], "Second");
130    }
131}
llm_worker/timeline/text_block_collector.rs

llm_worker/timeline/
text_block_collector.rs
