Refactor event source handling: Introduce NdJsonStream and Stream trait for improved stream management

Author: Rutik Thakre

llm/src/event_source/error.rs

@@ -1,6 +1,5 @@
-use crate::event_source::event_stream::EventStreamError;
 use core::fmt;
-use golem_rust::bindings::wasi::io::streams::StreamError;
+use golem_rust::bindings::wasi::io::streams::StreamError as WasiStreamError;
 use nom::error::Error as NomError;
 use reqwest::header::HeaderValue;
 use reqwest::Error as ReqwestError;
@@ -9,6 +8,8 @@ use reqwest::StatusCode;
 use std::string::FromUtf8Error;
 use thiserror::Error;
 
+use super::stream::StreamError;
+
 /// Error raised when a [`RequestBuilder`] cannot be cloned. See [`RequestBuilder::try_clone`] for
 /// more information
 #[derive(Debug, Clone, Copy)]
@@ -51,24 +52,24 @@ pub enum Error {
     StreamEnded,
 }
 
-impl From<EventStreamError<ReqwestError>> for Error {
-    fn from(err: EventStreamError<ReqwestError>) -> Self {
+impl From<StreamError<ReqwestError>> for Error {
+    fn from(err: StreamError<ReqwestError>) -> Self {
         match err {
-            EventStreamError::Utf8(err) => Self::Utf8(err),
-            EventStreamError::Parser(err) => Self::Parser(err),
-            EventStreamError::Transport(err) => Self::Transport(err),
+            StreamError::Utf8(err) => Self::Utf8(err),
+            StreamError::Parser(err) => Self::Parser(err),
+            StreamError::Transport(err) => Self::Transport(err),
         }
     }
 }
 
-impl From<EventStreamError<StreamError>> for Error {
-    fn from(err: EventStreamError<StreamError>) -> Self {
+impl From<StreamError<WasiStreamError>> for Error {
+    fn from(err: StreamError<WasiStreamError>) -> Self {
         match err {
-            EventStreamError::Utf8(err) => Self::Utf8(err),
-            EventStreamError::Parser(err) => Self::Parser(err),
-            EventStreamError::Transport(err) => match err {
-                StreamError::Closed => Self::StreamEnded,
-                StreamError::LastOperationFailed(err) => {
+            StreamError::Utf8(err) => Self::Utf8(err),
+            StreamError::Parser(err) => Self::Parser(err),
+            StreamError::Transport(err) => match err {
+                WasiStreamError::Closed => Self::StreamEnded,
+                WasiStreamError::LastOperationFailed(err) => {
                     Self::TransportStream(err.to_debug_string())
                 }
             },

llm/src/event_source/event_stream.rs

@@ -1,15 +1,14 @@
 use crate::event_source::parser::{is_bom, is_lf, line, RawEventLine};
 use crate::event_source::utf8_stream::{Utf8Stream, Utf8StreamError};
 use crate::event_source::MessageEvent;
-use core::fmt;
 use core::time::Duration;
 use golem_rust::bindings::wasi::io::streams::{InputStream, StreamError};
 use golem_rust::wasm_rpc::Pollable;
 use log::trace;
-use nom::error::Error as NomError;
-use std::string::FromUtf8Error;
 use std::task::Poll;
 
+use super::stream::{Stream, StreamError as EventStreamError};
+
 #[derive(Default, Debug)]
 struct EventBuilder {
     event: MessageEvent,
@@ -133,9 +132,9 @@ pub struct EventStream {
     last_event_id: String,
 }
 
-impl EventStream {
+impl Stream  for EventStream {
     /// Initialize the EventStream with a Stream
-    pub fn new(stream: InputStream) -> Self {
+     fn new(stream: InputStream) -> Self {
         println!("EventStream::new");
         Self {
             stream: Utf8Stream::new(stream),
@@ -148,23 +147,23 @@ impl EventStream {
 
     /// Set the last event ID of the stream. Useful for initializing the stream with a previous
     /// last event ID
-    pub fn set_last_event_id(&mut self, id: impl Into<String>) {
+     fn set_last_event_id(&mut self, id: impl Into<String>) {
         println!("EventStream::set_last_event_id");
         self.last_event_id = id.into();
     }
 
     /// Get the last event ID of the stream
-    pub fn last_event_id(&self) -> &str {
+     fn last_event_id(&self) -> &str {
         println!("EventStream::last_event_id");
         &self.last_event_id
     }
 
-    pub fn subscribe(&self) -> Pollable {
+     fn subscribe(&self) -> Pollable {
         println!("EventStream::subscribe");
         self.stream.subscribe()
     }
 
-    pub fn poll_next(
+     fn poll_next(
         &mut self,
     ) -> Poll<Option<Result<MessageEvent, EventStreamError<StreamError>>>> {
         println!("EventStream::poll_next buffer {} ", self.buffer.as_str());
@@ -222,50 +221,6 @@ impl EventStream {
     }
 }
 
-/// Error thrown while parsing an event line
-#[derive(Debug, PartialEq)]
-pub enum EventStreamError<E> {
-    /// Source stream is not valid UTF8
-    Utf8(FromUtf8Error),
-    /// Source stream is not a valid EventStream
-    Parser(NomError<String>),
-    /// Underlying source stream error
-    Transport(E),
-}
-
-impl<E> From<Utf8StreamError<E>> for EventStreamError<E> {
-    fn from(err: Utf8StreamError<E>) -> Self {
-        println!("EventStreamError::from");
-        match err {
-            Utf8StreamError::Utf8(err) => Self::Utf8(err),
-            Utf8StreamError::Transport(err) => Self::Transport(err),
-        }
-    }
-}
-
-impl<E> From<NomError<&str>> for EventStreamError<E> {
-    fn from(err: NomError<&str>) -> Self {
-        println!("EventStreamError::from");
-        EventStreamError::Parser(NomError::new(err.input.to_string(), err.code))
-    }
-}
-
-impl<E> fmt::Display for EventStreamError<E>
-where
-    E: fmt::Display,
-{
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        println!("EventStreamError::fmt");
-        match self {
-            Self::Utf8(err) => f.write_fmt(format_args!("UTF8 error: {}", err)),
-            Self::Parser(err) => f.write_fmt(format_args!("Parse error: {}", err)),
-            Self::Transport(err) => f.write_fmt(format_args!("Transport error: {}", err)),
-        }
-    }
-}
-
-impl<E> std::error::Error for EventStreamError<E> where E: fmt::Display + fmt::Debug + Send + Sync {}
-
 fn parse_event<E>(
     buffer: &mut String,
     builder: &mut EventBuilder,

llm/src/event_source/mod.rs

@@ -4,16 +4,20 @@
 pub mod error;
 mod event_stream;
 mod message_event;
+mod ndjson_stream;
 mod parser;
+mod stream;
 mod utf8_stream;
 
 use crate::event_source::error::Error;
 use crate::event_source::event_stream::EventStream;
 use golem_rust::wasm_rpc::Pollable;
 pub use message_event::MessageEvent;
+use ndjson_stream::NdJsonStream;
 use reqwest::header::HeaderValue;
 use reqwest::{Response, StatusCode};
 use std::task::Poll;
+use stream::{Stream, StreamType};
 
 /// The ready state of an [`EventSource`]
 #[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd)]
@@ -28,7 +32,8 @@ pub enum ReadyState {
 }
 
 pub struct EventSource {
-    stream: EventStream,
+    /// stream is the type which implements Stream trait
+    stream: StreamType,
     response: Response,
     is_closed: bool,
 }
@@ -45,7 +50,19 @@ impl EventSource {
                         golem_rust::bindings::wasi::io::streams::InputStream,
                     >(response.get_raw_input_stream())
                 };
-                let stream = EventStream::new(handle);
+
+                let stream = if response
+                    .headers()
+                    .get(&reqwest::header::CONTENT_TYPE)
+                    .unwrap()
+                    .to_str()
+                    .unwrap()
+                    .contains("ndjson")
+                {
+                    StreamType::NdJsonStream(NdJsonStream::new(handle))
+                } else {
+                    StreamType::EventStream(EventStream::new(handle))
+                };
                 Ok(Self {
                     response,
                     stream,
@@ -74,7 +91,10 @@ impl EventSource {
 
     pub fn subscribe(&self) -> Pollable {
         println!("EventSource::subscribe");
-        self.stream.subscribe()
+        match &self.stream {
+            StreamType::EventStream(stream) => stream.subscribe(),
+            StreamType::NdJsonStream(stream) => stream.subscribe(),
+        }
     }
 
     pub fn poll_next(&mut self) -> Poll<Option<Result<Event, Error>>> {
@@ -83,19 +103,23 @@ impl EventSource {
             return Poll::Ready(None);
         }
 
-        match self.stream.poll_next() {
-            Poll::Ready(Some(Err(err))) => {
-                let err = err.into();
-                self.is_closed = true;
-                Poll::Ready(Some(Err(err)))
+        match &mut self.stream {
+            StreamType::EventStream( stream) => {
+                match stream.poll_next() {
+                    Poll::Ready(Some(Ok(event))) => Poll::Ready(Some(Ok(Event::Message(event)))),
+                    Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(err.into()))),
+                    Poll::Ready(None) => Poll::Ready(None),
+                    Poll::Pending => Poll::Pending,
+                }
             }
-            Poll::Ready(Some(Ok(event))) => Poll::Ready(Some(Ok(event.into()))),
-            Poll::Ready(None) => {
-                let err = Error::StreamEnded;
-                self.is_closed = true;
-                Poll::Ready(Some(Err(err)))
+            StreamType::NdJsonStream(stream) => {
+                match stream.poll_next() {
+                    Poll::Ready(Some(Ok(event))) => Poll::Ready(Some(Ok(Event::Message(event)))),
+                    Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(err.into()))),
+                    Poll::Ready(None) => Poll::Ready(None),
+                    Poll::Pending => Poll::Pending,
+                }
             }
-            Poll::Pending => Poll::Pending,
         }
     }
 }
@@ -126,7 +150,7 @@ fn check_response(response: Response) -> Result<Response, Error> {
             matches!(
                 (mime_type.type_(), mime_type.subtype()),
                 (mime::TEXT, mime::EVENT_STREAM)
-            ) || mime_type.subtype().as_str().contains("ndjson") 
+            ) || mime_type.subtype().as_str().contains("ndjson")
         })
         .unwrap_or(false)
     {