starred/windmill
1
0
Fork 0
mirror of https://github.com/windmill-labs/windmill.git synced 2026-09-05 08:02:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: ai agent streaming (#6644)

Browse Source 
* feat: ai agent step streaming

* refactor

* all

* nits

* fix other providers

* nits

* adapt to new streaming process
This commit is contained in:
hugocasa
2025-09-26 17:09:34 +00:00
committed by GitHub
parent bce25bc283
commit f990107c45
15 changed files with 555 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/windmill-api/src/jobs.rs
+37 -13
View File
@@ -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<UserDB>,
Path((w_id, flow_path)): Path<(String, StripPath)>,
Query(run_query): Query<RunJobQuery>,
method: hyper::http::Method,
args: RawWebhookArgs,
) -> error::Result<Response> {
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<UserDB>,
Path((w_id, script_path)): Path<(String, StripPath)>,
Query(run_query): Query<RunJobQuery>,
method: hyper::http::Method,
args: RawWebhookArgs,
) -> error::Result<Response> {
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<UserDB>,
Path((w_id, script_hash)): Path<(String, ScriptHash)>,
Query(run_query): Query<RunJobQuery>,
method: hyper::http::Method,
args: RawWebhookArgs,
) -> error::Result<Response> {
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<Response> {
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))
backend/windmill-worker/src/ai/mod.rs
+1
View File
@@ -4,4 +4,5 @@
pub mod image_handler;
pub mod providers;
pub mod query_builder;
pub mod sse;
pub mod types;
backend/windmill-worker/src/ai/providers/google_ai.rs
+19 -2
View File
@@ -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<String, Error> {
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<ParsedResponse, Error> {
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
backend/windmill-worker/src/ai/providers/openai.rs
+53 -2
View File
@@ -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<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub response_format: Option<ResponseFormat>,
pub stream: bool,
}
pub struct OpenAIQueryBuilder {
@@ -162,6 +164,7 @@ impl OpenAIQueryBuilder {
args: &BuildRequestArgs<'_>,
client: &AuthedClient,
workspace_id: &str,
stream: bool,
) -> Result<String, Error> {
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<String, Error> {
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<ParsedResponse, Error> {
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",
backend/windmill-worker/src/ai/providers/openrouter.rs
+19 -2
View File
@@ -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<String, Error> {
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<ParsedResponse, Error> {
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)
backend/windmill-worker/src/ai/query_builder.rs
+103 -8
View File
@@ -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<String>, tool_calls: Vec<OpenAIToolCall> },
Text { content: Option<String>, tool_calls: Vec<OpenAIToolCall>, events_str: Option<String> },
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<String, Error>;
/// Parse the response from the provider
async fn parse_response(&self, response: reqwest::Response) -> Result<ParsedResponse, Error>;
/// Parse streaming response from the provider
async fn parse_streaming_response(
&self,
_response: reqwest::Response,
_stream_event_processor: StreamEventProcessor,
) -> Result<ParsedResponse, Error> {
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<dyn QueryBui
_ => Box::new(OpenAIQueryBuilder::new(provider.kind.clone())), // Pass provider kind for Azure handling
}
}
pub struct StreamEventProcessor {
tx: tokio::sync::mpsc::Sender<String>,
pub handle: Option<tokio::task::JoinHandle<()>>,
}
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::<String>(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<tokio::task::JoinHandle<()>> {
self.handle
}
}
backend/windmill-worker/src/ai/sse.rs
+158
View File
@@ -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<String>,
pub arguments: Option<String>,
}
#[derive(Deserialize)]
pub struct OpenAIChoiceDeltaToolCall {
pub index: Option<i64>,
pub id: Option<String>,
pub function: Option<OpenAIChoiceDeltaToolCallFunction>,
}
#[derive(Deserialize)]
pub struct OpenAIChoiceDelta {
pub content: Option<String>,
pub tool_calls: Option<Vec<OpenAIChoiceDeltaToolCall>>,
}
#[derive(Deserialize)]
pub struct OpenAIChoice {
pub delta: Option<OpenAIChoiceDelta>,
}
#[derive(Deserialize)]
pub struct OpenAISSEEvent {
pub choices: Option<Vec<OpenAIChoice>>,
}
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::<String>();
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<i64, OpenAIToolCall>,
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(())
}
}
backend/windmill-worker/src/ai/types.rs
+19
View File
@@ -116,6 +116,7 @@ pub struct AIAgentArgs {
pub output_schema: Option<OpenAPISchema>,
pub output_type: Option<OutputType>,
pub user_images: Option<Vec<S3Object>>,
pub streaming: Option<bool>,
}
#[derive(Deserialize, Debug)]
@@ -153,6 +154,24 @@ impl ProviderWithResource {
pub struct AIAgentResult<'a> {
pub output: Box<RawValue>,
pub messages: Vec<Message<'a>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub wm_stream: Option<String>,
}
/// 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)]
backend/windmill-worker/src/ai_executor.rs
+130 -18
View File
@@ -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<Box<RawValue>, 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<Box<RawValue>, 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<Box<RawValue>> {
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::<HashMap<String, Box<RawValue>>>(
&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
},
}))
}
backend/windmill-worker/src/bun_executor.rs
+1 -1
View File
@@ -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;
}}
backend/windmill-worker/src/deno_executor.rs
+1 -1
View File
@@ -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;
}}
backend/windmill-worker/src/js_eval.rs
+3 -6
View File
@@ -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 {{
backend/windmill-worker/src/worker.rs
+1
View File
@@ -2588,6 +2588,7 @@ pub async fn handle_queued_job(
worker_name,
hostname,
killpill_rx,
&mut has_stream,
)
.await
}
frontend/src/lib/components/AIAgentLogViewer.svelte
+2 -2
View File
@@ -13,14 +13,14 @@
import { onMount } from 'svelte'
type AgentActionWithContent = NonNullable<FlowStatusModule['agent_actions']>[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({
frontend/src/lib/components/flows/flowInfers.ts
+8 -1
View File
@@ -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'"
}