diff --git a/backend/windmill-api/src/jobs.rs b/backend/windmill-api/src/jobs.rs index 40cb11c8a5..52e3b23f06 100644 --- a/backend/windmill-api/src/jobs.rs +++ b/backend/windmill-api/src/jobs.rs @@ -3605,7 +3605,7 @@ where } fn decode_payload(t: String) -> anyhow::Result { - let vec = base64::engine::general_purpose::URL_SAFE + let vec = base64::engine::general_purpose::STANDARD .decode(t) .context("invalid base64")?; serde_json::from_slice(vec.as_slice()).context("invalid json") diff --git a/backend/windmill-common/src/cache.rs b/backend/windmill-common/src/cache.rs index a6305b784f..f46259f457 100644 --- a/backend/windmill-common/src/cache.rs +++ b/backend/windmill-common/src/cache.rs @@ -282,12 +282,25 @@ pub mod future { pub struct FlowData { pub raw_flow: Box, pub flow: FlowValue, + pub summary: Option, } impl FlowData { pub fn from_raw(raw_flow: Box) -> error::Result { - let flow = serde_json::from_str(raw_flow.get())?; - Ok(Self { raw_flow, flow }) + let (flow, summary) = if let Ok(parsed) = + serde_json::from_str::(raw_flow.get()) + { + (parsed.value, parsed.summary) + } else { + // fallback to plain FlowValue + ( + serde_json::from_str::(raw_flow.get()).map_err(|e| { + error::Error::internal_err(format!("Failed to parse as FlowValue: {}", e)) + })?, + None, + ) + }; + Ok(Self { raw_flow, flow, summary }) } pub fn value(&self) -> &FlowValue { @@ -837,10 +850,12 @@ pub mod job { match (kind, hash.map(|ScriptHash(id)| id)) { (FlowDependencies, Some(id)) => flow::fetch_version(db, id).await, (FlowNode, Some(id)) => flow::fetch_flow(db, FlowNodeId(id)).await, - (Flow, Some(id)) | (SingleStepFlow, Some(id)) => match flow::fetch_version_lite(db, id).await { - Ok(raw_flow) => Ok(raw_flow), - Err(_) => flow::fetch_version(db, id).await, - }, + (Flow, Some(id)) | (SingleStepFlow, Some(id)) => { + match flow::fetch_version_lite(db, id).await { + Ok(raw_flow) => Ok(raw_flow), + Err(_) => flow::fetch_version(db, id).await, + } + } _ => Err(error::Error::internal_err(format!( "Isn't a flow job {:?}", kind diff --git a/backend/windmill-common/src/flows.rs b/backend/windmill-common/src/flows.rs index 810368f901..99093cc107 100644 --- a/backend/windmill-common/src/flows.rs +++ b/backend/windmill-common/src/flows.rs @@ -141,6 +141,13 @@ pub struct FlowValue { pub chat_input_enabled: Option, } +#[derive(Serialize, Deserialize)] +pub struct FlowNodeFlow { + pub value: FlowValue, + #[serde(skip_serializing_if = "Option::is_none")] + pub summary: Option, +} + impl FlowValue { pub fn get_flow_module_at_step(&self, step: Step) -> anyhow::Result<&FlowModule> { let flow_module = match step { @@ -726,6 +733,8 @@ pub enum FlowModuleValue { AIAgent { input_transforms: HashMap, tools: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + modules_node: Option, }, } @@ -863,6 +872,7 @@ impl<'de> Deserialize<'de> for FlowModuleValue { tools: untagged .tools .ok_or_else(|| serde::de::Error::missing_field("tools"))?, + modules_node: untagged.modules_node, }), other => Err(serde::de::Error::unknown_variant( other, @@ -1054,6 +1064,9 @@ pub async fn resolve_module( .await?; } } + AIAgent { tools, modules_node, .. } => { + resolve_modules(db, workspace_id, tools, modules_node.take(), with_code).await?; + } _ => {} } *value = to_raw_value(&val); diff --git a/backend/windmill-common/src/jobs.rs b/backend/windmill-common/src/jobs.rs index 4608fe72b3..b4d4f0acc7 100644 --- a/backend/windmill-common/src/jobs.rs +++ b/backend/windmill-common/src/jobs.rs @@ -425,6 +425,7 @@ pub enum JobPayload { Noop, AIAgent { path: String, + flow_node_id: Option, }, } diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index dddeb46851..326c5a964c 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -3737,7 +3737,8 @@ pub async fn push<'c, 'd>( if let Some(skip_handler) = skip_handler { let mut skip_input_transforms = HashMap::::new(); for (arg_name, arg_value) in skip_handler.args { - skip_input_transforms.insert(arg_name, InputTransform::Static { value: arg_value }); + skip_input_transforms + .insert(arg_name, InputTransform::Static { value: arg_value }); } modules.push(FlowModule { @@ -3872,7 +3873,7 @@ pub async fn push<'c, 'd>( // this is a new flow being pushed, flow_status is set to flow_value: let flow_status: FlowStatus = FlowStatus::new(&flow_value); ( - None, // No version needed - flow is stored in raw_flow like FlowPreview + None, // No version needed - flow is stored in raw_flow like FlowPreview Some(path), None, JobKind::SingleStepFlow, @@ -4071,8 +4072,8 @@ pub async fn push<'c, 'd>( None, None, ), - JobPayload::AIAgent { path } => ( - None, + JobPayload::AIAgent { path, flow_node_id } => ( + flow_node_id.map(|id| id.0), Some(path), None, JobKind::AIAgent, diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs index 01a4b36b71..8fda760a83 100644 --- a/backend/windmill-worker/src/ai_executor.rs +++ b/backend/windmill-worker/src/ai_executor.rs @@ -1,4 +1,7 @@ -use crate::memory_oss::{read_from_memory, write_to_memory}; +use crate::{ + memory_oss::{read_from_memory, write_to_memory}, + worker_flow::JobPayloadWithTag, +}; use anyhow::Context; use async_recursion::async_recursion; use regex::Regex; @@ -14,9 +17,9 @@ use windmill_common::{ error::{self, to_anyhow, Error}, flow_conversations::{add_message_to_conversation_tx, MessageType}, flow_status::AgentAction, - flows::{FlowModuleValue, FlowValue, Step}, + flows::{FlowModuleValue, FlowNodeId, Step}, get_latest_hash_for_path, - jobs::JobKind, + jobs::{JobKind, JobPayload}, scripts::{get_full_hub_script_by_path, ScriptHash, ScriptLang}, utils::{StripPath, HTTP_CLIENT}, worker::{to_raw_value, Connection}, @@ -206,36 +209,51 @@ pub async fn handle_ai_agent_job( )); }; - let flow_job = get_flow_job_runnable_and_raw_flow(db, &parent_job).await?; + let (tools, summary) = if let Some(ScriptHash(flow_node_id)) = job.runnable_id { + tracing::debug!( + "Fetching AI Agent flow data using flow node id {}", + flow_node_id + ); + let flow_data = cache::flow::fetch_flow(db, FlowNodeId(flow_node_id)).await?; - let flow_data = match flow_job.kind { - JobKind::Flow | JobKind::FlowNode => { - cache::job::fetch_flow(db, &flow_job.kind, flow_job.runnable_id).await? - } - JobKind::FlowPreview => { - cache::job::fetch_preview_flow(db, &parent_job, flow_job.raw_flow).await? - } - _ => { + let value = flow_data.value(); + + (value.modules.clone(), flow_data.summary.clone()) + } else { + tracing::debug!("Fetching flow data for parent job of AI Agent job"); + let flow_job = get_flow_job_runnable_and_raw_flow(db, &parent_job).await?; + + let flow_data = match flow_job.kind { + JobKind::Flow | JobKind::FlowNode => { + cache::job::fetch_flow(db, &flow_job.kind, flow_job.runnable_id).await? + } + JobKind::FlowPreview => { + cache::job::fetch_preview_flow(db, &parent_job, flow_job.raw_flow).await? + } + _ => { + return Err(Error::internal_err( + "expected parent flow, flow preview or flow node for ai agent job".to_string(), + )); + } + }; + + let value = flow_data.value(); + + let module = value.modules.iter().find(|m| m.id == *flow_step_id); + + let Some(module) = module else { return Err(Error::internal_err( - "expected parent flow, flow preview or flow node for ai agent job".to_string(), + "AI agent module not found in flow".to_string(), )); - } - }; + }; - let value = flow_data.value(); + let FlowModuleValue::AIAgent { tools, .. } = module.get_value()? else { + return Err(Error::internal_err( + "AI agent module is not an AI agent".to_string(), + )); + }; - let module = value.modules.iter().find(|m| m.id == *flow_step_id); - - let Some(module) = module else { - return Err(Error::internal_err( - "AI agent module not found in flow".to_string(), - )); - }; - - let FlowModuleValue::AIAgent { tools, .. } = module.get_value()? else { - return Err(Error::internal_err( - "AI agent module is not an AI agent".to_string(), - )); + (tools, module.summary.clone()) }; let tools = futures::future::try_join_all(tools.into_iter().map(|mut t| { @@ -303,6 +321,10 @@ pub async fn handle_ai_agent_job( Ok(FlowModuleValue::RawScript { content, language, .. }) => { Ok(Some(parse_raw_script_schema(&content, &language)?)) } + Ok(FlowModuleValue::FlowScript { id, language, .. }) => { + let script_data = cache::flow::fetch_script(conn, id.clone()).await?; + Ok(Some(parse_raw_script_schema(&script_data.code, &language)?)) + } Err(e) => { return Err(Error::internal_err(format!( "Invalid tool {}: {}", @@ -354,7 +376,7 @@ pub async fn handle_ai_agent_job( parent_job, &args, &tools, - value, + summary.as_deref(), client, &mut inner_occupancy_metrics, job_completed_tx, @@ -469,14 +491,12 @@ async fn update_flow_status_module_with_actions_success( } /// Get step name from the flow module (summary if exists, else id) -fn get_step_name_from_flow(flow_value: &FlowValue, flow_step_id: Option<&str>) -> Option { +fn get_step_name_from_flow(summary: Option<&str>, flow_step_id: Option<&str>) -> Option { let flow_step_id = flow_step_id?; - let module = flow_value.modules.iter().find(|m| m.id == flow_step_id)?; Some( - module - .summary - .clone() - .unwrap_or_else(|| format!("AI Agent Step {}", module.id)), + summary + .map(|s| s.to_string()) + .unwrap_or_else(|| format!("AI Agent Step {}", flow_step_id)), ) } @@ -499,7 +519,7 @@ pub async fn run_agent( parent_job: &Uuid, args: &AIAgentArgs, tools: &[Tool], - flow_value: &FlowValue, + summary: Option<&str>, // job execution context client: &AuthedClient, @@ -751,7 +771,7 @@ pub async fn run_agent( let db_clone = db.clone(); let message_content = response_content.clone(); let step_name = get_step_name_from_flow( - flow_value, + summary, job.flow_step_id.as_deref(), ); @@ -914,6 +934,43 @@ pub async fn run_agent( ); payload } + FlowModuleValue::FlowScript { + id, + language, + custom_concurrency_key, + concurrent_limit, + concurrency_time_window_s, + tag, + .. + } => { + let path = format!( + "{}/tools/{}", + job.runnable_path(), + tool.module.id + ); + + let payload = JobPayloadWithTag { + payload: JobPayload::FlowScript { + id, + language, + custom_concurrency_key: custom_concurrency_key + .clone(), + concurrent_limit, + concurrency_time_window_s, + cache_ttl: tool.module.cache_ttl.map(|x| x as i32), + dedicated_worker: None, + path, + }, + tag: tag.clone(), + delete_after_use: tool + .module + .delete_after_use + .unwrap_or(false), + timeout: None, + on_behalf_of: None, + }; + payload + } _ => { return Err(Error::internal_err(format!( "Unsupported tool: {}", @@ -1121,7 +1178,7 @@ pub async fn run_agent( let tool_job_id = job_id; let db_clone = db.clone(); let step_name = get_step_name_from_flow( - flow_value, + summary, job.flow_step_id.as_deref(), ); @@ -1228,7 +1285,7 @@ pub async fn run_agent( let db_clone = db.clone(); let tool_name = tool_call.function.name.clone(); let step_name = get_step_name_from_flow( - flow_value, + summary, job.flow_step_id.as_deref(), ); let content = if success { diff --git a/backend/windmill-worker/src/worker_flow.rs b/backend/windmill-worker/src/worker_flow.rs index ad5b40ee4a..7a7f6cb52a 100644 --- a/backend/windmill-worker/src/worker_flow.rs +++ b/backend/windmill-worker/src/worker_flow.rs @@ -3701,9 +3701,9 @@ async fn compute_next_flow_transform( NextStatus::NextStep, )) } - FlowModuleValue::AIAgent { .. } => { + FlowModuleValue::AIAgent { modules_node, .. } => { let path = get_path(flow_job, status, module); - let payload = JobPayload::AIAgent { path }; + let payload = JobPayload::AIAgent { path, flow_node_id: modules_node }; Ok(NextFlowTransform::Continue( ContinuePayload::SingleJob(JobPayloadWithTag { payload, diff --git a/backend/windmill-worker/src/worker_lockfiles.rs b/backend/windmill-worker/src/worker_lockfiles.rs index 59f0f87552..621012aa9a 100644 --- a/backend/windmill-worker/src/worker_lockfiles.rs +++ b/backend/windmill-worker/src/worker_lockfiles.rs @@ -16,7 +16,7 @@ use uuid::Uuid; use windmill_common::assets::{clear_asset_usage, insert_asset_usage, AssetUsageKind}; use windmill_common::error::Error; use windmill_common::error::Result; -use windmill_common::flows::{FlowModule, FlowModuleValue, FlowNodeId}; +use windmill_common::flows::{FlowModule, FlowModuleValue, FlowNodeFlow, FlowNodeId}; use windmill_common::get_latest_deployed_hash_for_path; use windmill_common::jobs::JobPayload; use windmill_common::scripts::{hash_script, NewScript, ScriptHash}; @@ -1635,8 +1635,10 @@ async fn insert_flow_modules<'c>( workspace_id: &str, failure_module: Option<&Box>, same_worker: bool, + summary: Option, modules: &mut Vec, modules_node: &mut Option, + force_insert: bool, ) -> Result> { tx = Box::pin(reduce_flow( tx, @@ -1647,9 +1649,22 @@ async fn insert_flow_modules<'c>( same_worker, )) .await?; - if modules.is_empty() || crate::worker_flow::is_simple_modules(modules, failure_module) { + if !force_insert + && (modules.is_empty() || crate::worker_flow::is_simple_modules(modules, failure_module)) + { return Ok(tx); } + + let flow_node_flow = FlowNodeFlow { + value: FlowValue { + modules: std::mem::take(modules), + failure_module: failure_module.cloned(), + same_worker, + ..Default::default() + }, + summary, + }; + let id; (tx, id) = insert_flow_node( tx, @@ -1657,12 +1672,7 @@ async fn insert_flow_modules<'c>( workspace_id, None, None, - Some(&Json(to_raw_value(&FlowValue { - modules: std::mem::take(modules), - failure_module: failure_module.cloned(), - same_worker, - ..Default::default() - }))), + Some(&Json(to_raw_value(&flow_node_flow))), None, ) .await?; @@ -1739,8 +1749,10 @@ async fn reduce_flow<'c>( workspace_id, failure_module, same_worker, + None, modules, modules_node, + false, ) .await?; } @@ -1752,8 +1764,10 @@ async fn reduce_flow<'c>( workspace_id, failure_module, same_worker, + None, &mut branch.modules, &mut branch.modules_node, + false, ) .await?; } @@ -1763,8 +1777,10 @@ async fn reduce_flow<'c>( workspace_id, failure_module, same_worker, + None, default, default_node, + false, ) .await?; } @@ -1776,12 +1792,28 @@ async fn reduce_flow<'c>( workspace_id, failure_module, same_worker, + None, &mut branch.modules, &mut branch.modules_node, + false, ) .await?; } } + AIAgent { tools, modules_node, .. } => { + tx = insert_flow_modules( + tx, + path, + workspace_id, + failure_module, + same_worker, + module.summary.clone(), // we only include summary for ai agents modules + tools, + modules_node, + true, + ) + .await?; + } _ => {} } module.value = to_raw_value(&val);