aix_core/

streaming.rs

//! Streaming abstractions and utilities.
//!
//! This module provides types and utilities for handling streaming responses
//! from AI providers, including stream adapters and convenience functions.

use crate::error::{AixError, AixResult};
use crate::types::StreamChunk;
use futures_core::Stream;
use pin_project_lite::pin_project;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;

/// Type alias for a stream of stream chunks.
pub type TokenStream = Pin<Box<dyn Stream<Item = AixResult<StreamChunk>> + Send>>;

/// Extension trait for streams with convenience methods.
pub trait StreamExt: Stream {
    /// Collect all chunks into a single string.
    ///
    /// This method consumes the stream and collects all the delta content
    /// from each chunk into a single string.
    ///
    /// # Returns
    /// A future that resolves to the complete text or an error
    fn collect_text(self) -> CollectText<Self>
    where
        Self: Sized,
    {
        CollectText::new(self)
    }

    /// Filter out chunks with empty delta content.
    ///
    /// # Returns
    /// A stream that only yields chunks with non-empty delta content
    fn filter_empty(self) -> FilterEmpty<Self>
    where
        Self: Sized,
    {
        FilterEmpty::new(self)
    }

    /// Buffer chunks for a given duration before yielding them.
    ///
    /// This can be useful to reduce the frequency of updates in UI applications.
    ///
    /// # Arguments
    /// * `duration` - The buffer duration
    ///
    /// # Returns
    /// A stream that yields buffered chunks
    fn buffer_chunks(self, duration: Duration) -> BufferChunks<Self>
    where
        Self: Sized,
    {
        BufferChunks::new(self, duration)
    }
}

// Blanket implementation for all streams
impl<T: ?Sized> StreamExt for T where T: Stream {}

/// Stream adapter that collects all text from chunks.
pin_project! {
    pub struct CollectText<S> {
        #[pin]
        stream: S,
        buffer: String,
    }
}

impl<S> CollectText<S> {
    fn new(stream: S) -> Self {
        Self {
            stream,
            buffer: String::new(),
        }
    }
}

impl<S> std::future::Future for CollectText<S>
where
    S: Stream<Item = AixResult<StreamChunk>>,
{
    type Output = AixResult<String>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut this = self.project();
        
        loop {
            match futures_core::ready!(this.stream.as_mut().poll_next(cx)) {
                Some(Ok(chunk)) => {
                    this.buffer.push_str(&chunk.delta);
                }
                Some(Err(error)) => {
                    return Poll::Ready(Err(error));
                }
                None => {
                    return Poll::Ready(Ok(this.buffer.clone()));
                }
            }
        }
    }
}

/// Stream adapter that filters out empty chunks.
pin_project! {
    pub struct FilterEmpty<S> {
        #[pin]
        stream: S,
    }
}

impl<S> FilterEmpty<S> {
    fn new(stream: S) -> Self {
        Self { stream }
    }
}

impl<S> Stream for FilterEmpty<S>
where
    S: Stream<Item = AixResult<StreamChunk>>,
{
    type Item = AixResult<StreamChunk>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let mut this = self.project();
        
        loop {
            match futures_core::ready!(this.stream.as_mut().poll_next(cx)) {
                Some(Ok(chunk)) => {
                    if chunk.delta.is_empty() && chunk.finish_reason.is_none() {
                        // Skip empty chunks without finish reason
                        continue;
                    }
                    return Poll::Ready(Some(Ok(chunk)));
                }
                other => return Poll::Ready(other),
            }
        }
    }
}

/// Stream adapter that buffers chunks for a duration.
pin_project! {
    pub struct BufferChunks<S> {
        #[pin]
        stream: S,
        buffer: Vec<StreamChunk>,
        last_flush: Option<tokio::time::Instant>,
        duration: Duration,
        #[pin]
        delay: Option<tokio::time::Sleep>,
    }
}

