update

Author: ifsheldon

async-openai-wasm/src/client.rs

@@ -446,10 +446,10 @@ impl<C: Config> Client<C> {
         path: &str,
         request: I,
         event_mapper: impl Fn(eventsource_stream::Event) -> Result<O, OpenAIError> + Send + 'static,
-    ) -> Pin<Box<dyn Stream<Item=Result<O, OpenAIError>> + Send>>
+    ) -> OpenAIEventMappedStream<O>
         where
             I: Serialize,
-            O: DeserializeOwned + Send + 'static,
+            O: DeserializeOwned + Send + 'static
     {
         let event_source = self
             .http_client
@@ -460,8 +460,7 @@ impl<C: Config> Client<C> {
             .eventsource()
             .unwrap();
 
-        // stream_mapped_raw_events(event_source, event_mapper).await
-        todo!()
+        OpenAIEventMappedStream::new(event_source, event_mapper)
     }
 
     /// Make HTTP GET request to receive SSE
@@ -491,13 +490,13 @@ impl<C: Config> Client<C> {
 /// Request which responds with SSE.
 /// [server-sent events](https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#event_stream_format)
 #[pin_project]
-pub struct OpenAIEventStream<O> {
+pub struct OpenAIEventStream<O: DeserializeOwned + Send + 'static> {
     #[pin]
     stream: Filter<EventSource, future::Ready<bool>, fn(&Result<Event, reqwest_eventsource::Error>) -> future::Ready<bool>>,
     _phantom_data: PhantomData<O>,
 }
 
-impl<O> OpenAIEventStream<O> {
+impl<O: DeserializeOwned + Send + 'static> OpenAIEventStream<O> {
     pub(crate) fn new(event_source: EventSource) -> Self {
         Self {
             stream: event_source.filter(|result|
@@ -543,6 +542,66 @@ impl<O: DeserializeOwned + Send + 'static> Stream for OpenAIEventStream<O> {
     }
 }
 
+#[pin_project]
+pub struct OpenAIEventMappedStream<O>
+    where O: Send + 'static
+{
+    #[pin]
+    stream: Filter<EventSource, future::Ready<bool>, fn(&Result<Event, reqwest_eventsource::Error>) -> future::Ready<bool>>,
+    event_mapper: Box<dyn Fn(eventsource_stream::Event) -> Result<O, OpenAIError> + Send + 'static>,
+    _phantom_data: PhantomData<O>,
+}
+
+impl<O> OpenAIEventMappedStream<O>
+    where O: Send + 'static
+{
+    pub(crate) fn new<M>(event_source: EventSource, event_mapper: M) -> Self
+        where M: Fn(eventsource_stream::Event) -> Result<O, OpenAIError> + Send + 'static {
+        Self {
+            stream: event_source.filter(|result|
+                // filter out the first event which is always Event::Open
+                future::ready(!(result.is_ok() && result.as_ref().unwrap().eq(&Event::Open)))
+            ),
+            event_mapper: Box::new(event_mapper),
+            _phantom_data: PhantomData,
+        }
+    }
+}
+
+
+impl<O> Stream for OpenAIEventMappedStream<O>
+    where O: Send + 'static
+{
+    type Item = Result<O, OpenAIError>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.project();
+        let stream: Pin<&mut _> = this.stream;
+        match stream.poll_next(cx) {
+            Poll::Ready(response) => {
+                match response {
+                    None => Poll::Ready(None), // end of the stream
+                    Some(result) => match result {
+                        Ok(event) => match event {
+                            Event::Open => unreachable!(), // it has been filtered out
+                            Event::Message(message) => {
+                                if message.data == "[DONE]" {
+                                    Poll::Ready(None)  // end of the stream, defined by OpenAI
+                                } else {
+                                    todo!()
+                                }
+                            }
+                        }
+                        Err(e) => Poll::Ready(Some(Err(OpenAIError::StreamError(e.to_string()))))
+                    }
+                }
+            }
+            Poll::Pending => Poll::Pending
+        }
+    }
+}
+
+
 // pub(crate) async fn stream_mapped_raw_events<O>(
 //     mut event_source: EventSource,
 //     event_mapper: impl Fn(eventsource_stream::Event) -> Result<O, OpenAIError> + Send + 'static,

async-openai-wasm/src/types/assistant_stream.rs

@@ -1,8 +1,6 @@
-use std::pin::Pin;
-
-use futures::Stream;
 use serde::Deserialize;
 
+use crate::client::OpenAIEventMappedStream;
 use crate::error::{ApiError, map_deserialization_error, OpenAIError};
 
 use super::{
@@ -28,7 +26,6 @@ use super::{
 /// We may add additional events over time, so we recommend handling unknown events gracefully
 /// in your code. See the [Assistants API quickstart](https://platform.openai.com/docs/assistants/overview) to learn how to
 /// integrate the Assistants API with streaming.
-
 #[derive(Debug, Deserialize, Clone)]
 #[serde(tag = "event", content = "data")]
 #[non_exhaustive]
@@ -110,8 +107,7 @@ pub enum AssistantStreamEvent {
     Done(String),
 }
 
-pub type AssistantEventStream =
-    Pin<Box<dyn Stream<Item = Result<AssistantStreamEvent, OpenAIError>> + Send>>;
+pub type AssistantEventStream = OpenAIEventMappedStream<AssistantStreamEvent>;
 
 impl TryFrom<eventsource_stream::Event> for AssistantStreamEvent {
     type Error = OpenAIError;