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

feat(backend): use flow nodes opti for ai agent steps (#6808)

Browse Source
This commit is contained in:
hugocasa
2025-10-14 20:06:27 +02:00
committed by GitHub
parent 9d9c29fdfa
commit 8d5acda340
8 changed files with 179 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/windmill-api/src/jobs.rs
+1 -1
View File
@@ -3605,7 +3605,7 @@ where
}
fn decode_payload<D: DeserializeOwned>(t: String) -> anyhow::Result<D> {
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")
backend/windmill-common/src/cache.rs
+21 -6
View File
@@ -282,12 +282,25 @@ pub mod future {
pub struct FlowData {
pub raw_flow: Box<RawValue>,
pub flow: FlowValue,
pub summary: Option<String>,
}
impl FlowData {
pub fn from_raw(raw_flow: Box<RawValue>) -> error::Result<Self> {
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::<crate::flows::FlowNodeFlow>(raw_flow.get())
{
(parsed.value, parsed.summary)
} else {
// fallback to plain FlowValue
(
serde_json::from_str::<FlowValue>(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
backend/windmill-common/src/flows.rs
+13
View File
@@ -141,6 +141,13 @@ pub struct FlowValue {
pub chat_input_enabled: Option<bool>,
}
#[derive(Serialize, Deserialize)]
pub struct FlowNodeFlow {
pub value: FlowValue,
#[serde(skip_serializing_if = "Option::is_none")]
pub summary: Option<String>,
}
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<String, InputTransform>,
tools: Vec<FlowModule>,
#[serde(skip_serializing_if = "Option::is_none")]
modules_node: Option<FlowNodeId>,
},
}
@@ -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);
backend/windmill-common/src/jobs.rs
+1
View File
@@ -425,6 +425,7 @@ pub enum JobPayload {
Noop,
AIAgent {
path: String,
flow_node_id: Option<FlowNodeId>,
},
}
backend/windmill-queue/src/jobs.rs
+5 -4
View File
@@ -3737,7 +3737,8 @@ pub async fn push<'c, 'd>(
if let Some(skip_handler) = skip_handler {
let mut skip_input_transforms = HashMap::<String, InputTransform>::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,
backend/windmill-worker/src/ai_executor.rs
+96 -39
View File
@@ -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<String> {
fn get_step_name_from_flow(summary: Option<&str>, flow_step_id: Option<&str>) -> Option<String> {
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 {
backend/windmill-worker/src/worker_flow.rs
+2 -2
View File
@@ -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,
backend/windmill-worker/src/worker_lockfiles.rs
+40 -8
View File
@@ -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<FlowModule>>,
same_worker: bool,
summary: Option<String>,
modules: &mut Vec<FlowModule>,
modules_node: &mut Option<FlowNodeId>,
force_insert: bool,
) -> Result<sqlx::Transaction<'c, sqlx::Postgres>> {
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);