jules_core/streaming/
mod.rs

1//! Streaming abstractions for incremental responses.
2
3/// Reconnection logic for streams.
4pub mod reconnect;
5use serde::{Deserialize, Serialize};
6use std::future::Future;
7
8/// Represents an event in a stream of incremental responses.
9#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
10pub enum StreamEvent {
11    /// A chunk of text yielded incrementally.
12    TextChunk(String),
13    /// The stream has finished successfully.
14    Done,
15}
16
17/// A stream of incremental responses.
18pub trait Stream {
19    /// The type of items yielded by the stream.
20    type Item;
21
22    /// Attempts to pull out the next value of this stream, returning `None` if the stream is exhausted.
23    fn next(&mut self) -> impl Future<Output = Option<Self::Item>> + Send;
24}
25
26#[cfg(test)]
27mod tests {
28    use super::*;
29
30    struct MockStream {
31        count: u32,
32    }
33
34    impl Stream for MockStream {
35        type Item = StreamEvent;
36
37        fn next(&mut self) -> impl Future<Output = Option<Self::Item>> + Send {
38            let res = match self.count {
39                0..=2 => {
40                    self.count += 1;
41                    Some(StreamEvent::TextChunk(format!("chunk {}", self.count)))
42                }
43                3 => {
44                    self.count += 1;
45                    Some(StreamEvent::Done)
46                }
47                _ => None,
48            };
49            async move { res }
50        }
51    }
52
53    #[tokio::test]
54    async fn test_mock_stream() {
55        let mut stream = MockStream { count: 0 };
56        assert_eq!(
57            stream.next().await,
58            Some(StreamEvent::TextChunk("chunk 1".to_string()))
59        );
60        assert_eq!(
61            stream.next().await,
62            Some(StreamEvent::TextChunk("chunk 2".to_string()))
63        );
64        assert_eq!(
65            stream.next().await,
66            Some(StreamEvent::TextChunk("chunk 3".to_string()))
67        );
68        assert_eq!(stream.next().await, Some(StreamEvent::Done));
69        assert_eq!(stream.next().await, None);
70    }
71}
jules_core/streaming/mod.rs

jules_core/streaming/
mod.rs
