diff --git a/backend/Cargo.lock b/backend/Cargo.lock index c0129f7f57..d0b443e365 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -5144,6 +5144,17 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "eventsource-stream" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74fef4569247a5f429d9156b9d0a2599914385dd189c539334c625d8099d90ab" +dependencies = [ + "futures-core", + "nom 7.1.3", + "pin-project-lite", +] + [[package]] name = "fallible-iterator" version = "0.2.0" @@ -15777,6 +15788,7 @@ dependencies = [ "deno_webidl", "dotenv", "dyn-iter", + "eventsource-stream", "flume", "futures", "gcp_auth", diff --git a/backend/Cargo.toml b/backend/Cargo.toml index 41e33efd9a..86c994dd8b 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -252,6 +252,7 @@ urlencoding = "^2" url = { version = "^2" , features = ["serde"]} async-oauth2 = "0.5.1" reqwest = { version = "^0.12", features = ["json", "stream", "gzip", "multipart"] } +eventsource-stream = "0.2.3" time = "^0" serde_urlencoded = "^0" tokio-tar = "^0" diff --git a/backend/windmill-worker/Cargo.toml b/backend/windmill-worker/Cargo.toml index f36d2c5d58..d7e0f66685 100644 --- a/backend/windmill-worker/Cargo.toml +++ b/backend/windmill-worker/Cargo.toml @@ -115,6 +115,7 @@ nix.workspace = true bytes.workspace = true reqwest.workspace = true reqwest-middleware.workspace = true +eventsource-stream.workspace = true mime_guess.workspace = true hex.workspace = true tiberius = { workspace = true, optional = true } diff --git a/backend/windmill-worker/src/ai/sse.rs b/backend/windmill-worker/src/ai/sse.rs index 42c5d73603..d34bf00e71 100644 --- a/backend/windmill-worker/src/ai/sse.rs +++ b/backend/windmill-worker/src/ai/sse.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; +use eventsource_stream::Eventsource; use reqwest::Response; use serde::Deserialize; use serde_json; @@ -51,38 +52,23 @@ pub trait SSEParser { async fn parse_event_data(&mut self, data: &str) -> Result<(), Error>; async fn parse_events(&mut self, response: Response) -> Result<(), Error> { - let mut stream = response.bytes_stream(); - let mut buffer = String::new(); + let mut stream = response.bytes_stream().eventsource(); - while let Some(chunk_result) = stream.next().await { - let chunk = chunk_result - .map_err(|e| Error::internal_err(format!("Failed to read chunk: {}", e)))?; + while let Some(event) = stream.next().await { + match event { + Ok(event) => { + if *DEBUG_SSE_STREAM { + tracing::info!("SSE event: {:?}", event); + } - // Convert chunk to string and add to buffer - let chunk_str = String::from_utf8_lossy(&chunk); - if *DEBUG_SSE_STREAM { - tracing::info!("SSE chunk: {}", chunk_str); - } - buffer.push_str(&chunk_str); - - // Process complete lines from buffer - while let Some(newline_pos) = buffer.find("\n\n") { - let line = buffer.drain(..newline_pos + 2).collect::(); - let line = line.trim_end_matches('\n'); - - // Skip empty lines and comments - if line.is_empty() || line.starts_with(':') { - continue; - } - - // Parse SSE data field - if let Some(data) = line.strip_prefix("data: ") { - if data == "[DONE]" { - // OpenAI sends [DONE] to indicate end of stream + if event.data == "[DONE]" { return Ok(()); } - self.parse_event_data(data).await?; + self.parse_event_data(&event.data).await?; + } + Err(e) => { + tracing::error!("Failed to parse SSE event: {}", e); } } }