impl<S> BufferChunks<S> {
    fn new(stream: S, duration: Duration) -> Self {
        Self {
            stream,
            buffer: Vec::new(),
            last_flush: None,
            duration,
            delay: None,
        }
    }
}

impl<S> Stream for BufferChunks<S>
where
    S: Stream<Item = AixResult<StreamChunk>>,
{
    type Item = AixResult<StreamChunk>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let mut this = self.project();
        let now = tokio::time::Instant::now();

        // Check if we have buffered chunks and it's time to flush
        if !this.buffer.is_empty() {
            let should_flush = if let Some(last_flush) = this.last_flush {
                now.duration_since(*last_flush) >= *this.duration
            } else {
                true // Flush immediately on first chunk
            };

            if should_flush {
                // Combine all buffered chunks into one
                let combined_id = this.buffer
                    .first()
                    .map(|c| c.id.clone())
                    .unwrap_or_else(|| "buffered".to_string());
                
                let combined_delta: String = this.buffer
                    .iter()
                    .map(|c| c.delta.as_str())
                    .collect();
                
                let finish_reason = this.buffer
                    .iter()
                    .find_map(|c| c.finish_reason.clone());

                let combined_chunk = StreamChunk {
                    id: combined_id,
                    delta: combined_delta,
                    finish_reason,
                };

                this.buffer.clear();
                *this.last_flush = Some(now);

                return Poll::Ready(Some(Ok(combined_chunk)));
            }
        }

        // Poll for new chunks
        match futures_core::ready!(this.stream.as_mut().poll_next(cx)) {
            Some(Ok(chunk)) => {
                // Buffer the chunk
                this.buffer.push(chunk);

                // If this is the first chunk, set up the delay
                if this.delay.is_none() {
                    this.delay.set(Some(tokio::time::sleep(*this.duration)));
                }

                // Try to poll the delay
                if let Some(delay) = this.delay.as_mut().as_pin_mut() {
                    match delay.poll(cx) {
                        std::task::Poll::Ready(_) => {
                            this.delay.set(None);
                            // The next poll will flush the buffer
                        }
                        std::task::Poll::Pending => {
                            // Still waiting
                        }
                    }
                }

                Poll::Pending
            }
            Some(Err(error)) => {
                Poll::Ready(Some(Err(error)))
            }
            None => {
                // Stream ended, flush any remaining chunks
                if !this.buffer.is_empty() {
                    let combined_id = this.buffer
                        .first()
                        .map(|c| c.id.clone())
                        .unwrap_or_else(|| "buffered".to_string());
                    
                    let combined_delta: String = this.buffer
                        .iter()
                        .map(|c| c.delta.as_str())
                        .collect();
                    
                    let finish_reason = this.buffer
                        .iter()
                        .find_map(|c| c.finish_reason.clone());

                    let combined_chunk = StreamChunk {
                        id: combined_id,
                        delta: combined_delta,
                        finish_reason,
                    };

                    this.buffer.clear();

                    Poll::Ready(Some(Ok(combined_chunk)))
                } else {
                    Poll::Ready(None)
                }
            }
        }
    }
}

/// Create a stream from an iterator of results.
pub fn from_iter<I>(iter: I) -> TokenStream
where
    I: IntoIterator<Item = AixResult<StreamChunk>>,
    I::IntoIter: Send + 'static,
{
    let stream = futures_util::stream::iter(iter);
    Box::pin(stream)
}

/// Create a stream that immediately yields an error.
pub fn error_stream(error: AixError) -> TokenStream {
    let stream = futures_util::stream::once(async move { Err(error) });
    Box::pin(stream)
}

/// Create a stream that yields a single chunk.
pub fn single_chunk(chunk: StreamChunk) -> TokenStream {
    let stream = futures_util::stream::once(async move { Ok(chunk) });
    Box::pin(stream)
}

/// Create a stream that yields multiple chunks.
pub fn chunks<I>(chunks: I) -> TokenStream
where
    I: IntoIterator<Item = StreamChunk>,
    I::IntoIter: Send + 'static,
{
    let results = chunks.into_iter().map(Ok);
    from_iter(results)
}

