From f990107c45fbb2e955ef67439e92328976091eb0 Mon Sep 17 00:00:00 2001 From: hugocasa Date: Fri, 26 Sep 2025 19:09:34 +0200 Subject: [PATCH] feat: ai agent streaming (#6644) * feat: ai agent step streaming * refactor * all * nits * fix other providers * nits * adapt to new streaming process --- backend/windmill-api/src/jobs.rs | 50 ++++-- backend/windmill-worker/src/ai/mod.rs | 1 + .../src/ai/providers/google_ai.rs | 21 ++- .../src/ai/providers/openai.rs | 55 +++++- .../src/ai/providers/openrouter.rs | 21 ++- .../windmill-worker/src/ai/query_builder.rs | 111 +++++++++++- backend/windmill-worker/src/ai/sse.rs | 158 ++++++++++++++++++ backend/windmill-worker/src/ai/types.rs | 19 +++ backend/windmill-worker/src/ai_executor.rs | 148 ++++++++++++++-- backend/windmill-worker/src/bun_executor.rs | 2 +- backend/windmill-worker/src/deno_executor.rs | 2 +- backend/windmill-worker/src/js_eval.rs | 9 +- backend/windmill-worker/src/worker.rs | 1 + .../lib/components/AIAgentLogViewer.svelte | 4 +- .../src/lib/components/flows/flowInfers.ts | 9 +- 15 files changed, 555 insertions(+), 56 deletions(-) create mode 100644 backend/windmill-worker/src/ai/sse.rs diff --git a/backend/windmill-api/src/jobs.rs b/backend/windmill-api/src/jobs.rs index e7ed78b257..4e8beecfa5 100644 --- a/backend/windmill-api/src/jobs.rs +++ b/backend/windmill-api/src/jobs.rs @@ -175,6 +175,7 @@ pub fn workspaced_service() -> Router { .route( "/run_and_stream/f/*script_path", get(stream_flow_by_path) + .post(stream_flow_by_path) .head(|| async { "" }) .layer(cors.clone()) .layer(ce_headers.clone()), @@ -182,6 +183,7 @@ pub fn workspaced_service() -> Router { .route( "/run_and_stream/p/*script_path", get(stream_script_by_path) + .post(stream_script_by_path) .head(|| async { "" }) .layer(cors.clone()) .layer(ce_headers.clone()), @@ -189,6 +191,7 @@ pub fn workspaced_service() -> Router { .route( "/run_and_stream/h/:hash", get(stream_script_by_hash) + .post(stream_script_by_hash) .head(|| async { "" }) .layer(cors.clone()) .layer(ce_headers.clone()), @@ -5220,6 +5223,7 @@ pub async fn stream_flow_by_path( Extension(user_db): Extension, Path((w_id, flow_path)): Path<(String, StripPath)>, Query(run_query): Query, + method: hyper::http::Method, args: RawWebhookArgs, ) -> error::Result { stream_job( @@ -5230,6 +5234,7 @@ pub async fn stream_flow_by_path( RunnableId::from_flow_path(flow_path.to_path()), args, run_query, + method == http::Method::GET, ) .await } @@ -5240,6 +5245,7 @@ pub async fn stream_script_by_path( Extension(user_db): Extension, Path((w_id, script_path)): Path<(String, StripPath)>, Query(run_query): Query, + method: hyper::http::Method, args: RawWebhookArgs, ) -> error::Result { stream_job( @@ -5250,6 +5256,7 @@ pub async fn stream_script_by_path( RunnableId::from_script_path(script_path.to_path()), args, run_query, + method == http::Method::GET, ) .await } @@ -5260,6 +5267,7 @@ pub async fn stream_script_by_hash( Extension(user_db): Extension, Path((w_id, script_hash)): Path<(String, ScriptHash)>, Query(run_query): Query, + method: hyper::http::Method, args: RawWebhookArgs, ) -> error::Result { stream_job( @@ -5270,6 +5278,7 @@ pub async fn stream_script_by_hash( RunnableId::from_script_hash(script_hash), args, run_query, + method == http::Method::GET, ) .await } @@ -5282,24 +5291,39 @@ pub async fn stream_job( runnable_id: RunnableId, args: RawWebhookArgs, run_query: RunJobQuery, + is_get: bool, ) -> error::Result { - let payload_r = run_query.payload.clone().map(decode_payload).map(|x| { - x.map_err(|e| Error::internal_err(format!("Impossible to decode query payload: {e:#?}"))) - }); + let args = if is_get { + let payload_r = run_query.payload.clone().map(decode_payload).map(|x| { + x.map_err(|e| { + Error::internal_err(format!("Impossible to decode query payload: {e:#?}")) + }) + }); - let payload_args = if let Some(payload) = payload_r { - payload? + let payload_args = if let Some(payload) = payload_r { + payload? + } else { + HashMap::new() + }; + + let mut args = args.process_args(&authed, &db, &w_id, None).await?; + args.body = args::Body::HashMap(payload_args); + + let args = args + .to_args_from_runnable(&db, &w_id, runnable_id.clone(), run_query.skip_preprocessor) + .await?; + args } else { - HashMap::new() + args.to_args_from_runnable( + &authed, + &db, + &w_id, + runnable_id.clone(), + run_query.skip_preprocessor, + ) + .await? }; - let mut args = args.process_args(&authed, &db, &w_id, None).await?; - args.body = args::Body::HashMap(payload_args); - - let args = args - .to_args_from_runnable(&db, &w_id, runnable_id.clone(), run_query.skip_preprocessor) - .await?; - let poll_delay_ms = run_query.poll_delay_ms; let uuid = match runnable_id { RunnableId::ScriptId(ScriptId::ScriptPath(script_path)) diff --git a/backend/windmill-worker/src/ai/mod.rs b/backend/windmill-worker/src/ai/mod.rs index b381e51693..4d11f100a5 100644 --- a/backend/windmill-worker/src/ai/mod.rs +++ b/backend/windmill-worker/src/ai/mod.rs @@ -4,4 +4,5 @@ pub mod image_handler; pub mod providers; pub mod query_builder; +pub mod sse; pub mod types; diff --git a/backend/windmill-worker/src/ai/providers/google_ai.rs b/backend/windmill-worker/src/ai/providers/google_ai.rs index b5a27a5f99..b94d4f29e5 100644 --- a/backend/windmill-worker/src/ai/providers/google_ai.rs +++ b/backend/windmill-worker/src/ai/providers/google_ai.rs @@ -5,7 +5,7 @@ use windmill_common::{ai_providers::AIProvider, client::AuthedClient, error::Err use crate::ai::{ image_handler::download_and_encode_s3_image, - query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder}, + query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventProcessor}, types::*, }; @@ -157,18 +157,24 @@ impl QueryBuilder for GoogleAIQueryBuilder { matches!(output_type, OutputType::Text) } + fn supports_streaming(&self) -> bool { + // Google AI supports streaming for text output + true + } + async fn build_request( &self, args: &BuildRequestArgs<'_>, client: &AuthedClient, workspace_id: &str, + stream: bool, ) -> Result { match args.output_type { OutputType::Text => { // For text output, use OpenAI-compatible format let openai_builder = super::openai::OpenAIQueryBuilder::new(AIProvider::GoogleAI); openai_builder - .build_request(args, client, workspace_id) + .build_request(args, client, workspace_id, stream) .await } OutputType::Image => self.build_image_request(args, client, workspace_id).await, @@ -238,6 +244,17 @@ impl QueryBuilder for GoogleAIQueryBuilder { } } + async fn parse_streaming_response( + &self, + response: reqwest::Response, + stream_event_processor: StreamEventProcessor, + ) -> Result { + let openai_builder = super::openai::OpenAIQueryBuilder::new(AIProvider::GoogleAI); + openai_builder + .parse_streaming_response(response, stream_event_processor) + .await + } + fn get_endpoint(&self, base_url: &str, model: &str, output_type: &OutputType) -> String { match output_type { OutputType::Text => format!("{}/chat/completions", base_url), // Use OpenAI-compatible endpoint diff --git a/backend/windmill-worker/src/ai/providers/openai.rs b/backend/windmill-worker/src/ai/providers/openai.rs index 421c5761f6..4fc4065766 100644 --- a/backend/windmill-worker/src/ai/providers/openai.rs +++ b/backend/windmill-worker/src/ai/providers/openai.rs @@ -5,7 +5,8 @@ use windmill_common::{ai_providers::AIProvider, client::AuthedClient, error::Err use crate::ai::{ image_handler::download_and_encode_s3_image, - query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder}, + query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventProcessor}, + sse::{OpenAISSEParser, SSEParser}, types::*, }; @@ -90,6 +91,7 @@ pub struct OpenAIRequest<'a> { pub max_completion_tokens: Option, #[serde(skip_serializing_if = "Option::is_none")] pub response_format: Option, + pub stream: bool, } pub struct OpenAIQueryBuilder { @@ -162,6 +164,7 @@ impl OpenAIQueryBuilder { args: &BuildRequestArgs<'_>, client: &AuthedClient, workspace_id: &str, + stream: bool, ) -> Result { let prepared_messages = self .prepare_messages_for_api(args.messages, client, workspace_id) @@ -196,6 +199,7 @@ impl OpenAIQueryBuilder { temperature: args.temperature, max_completion_tokens: args.max_tokens, response_format, + stream, }; serde_json::to_string(&request) @@ -255,14 +259,23 @@ impl QueryBuilder for OpenAIQueryBuilder { true } + fn supports_streaming(&self) -> bool { + // OpenAI supports streaming for text output + true + } + async fn build_request( &self, args: &BuildRequestArgs<'_>, client: &AuthedClient, workspace_id: &str, + stream: bool, ) -> Result { match args.output_type { - OutputType::Text => self.build_text_request(args, client, workspace_id).await, + OutputType::Text => { + self.build_text_request(args, client, workspace_id, stream) + .await + } OutputType::Image => self.build_image_request(args, client, workspace_id).await, } } @@ -330,10 +343,48 @@ impl QueryBuilder for OpenAIQueryBuilder { } }), tool_calls: first_choice.message.tool_calls.unwrap_or_default(), + events_str: None, }) } } + async fn parse_streaming_response( + &self, + response: reqwest::Response, + stream_event_processor: StreamEventProcessor, + ) -> Result { + let mut openai_sse_parser = OpenAISSEParser::new(stream_event_processor); + openai_sse_parser.parse_events(response).await?; + + let OpenAISSEParser { + accumulated_content, + accumulated_tool_calls, + mut events_str, + stream_event_processor, + } = openai_sse_parser; + + // Process streaming events with error handling + + for tool_call in accumulated_tool_calls.values() { + let event = StreamingEvent::ToolCallArguments { + call_id: tool_call.id.clone(), + function_name: tool_call.function.name.clone(), + arguments: tool_call.function.arguments.clone(), + }; + stream_event_processor.send(event, &mut events_str).await?; + } + + Ok(ParsedResponse::Text { + content: if accumulated_content.is_empty() { + None + } else { + Some(accumulated_content) + }, + tool_calls: accumulated_tool_calls.into_values().collect(), + events_str: Some(events_str), + }) + } + fn get_endpoint(&self, base_url: &str, model: &str, output_type: &OutputType) -> String { let path = match output_type { OutputType::Text => "chat/completions", diff --git a/backend/windmill-worker/src/ai/providers/openrouter.rs b/backend/windmill-worker/src/ai/providers/openrouter.rs index 3d5dd0314a..9e22a63552 100644 --- a/backend/windmill-worker/src/ai/providers/openrouter.rs +++ b/backend/windmill-worker/src/ai/providers/openrouter.rs @@ -5,7 +5,7 @@ use windmill_common::{ai_providers::AIProvider, client::AuthedClient, error::Err use crate::ai::{ providers::openai::{OpenAIQueryBuilder, OpenAIResponse}, - query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder}, + query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventProcessor}, types::*, }; @@ -70,17 +70,23 @@ impl QueryBuilder for OpenRouterQueryBuilder { true } + fn supports_streaming(&self) -> bool { + // OpenRouter supports streaming for text output + true + } + async fn build_request( &self, args: &BuildRequestArgs<'_>, client: &AuthedClient, workspace_id: &str, + stream: bool, ) -> Result { match args.output_type { OutputType::Text => { // For text, use standard OpenAI format without modalities self.openai_builder - .build_request(args, client, workspace_id) + .build_request(args, client, workspace_id, stream) .await } OutputType::Image => { @@ -184,9 +190,20 @@ impl QueryBuilder for OpenRouterQueryBuilder { .join(" "), }), tool_calls: first_choice.message.tool_calls.unwrap_or_default(), + events_str: None, }) } + async fn parse_streaming_response( + &self, + response: reqwest::Response, + stream_event_processor: StreamEventProcessor, + ) -> Result { + self.openai_builder + .parse_streaming_response(response, stream_event_processor) + .await + } + fn get_endpoint(&self, base_url: &str, _model: &str, _output_type: &OutputType) -> String { // OpenRouter uses the same endpoint for both text and image generation format!("{}/chat/completions", base_url) diff --git a/backend/windmill-worker/src/ai/query_builder.rs b/backend/windmill-worker/src/ai/query_builder.rs index 6b043726bc..e332c21f73 100644 --- a/backend/windmill-worker/src/ai/query_builder.rs +++ b/backend/windmill-worker/src/ai/query_builder.rs @@ -1,13 +1,19 @@ use async_trait::async_trait; -use windmill_common::{client::AuthedClient, error::Error, s3_helpers::S3Object}; +use windmill_common::{ + client::AuthedClient, error::Error, s3_helpers::S3Object, worker::Connection, +}; +use windmill_queue::MiniPulledJob; -use crate::ai::{ - providers::{ - google_ai::GoogleAIQueryBuilder, - openai::{OpenAIQueryBuilder, OpenAIToolCall}, - openrouter::OpenRouterQueryBuilder, +use crate::{ + ai::{ + providers::{ + google_ai::GoogleAIQueryBuilder, + openai::{OpenAIQueryBuilder, OpenAIToolCall}, + openrouter::OpenRouterQueryBuilder, + }, + types::*, }, - types::*, + job_logger::append_result_stream, }; /// Arguments for building an AI request @@ -26,7 +32,7 @@ pub struct BuildRequestArgs<'a> { /// Response from AI provider pub enum ParsedResponse { - Text { content: Option, tool_calls: Vec }, + Text { content: Option, tool_calls: Vec, events_str: Option }, Image { base64_data: String }, } @@ -36,17 +42,32 @@ pub trait QueryBuilder: Send + Sync { /// Check if this provider supports tools with the given output type fn supports_tools_with_output_type(&self, output_type: &OutputType) -> bool; + /// Check if this provider supports streaming + fn supports_streaming(&self) -> bool; + /// Build the request body for the provider async fn build_request( &self, args: &BuildRequestArgs<'_>, client: &AuthedClient, workspace_id: &str, + stream: bool, ) -> Result; /// Parse the response from the provider async fn parse_response(&self, response: reqwest::Response) -> Result; + /// Parse streaming response from the provider + async fn parse_streaming_response( + &self, + _response: reqwest::Response, + _stream_event_processor: StreamEventProcessor, + ) -> Result { + return Err(Error::internal_err( + "Missing implementation for parse_streaming_response for this provider".to_string(), + )); + } + /// Get the API endpoint for this provider fn get_endpoint(&self, base_url: &str, model: &str, output_type: &OutputType) -> String; @@ -69,3 +90,77 @@ pub fn create_query_builder(provider: &ProviderWithResource) -> Box Box::new(OpenAIQueryBuilder::new(provider.kind.clone())), // Pass provider kind for Azure handling } } + +pub struct StreamEventProcessor { + tx: tokio::sync::mpsc::Sender, + pub handle: Option>, +} + +impl Clone for StreamEventProcessor { + fn clone(&self) -> Self { + Self { tx: self.tx.clone(), handle: None } + } +} + +impl StreamEventProcessor { + pub fn new(conn: &Connection, job: &MiniPulledJob) -> Self { + let (tx, mut rx) = tokio::sync::mpsc::channel::(100); + let conn = conn.clone(); + let job_id = job.id.clone(); + let workspace_id = job.workspace_id.clone(); + let handle = tokio::spawn(async move { + let mut offset = -1; + while let Some(event) = rx.recv().await { + offset += 1; + match tokio::time::timeout( + std::time::Duration::from_secs(20), + append_result_stream(&conn, &workspace_id, &job_id, &event, offset), + ) + .await + { + Ok(res) => { + if let Err(err) = res { + tracing::error!("Failed to save stream event: {}", err); + } + } + Err(err) => { + tracing::error!("Did not manage to save stream event after 20 seconds, stopping stream event processor: {}", err); + break; + } + } + } + }); + + Self { tx, handle: Some(handle) } + } + + pub async fn send(&self, event: StreamingEvent, events_str: &mut String) -> Result<(), Error> { + match serde_json::to_string(&event) { + Ok(event_json) => { + let event_json = format!("{}\n", event_json); + events_str.push_str(&event_json); + if let Err(err) = self + .tx + .send(event_json.clone()) + .await + .map_err(|e| Error::internal_err(format!("Failed to send event: {}", e))) + { + tracing::error!( + "Failed to send event to stream event processor, skiping event: {}", + err + ); + } + + Ok(()) + } + Err(e) => Err(Error::internal_err(format!( + "Failed to serialize streaming event {:#?}, error is: {}", + event, e + ))), + } + } + + pub fn to_handle(self) -> Option> { + self.handle + } +} diff --git a/backend/windmill-worker/src/ai/sse.rs b/backend/windmill-worker/src/ai/sse.rs new file mode 100644 index 0000000000..fcafc9cd7e --- /dev/null +++ b/backend/windmill-worker/src/ai/sse.rs @@ -0,0 +1,158 @@ +use std::collections::HashMap; + +use reqwest::Response; +use serde::Deserialize; +use serde_json; +use tokio_stream::StreamExt; +use windmill_common::{error::Error, utils::rd_string}; + +use crate::ai::{ + providers::openai::{OpenAIFunction, OpenAIToolCall}, + query_builder::StreamEventProcessor, + types::StreamingEvent, +}; + +#[derive(Deserialize)] +pub struct OpenAIChoiceDeltaToolCallFunction { + pub name: Option, + pub arguments: Option, +} + +#[derive(Deserialize)] +pub struct OpenAIChoiceDeltaToolCall { + pub index: Option, + pub id: Option, + pub function: Option, +} + +#[derive(Deserialize)] +pub struct OpenAIChoiceDelta { + pub content: Option, + pub tool_calls: Option>, +} + +#[derive(Deserialize)] +pub struct OpenAIChoice { + pub delta: Option, +} + +#[derive(Deserialize)] +pub struct OpenAISSEEvent { + pub choices: Option>, +} + +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(); + + while let Some(chunk_result) = stream.next().await { + let chunk = chunk_result + .map_err(|e| Error::internal_err(format!("Failed to read chunk: {}", e)))?; + + // Convert chunk to string and add to buffer + let chunk_str = String::from_utf8_lossy(&chunk); + 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 + return Ok(()); + } + + self.parse_event_data(data).await?; + } + } + } + + Ok(()) + } +} + +pub struct OpenAISSEParser { + pub accumulated_content: String, + pub accumulated_tool_calls: HashMap, + pub events_str: String, + pub stream_event_processor: StreamEventProcessor, +} + +impl OpenAISSEParser { + pub fn new(stream_event_processor: StreamEventProcessor) -> Self { + Self { + accumulated_content: String::new(), + accumulated_tool_calls: HashMap::new(), + events_str: String::new(), + stream_event_processor, + } + } +} + +impl SSEParser for OpenAISSEParser { + async fn parse_event_data(&mut self, data: &str) -> Result<(), Error> { + let event: OpenAISSEEvent = serde_json::from_str(data).map_err(|e| { + Error::internal_err(format!("Failed to parse SSE chunk {}: {}", data, e)) + })?; + + if let Some(mut choices) = event.choices.filter(|s| !s.is_empty()) { + if let Some(delta) = choices.remove(0).delta { + if let Some(content) = delta.content.filter(|s| !s.is_empty()) { + self.accumulated_content.push_str(&content); + let event = StreamingEvent::TokenDelta { content }; + self.stream_event_processor + .send(event, &mut self.events_str) + .await?; + } + + if let Some(tool_calls) = delta.tool_calls { + for (idx, tool_call) in tool_calls.into_iter().enumerate() { + let idx = tool_call.index.unwrap_or_else(|| idx as i64); + + if let Some(function) = tool_call.function { + if let Some(tool_call) = self.accumulated_tool_calls.get_mut(&idx) { + if let Some(arguments) = function.arguments { + tool_call.function.arguments += &arguments; + } + } else { + let fun_name = function.name.unwrap_or_default(); + let call_id = tool_call.id.unwrap_or_else(|| rd_string(24)); + let event = StreamingEvent::ToolCall { + call_id: call_id.clone(), + function_name: fun_name.clone(), + }; + self.stream_event_processor + .send(event, &mut self.events_str) + .await?; + self.accumulated_tool_calls.insert( + idx, + OpenAIToolCall { + id: call_id, + function: OpenAIFunction { + name: fun_name, + arguments: function.arguments.unwrap_or_default(), + }, + r#type: "function".to_string(), + }, + ); + } + } + } + } + } + } + + Ok(()) + } +} diff --git a/backend/windmill-worker/src/ai/types.rs b/backend/windmill-worker/src/ai/types.rs index cc1417a92d..2ffab9713e 100644 --- a/backend/windmill-worker/src/ai/types.rs +++ b/backend/windmill-worker/src/ai/types.rs @@ -116,6 +116,7 @@ pub struct AIAgentArgs { pub output_schema: Option, pub output_type: Option, pub user_images: Option>, + pub streaming: Option, } #[derive(Deserialize, Debug)] @@ -153,6 +154,24 @@ impl ProviderWithResource { pub struct AIAgentResult<'a> { pub output: Box, pub messages: Vec>, + #[serde(skip_serializing_if = "Option::is_none")] + pub wm_stream: Option, +} + +/// Events for streaming AI responses +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum StreamingEvent { + /// Individual token from the AI response + TokenDelta { content: String }, + /// Tool call has started + ToolCall { call_id: String, function_name: String }, + /// Tool call arguments are complete + ToolCallArguments { call_id: String, function_name: String, arguments: String }, + /// Tool execution has started + ToolExecution { call_id: String, function_name: String }, + /// Tool execution result + ToolResult { call_id: String, function_name: String, result: String, success: bool }, } #[derive(Serialize, Deserialize, Clone, Debug)] diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs index 2eddb657b5..0e14b4f556 100644 --- a/backend/windmill-worker/src/ai_executor.rs +++ b/backend/windmill-worker/src/ai_executor.rs @@ -1,3 +1,4 @@ +use anyhow::Context; use async_recursion::async_recursion; use regex::Regex; use serde_json::value::RawValue; @@ -5,8 +6,7 @@ use std::{collections::HashMap, sync::Arc}; use ulid; use uuid::Uuid; use windmill_common::{ - ai_providers::AIProvider, - ai_providers::AZURE_API_VERSION, + ai_providers::{AIProvider, AZURE_API_VERSION}, cache, client::AuthedClient, db::DB, @@ -27,10 +27,14 @@ use windmill_queue::{ use crate::{ ai::{ image_handler::upload_image_to_s3, - query_builder::{create_query_builder, BuildRequestArgs, ParsedResponse}, + query_builder::{ + create_query_builder, BuildRequestArgs, ParsedResponse, StreamEventProcessor, + }, types::*, }, - common::{build_args_map, error_to_value, OccupancyMetrics}, + common::{ + build_args_map, error_to_value, resolve_job_timeout, OccupancyMetrics, StreamNotifier, + }, create_job_dir, handle_child::run_future_with_polling_update_job_poller, handle_queued_job, parse_sig_of_lang, @@ -44,7 +48,6 @@ lazy_static::lazy_static! { } const MAX_AGENT_ITERATIONS: usize = 10; -const REQUEST_TIMEOUT_SECONDS: u64 = 120; fn parse_raw_script_schema(content: &str, language: &ScriptLang) -> Result, Error> { let main_arg_signature = parse_sig_of_lang(content, Some(&language), None)?.unwrap(); // safe to unwrap as langauge is some @@ -114,6 +117,7 @@ pub async fn handle_ai_agent_job( worker_name: &str, hostname: &str, killpill_rx: &mut tokio::sync::broadcast::Receiver<()>, + has_stream: &mut bool, ) -> Result, Error> { let args = build_args_map(job, client, conn).await?; @@ -264,6 +268,12 @@ pub async fn handle_ai_agent_job( let mut inner_occupancy_metrics = occupancy_metrics.clone(); + let stream_notifier = StreamNotifier::new(conn, job); + + if let Some(stream_notifier) = stream_notifier { + stream_notifier.update_flow_status_with_stream_job(); + } + let agent_fut = run_agent( db, conn, @@ -279,6 +289,7 @@ pub async fn handle_ai_agent_job( worker_name, hostname, killpill_rx, + has_stream, ); let result = run_future_with_polling_update_job_poller( @@ -412,6 +423,7 @@ pub async fn run_agent( worker_name: &str, hostname: &str, killpill_rx: &mut tokio::sync::broadcast::Receiver<()>, + has_stream: &mut bool, ) -> error::Result> { let output_type = args.output_type.as_ref().unwrap_or(&OutputType::Text); let base_url = args.provider.get_base_url(db).await?; @@ -501,6 +513,21 @@ pub async fn run_agent( // For non-Anthropic providers, response_format is handled by the query builder } + // Check if streaming is enabled and supported + let should_stream = args.streaming.unwrap_or(false) + && query_builder.supports_streaming() + && output_type == &OutputType::Text; + + *has_stream = should_stream; + + let mut final_events_str = String::new(); + + let stream_event_processor = if should_stream { + Some(StreamEventProcessor::new(conn, job)) + } else { + None + }; + // Main agent loop for i in 0..MAX_AGENT_ITERATIONS { if used_structured_output_tool { @@ -522,21 +549,25 @@ pub async fn run_agent( }; let request_body = query_builder - .build_request(&build_args, client, &job.workspace_id) + .build_request(&build_args, client, &job.workspace_id, should_stream) .await?; let endpoint = query_builder.get_endpoint(&base_url, args.provider.get_model(), output_type); let auth_headers = query_builder.get_auth_headers(api_key, &base_url, output_type); + let timeout = resolve_job_timeout(conn, &job.workspace_id, job.id, job.timeout) + .await + .0; + let mut request = HTTP_CLIENT .post(&endpoint) - .timeout(std::time::Duration::from_secs(REQUEST_TIMEOUT_SECONDS)) + .timeout(timeout) .header("Content-Type", "application/json"); // Apply authentication headers - for (header_name, header_value) in auth_headers { - request = request.header(header_name, header_value); + for (header_name, header_value) in &auth_headers { + request = request.header(*header_name, header_value.clone()); } if args.provider.kind.is_azure_openai(&base_url) { @@ -551,10 +582,21 @@ pub async fn run_agent( match resp.error_for_status_ref() { Ok(_) => { - let parsed = query_builder.parse_response(resp).await?; + let parsed = if let Some(stream_event_processor) = stream_event_processor.clone() { + query_builder + .parse_streaming_response(resp, stream_event_processor) + .await? + } else { + // Handle non-streaming response + query_builder.parse_response(resp).await? + }; match parsed { - ParsedResponse::Text { content: response_content, tool_calls } => { + ParsedResponse::Text { content: response_content, tool_calls, events_str } => { + if let Some(events_str) = events_str { + final_events_str.push_str(&events_str); + } + if let Some(ref response_content) = response_content { actions.push(AgentAction::Message {}); messages.push(OpenAIMessage { @@ -589,6 +631,17 @@ pub async fn run_agent( // Handle tool calls (keeping existing tool execution logic) for tool_call in tool_calls.iter() { + // Stream tool call progress + if let Some(ref stream_event_processor) = stream_event_processor { + let event = StreamingEvent::ToolExecution { + call_id: tool_call.id.clone(), + function_name: tool_call.function.name.clone(), + }; + stream_event_processor + .send(event, &mut final_events_str) + .await?; + } + // Check if this is the structured output tool if structured_output_tool_name .as_ref() @@ -631,10 +684,22 @@ pub async fn run_agent( update_flow_status_module_with_actions(db, parent_job, &actions) .await?; + let raw_tool_call_args = if tool_call.function.arguments.is_empty() + { + "{}".to_string() + } else { + tool_call.function.arguments.clone() + }; let tool_call_args = serde_json::from_str::>>( - &tool_call.function.arguments, - )?; + &raw_tool_call_args, + ) + .with_context(|| { + format!( + "Failed to parse tool call arguments for tool call {}: {}", + tool_call.function.name, tool_call.function.arguments + ) + })?; let job_payload = match tool.module.get_value()? { FlowModuleValue::Script { @@ -806,12 +871,13 @@ pub async fn run_agent( worker_name, ) .await; + let error_message = + format!("Error running tool: {}", err_string); messages.push(OpenAIMessage { role: "tool".to_string(), - content: Some(OpenAIContent::Text(format!( - "Error running tool: {}", - err_string - ))), + content: Some(OpenAIContent::Text( + error_message.clone(), + )), tool_call_id: Some(tool_call.id.clone()), agent_action: Some(AgentAction::ToolCall { job_id, @@ -820,6 +886,21 @@ pub async fn run_agent( }), ..Default::default() }); + // Stream tool result (error case) + if let Some(ref stream_event_processor) = + stream_event_processor + { + let tool_result_event = StreamingEvent::ToolResult { + call_id: tool_call.id.clone(), + function_name: tool_call.function.name.clone(), + result: error_message, + success: false, + }; + stream_event_processor + .send(tool_result_event, &mut final_events_str) + .await?; + } + update_flow_status_module_with_actions_success( db, parent_job, false, ) @@ -857,7 +938,6 @@ pub async fn run_agent( "Tool job completed but no result".to_string(), )); }; - messages.push(OpenAIMessage { role: "tool".to_string(), content: Some(OpenAIContent::Text( @@ -871,6 +951,22 @@ pub async fn run_agent( }), ..Default::default() }); + + // Stream tool result (success case) + if let Some(ref stream_event_processor) = + stream_event_processor + { + let tool_result_event = StreamingEvent::ToolResult { + call_id: tool_call.id.clone(), + function_name: tool_call.function.name.clone(), + result: result.get().to_string(), + success: true, + }; + stream_event_processor + .send(tool_result_event, &mut final_events_str) + .await?; + } + update_flow_status_module_with_actions_success( db, parent_job, success, ) @@ -930,8 +1026,24 @@ pub async fn run_agent( None => to_raw_value(&""), }; + if let Some(stream_event_processor) = stream_event_processor { + if let Some(handle) = stream_event_processor.to_handle() { + if let Err(e) = handle.await { + return Err(Error::internal_err(format!( + "Error waiting for stream event processor: {}", + e + ))); + } + } + } + Ok(to_raw_value(&AIAgentResult { output: output_value, messages: final_messages, + wm_stream: if !final_events_str.is_empty() { + Some(final_events_str) + } else { + None + }, })) } diff --git a/backend/windmill-worker/src/bun_executor.rs b/backend/windmill-worker/src/bun_executor.rs index 492b8f3beb..87f15c3294 100644 --- a/backend/windmill-worker/src/bun_executor.rs +++ b/backend/windmill-worker/src/bun_executor.rs @@ -1116,7 +1116,7 @@ async function run() {{ let res = await Main.{main_name}(...argsArr); if (isAsyncIterable(res)) {{ for await (const chunk of res) {{ - console.log("WM_STREAM: " + chunk.replace('\n', '\\n')); + console.log("WM_STREAM: " + chunk.replace(/\n/g, '\\n')); }} res = null; }} diff --git a/backend/windmill-worker/src/deno_executor.rs b/backend/windmill-worker/src/deno_executor.rs index 3433eef3ae..fffe236fe9 100644 --- a/backend/windmill-worker/src/deno_executor.rs +++ b/backend/windmill-worker/src/deno_executor.rs @@ -299,7 +299,7 @@ async function run() {{ let res: any = await {main_name}(...argsArr); if (isAsyncIterable(res)) {{ for await (const chunk of res) {{ - console.log("WM_STREAM: " + chunk.replace('\n', '\\n')); + console.log("WM_STREAM: " + chunk.replace(/\n/g, '\\n')); }} res = null; }} diff --git a/backend/windmill-worker/src/js_eval.rs b/backend/windmill-worker/src/js_eval.rs index b935b017f9..ec6195a5cc 100644 --- a/backend/windmill-worker/src/js_eval.rs +++ b/backend/windmill-worker/src/js_eval.rs @@ -810,14 +810,11 @@ pub async fn eval_fetch_timeout( let conn_ = conn.clone(); let w_id_ = w_id.to_string(); tokio::spawn(async move { - let mut offset = 0; + let mut offset = -1; while let Some(stream) = result_stream_receiver.recv().await { use crate::job_logger::append_result_stream; - let curr_offset = offset.clone(); offset += 1; - if let Err(e) = - append_result_stream(&conn_, &w_id_, &job_id, &stream, curr_offset).await - { + if let Err(e) = append_result_stream(&conn_, &w_id_, &job_id, &stream, offset).await { tracing::error!("failed to append result stream: {e}"); } } @@ -1119,7 +1116,7 @@ function processStreamIterative(res) {{ iterator.next().then(function(result) {{ if (!result.done) {{ const chunk = result.value; - console.log("WM_STREAM: " + chunk.replace('\n', '\\n')); + console.log("WM_STREAM: " + chunk.replace(/\n/g, '\\n')); // Continue the loop step(); }} else {{ diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs index 6a70d250df..6c9f278d59 100644 --- a/backend/windmill-worker/src/worker.rs +++ b/backend/windmill-worker/src/worker.rs @@ -2588,6 +2588,7 @@ pub async fn handle_queued_job( worker_name, hostname, killpill_rx, + &mut has_stream, ) .await } diff --git a/frontend/src/lib/components/AIAgentLogViewer.svelte b/frontend/src/lib/components/AIAgentLogViewer.svelte index 2640ea52eb..c8fd14d6fc 100644 --- a/frontend/src/lib/components/AIAgentLogViewer.svelte +++ b/frontend/src/lib/components/AIAgentLogViewer.svelte @@ -13,14 +13,14 @@ import { onMount } from 'svelte' type AgentActionWithContent = NonNullable[number] & { - content: string + content?: unknown } const resultSchema = z.object({ messages: z.array( z.object({ role: z.string(), - content: z.string().optional(), + content: z.unknown(), agent_action: z .union([ z.object({ diff --git a/frontend/src/lib/components/flows/flowInfers.ts b/frontend/src/lib/components/flows/flowInfers.ts index 7c24278cc2..ee25a6338e 100644 --- a/frontend/src/lib/components/flows/flowInfers.ts +++ b/frontend/src/lib/components/flows/flowInfers.ts @@ -77,6 +77,13 @@ export async function loadSchemaFromModule(module: FlowModule): Promise<{ type: 'string', description: 'The system prompt to give as input to the AI agent.' }, + streaming: { + type: 'boolean', + description: + 'Whether to stream the output of the AI agent (only used if output_type is text).', + default: false, + showExpr: "fields.output_type === 'text'" + }, user_images: { type: 'array', description: @@ -99,7 +106,7 @@ export async function loadSchemaFromModule(module: FlowModule): Promise<{ output_schema: { type: 'object', description: - 'JSON schema that the AI agent will follow for its response format (only used if output_type is text)', + 'JSON schema that the AI agent will follow for its response format (only used if output_type is text).', format: 'json-schema', showExpr: "fields.output_type === 'text'" }