/// Create a stream from a string that yields character by character.
pub fn from_string<S>(id: S, text: S) -> TokenStream
where
    S: Into<String> + Clone,
{
    let id = id.into();
    let text = text.into();
    let chars: Vec<char> = text.chars().collect();
    let stream = futures_util::stream::iter(chars.into_iter().map(move |c| {
        let id = id.clone();
        Ok(StreamChunk::new(id, c.to_string()))
    }));
    Box::pin(stream)
}

/// Create a stream from a string that yields word by word.
pub fn from_string_words<S>(id: S, text: S) -> TokenStream
where
    S: Into<String> + Clone,
{
    let id = id.into();
    let text = text.into();
    let words: Vec<String> = text.split_whitespace().map(|s| s.to_string()).collect();
    let stream = futures_util::stream::iter(words.into_iter().map(move |word| {
        let id = id.clone();
        Ok(StreamChunk::new(id, format!("{} ", word)))
    }));
    Box::pin(stream)
}

/// Utility for parsing Server-Sent Events (SSE).
pub struct SseParser {
    buffer: String,
}

impl SseParser {
    /// Create a new SSE parser.
    pub fn new() -> Self {
        Self {
            buffer: String::new(),
        }
    }

    /// Parse a chunk of SSE data.
    ///
    /// # Arguments
    /// * `chunk` - A chunk of bytes from the SSE stream
    ///
    /// # Returns
    /// A vector of parsed events, or an error if parsing fails
    pub fn parse_chunk(&mut self, chunk: &[u8]) -> AixResult<Vec<String>> {
        let chunk_str = std::str::from_utf8(chunk)
            .map_err(|e| AixError::serialization(e.to_string(), "SSE chunk parsing"))?;

        self.buffer.push_str(chunk_str);
        self.extract_events()
    }

    /// Extract complete events from the buffer.
    fn extract_events(&mut self) -> AixResult<Vec<String>> {
        let mut events = Vec::new();
        let mut lines = self.buffer.lines().peekable();

        while let Some(line) = lines.next() {
            if line.starts_with("data:") {
                let mut event_data = line[5..].trim().to_string();
                
                // Look for additional data lines
                while let Some(&next_line) = lines.peek() {
                    if next_line.starts_with("data:") {
                        event_data.push_str(&next_line[5..].trim());
                        lines.next(); // Consume the line
                    } else {
                        break;
                    }
                }

                // Check if this is the end of an event (empty line or [DONE])
                if event_data == "[DONE]" {
                    events.push("[DONE]".to_string());
                } else if !event_data.is_empty() {
                    events.push(event_data);
                }
            }
        }

        // Clear processed data from buffer
        // Keep any incomplete data that might be waiting for more chunks
        let last_complete_pos = self.buffer.rfind("\n\n").unwrap_or(0);
        if last_complete_pos > 0 {
            self.buffer.drain(0..=last_complete_pos + 1);
        }

        Ok(events)
    }

    /// Get any remaining data in the buffer.
    pub fn remaining_data(&self) -> &str {
        &self.buffer
    }

    /// Clear the buffer.
    pub fn clear(&mut self) {
        self.buffer.clear();
    }
}

impl Default for SseParser {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures_util::StreamExt as FuturesStreamExt;

    #[tokio::test]
    async fn test_collect_text() {
        let chunks = vec![
            Ok(StreamChunk::new("1", "Hello")),
            Ok(StreamChunk::new("2", ", ")),
            Ok(StreamChunk::new("3", "world")),
            Ok(StreamChunk::new("4", "!")),
        ];

        let stream = from_iter(chunks);
        let text = stream.collect_text().await.unwrap();
        assert_eq!(text, "Hello, world!");
    }

    #[tokio::test]
    async fn test_filter_empty() {
        let chunks = vec![
            Ok(StreamChunk::new("1", "Hello")),
            Ok(StreamChunk::new("2", "")), // Should be filtered out
            Ok(StreamChunk::new("3", "world")),
            Ok(StreamChunk::new("4", "")), // Should be filtered out
        ];

        let stream = from_iter(chunks).filter_empty();
        let collected: Vec<_> = stream.collect().await;
        
        assert_eq!(collected.len(), 2);
        assert_eq!(collected[0].as_ref().unwrap().delta, "Hello");
        assert_eq!(collected[1].as_ref().unwrap().delta, "world");
    }

    #[tokio::test]
    async fn test_from_string() {
        let stream = from_string("test", "Hello world");
        let collected: Vec<_> = stream.collect().await;
        
        assert_eq!(collected.len(), 11); // "Hello world" + space = 11 chars
        assert_eq!(collected[0].as_ref().unwrap().delta, "H");
        assert_eq!(collected[1].as_ref().unwrap().delta, "e");
    }

    #[tokio::test]
    async fn test_from_string_words() {
        let stream = from_string_words("test", "Hello world from Rust");
        let collected: Vec<_> = stream.collect().await;
        
        assert_eq!(collected.len(), 4);
        assert_eq!(collected[0].as_ref().unwrap().delta, "Hello ");
        assert_eq!(collected[1].as_ref().unwrap().delta, "world ");
        assert_eq!(collected[2].as_ref().unwrap().delta, "from ");
        assert_eq!(collected[3].as_ref().unwrap().delta, "Rust");
    }

    #[test]
    fn test_sse_parser() {
        let mut parser = SseParser::new();
        
        // Test parsing a complete event
        let chunk = b"data: {\"content\": \"Hello\"}\n\n";
        let events = parser.parse_chunk(chunk).unwrap();
        assert_eq!(events.len(), 1);
        assert_eq!(events[0], "{\"content\": \"Hello\"}");

        // Test parsing [DONE] event
        let chunk = b"data: [DONE]\n\n";
        let events = parser.parse_chunk(chunk).unwrap();
        assert_eq!(events.len(), 1);
        assert_eq!(events[0], "[DONE]");
    }

    #[test]
    fn test_sse_parser_incomplete_event() {
        let mut parser = SseParser::new();
        
        // Send incomplete event
        let chunk = b"data: {\"content\":";
        let events = parser.parse_chunk(chunk).unwrap();
        assert_eq!(events.len(), 0); // Should not yield events yet

        // Complete the event
        let chunk = b" \"Hello\"}\n\n";
        let events = parser.parse_chunk(chunk).unwrap();
        assert_eq!(events.len(), 1);
        assert_eq!(events[0], "{\"content\": \"Hello\"}");
    }

    #[test]
    fn test_sse_parser_multiple_events() {
        let mut parser = SseParser::new();
        
        let chunk = b"data: {\"content\": \"Hello\"}\n\ndata: {\"content\": \"world\"}\n\ndata: [DONE]\n\n";
        let events = parser.parse_chunk(chunk).unwrap();
        assert_eq!(events.len(), 3);
        assert_eq!(events[0], "{\"content\": \"Hello\"}");
        assert_eq!(events[1], "{\"content\": \"world\"}");
        assert_eq!(events[2], "[DONE]");
    }

    #[tokio::test]
    async fn test_error_stream() {
        let error = AixError::other("test error");
        let stream = error_stream(error);
        let collected: Vec<_> = stream.collect().await;
        
        assert_eq!(collected.len(), 1);
        assert!(collected[0].is_err());
        assert_eq!(collected[0].as_ref().unwrap_err().to_string(), "Error: test error");
    }

    #[tokio::test]
    async fn test_single_chunk() {
        let chunk = StreamChunk::new("test", "Hello");
        let stream = single_chunk(chunk);
        let collected: Vec<_> = stream.collect().await;
        
        assert_eq!(collected.len(), 1);
        assert_eq!(collected[0].as_ref().unwrap().delta, "Hello");
    }
}