diff --git a/backend/migrations/20260921150919_typesafe_resource_type.up.sql b/backend/migrations/20260921150919_typesafe_resource_type.up.sql index f97d8009c3..e3b31c24fa 100644 --- a/backend/migrations/20260921150919_typesafe_resource_type.up.sql +++ b/backend/migrations/20260921150919_typesafe_resource_type.up.sql @@ -1,4 +1,4 @@ --- The credential an AI agent step with decision output runs TypeSafe's Jev with. Seeded under +-- The credential an AI decision step runs TypeSafe's Jev with. Seeded under -- 'admins' so every workspace sees it, in the shape the hub gives the other AI providers, which -- the step reads as `api_key` and `base_url`. A hub sync that later publishes the same name -- updates the schema in place. @@ -26,7 +26,7 @@ VALUES ( }, "required": ["api_key"] }', - 'API key for TypeSafe''s Jev decision model, which answers typed questions with calibrated probabilities. Used by AI agent steps with decision output.', + 'API key for TypeSafe''s Jev decision model, which answers typed questions with calibrated probabilities. Used by AI decision steps and tools.', 'system', now() ) diff --git a/backend/tests/ai_decision.rs b/backend/tests/ai_decision.rs index fa34b5544d..35c7c813aa 100644 --- a/backend/tests/ai_decision.rs +++ b/backend/tests/ai_decision.rs @@ -52,24 +52,25 @@ async fn test_ai_decision_runs_the_branch_its_answer_picks( let server = ApiServer::start(db.clone()).await?; let port = server.addr.port(); + let decision_inputs = json!({ + "provider": {"type": "static", "value": { + "kind": "typesafe", + "resource": {"api_key": "key", "base_url": format!("http://127.0.0.1:{stub_port}/v1")}, + "model": "jev-latest" + }}, + "state": {"type": "static", "value": "I was charged twice"}, + "questions": {"type": "static", "value": {"intent": { + "type": "choice", + "instructions": "What does the customer want?", + "criteria": {"refund": "Money back", "bug": "Something is broken"} + }}} + }); let flow: FlowValue = serde_json::from_value(json!({ "modules": [{ "id": "d", "value": { "type": "aidecision", - "input_transforms": { - "provider": {"type": "static", "value": { - "kind": "typesafe", - "resource": {"api_key": "key", "base_url": format!("http://127.0.0.1:{stub_port}/v1")}, - "model": "jev-latest" - }}, - "state": {"type": "static", "value": "I was charged twice"}, - "questions": {"type": "static", "value": {"intent": { - "type": "choice", - "instructions": "What does the customer want?", - "criteria": {"refund": "Money back", "bug": "Something is broken"} - }}} - }, + "input_transforms": decision_inputs, "branches": [ { "expr": "previous_result.output.intent.choice === 'bug'", @@ -114,5 +115,57 @@ async fn test_ai_decision_runs_the_branch_its_answer_picks( decision_job.is_some(), "the decision job is kept once the branch has run" ); + + // A branch condition still reads the step before the decision as `results.`, while + // `previous_result` is the answers. Once an empty branch is chosen, `results.d` stays the + // answers for every later step, not only the next one. + let flow: FlowValue = serde_json::from_value(json!({ + "modules": [ + {"id": "pre", "value": { + "type": "rawscript", + "language": "deno", + "content": "export function main(){ return { allow: true } }", + "input_transforms": {} + }}, + {"id": "d", "value": { + "type": "aidecision", + "input_transforms": decision_inputs, + "branches": [{ + "expr": "results.pre.allow && previous_result.output.intent.choice === 'refund'", + "modules": [] + }] + }}, + {"id": "next", "value": {"type": "identity"}}, + {"id": "after", "value": { + "type": "rawscript", + "language": "deno", + "content": "export function main(choice){ return choice }", + "input_transforms": {"choice": { + "type": "javascript", + "expr": "results.d.output.intent.choice" + }} + }} + ] + }))?; + let job = run_job_in_new_worker_until_complete( + &db, + false, + JobPayload::RawFlow { value: flow, path: None, restarted_from: None }, + port, + ) + .await; + assert_eq!(job.json_result(), Some(json!("refund"))); + let status: FlowStatus = serde_json::from_value(job.flow_status.expect("flow status"))?; + assert!( + matches!( + &status.modules[1], + FlowStatusModule::Success { + branch_chosen: Some(BranchChosen::Branch { branch: 0 }), + .. + } + ), + "the branch reading the previous step was not chosen: {:?}", + status.modules[1] + ); Ok(()) } diff --git a/backend/windmill-ai/src/ai_providers.rs b/backend/windmill-ai/src/ai_providers.rs index 776d52f0ca..dcb57ffbd7 100644 --- a/backend/windmill-ai/src/ai_providers.rs +++ b/backend/windmill-ai/src/ai_providers.rs @@ -124,7 +124,7 @@ pub enum AIProvider { #[serde(rename = "aws_bedrock")] AWSBedrock, /// TypeSafe's Jev decision model. It answers typed questions rather than chatting, so only - /// an agent step with decision output runs it (`run_decision`), never the agent loop. + /// an AI decision runs it (`run_systemone`), never the agent loop. TypeSafe, } diff --git a/backend/windmill-types/src/flows.rs b/backend/windmill-types/src/flows.rs index 8c87034a69..7cfcaaeffd 100644 --- a/backend/windmill-types/src/flows.rs +++ b/backend/windmill-types/src/flows.rs @@ -1150,13 +1150,29 @@ impl BranchChoiceShell { FlowModuleValue::BranchOne { branches, default, default_node } } BranchChoiceShell::AIDecision { input_transforms, default_node, tag } => { - FlowModuleValue::AIDecision { input_transforms, branches, default, default_node, tag } + FlowModuleValue::AIDecision { + input_transforms, + branches, + default, + default_node, + tag, + } } } } } impl FlowModuleValue { + /// Whether this is an AI decision that runs a branch once it has answered. A deployed flow + /// moves a default's steps into `default_node`, leaving `default` empty. + pub fn is_branched_ai_decision(&self) -> bool { + matches!( + self, + FlowModuleValue::AIDecision { branches, default, default_node, .. } + if !branches.is_empty() || !default.is_empty() || default_node.is_some() + ) + } + /// Splits a step that runs one of its branches into its branches, its default and the rest, /// so a pass rewriting branches handles `BranchOne` and AI decisions alike. Any other step is /// given back unchanged. @@ -1164,16 +1180,22 @@ impl FlowModuleValue { self, ) -> Result<(Vec, Vec, BranchChoiceShell), FlowModuleValue> { match self { - FlowModuleValue::BranchOne { branches, default, default_node } => { - Ok((branches, default, BranchChoiceShell::BranchOne { default_node })) - } - FlowModuleValue::AIDecision { input_transforms, branches, default, default_node, tag } => { - Ok(( - branches, - default, - BranchChoiceShell::AIDecision { input_transforms, default_node, tag }, - )) - } + FlowModuleValue::BranchOne { branches, default, default_node } => Ok(( + branches, + default, + BranchChoiceShell::BranchOne { default_node }, + )), + FlowModuleValue::AIDecision { + input_transforms, + branches, + default, + default_node, + tag, + } => Ok(( + branches, + default, + BranchChoiceShell::AIDecision { input_transforms, default_node, tag }, + )), other => Err(other), } } @@ -1707,6 +1729,13 @@ mod tests { output["branches"][0]["expr"], json!("previous_result.output.intent.choice === 'refund'") ); + assert!(val.is_branched_ai_decision()); + + // A deployed default-only decision keeps its default as a flow node. + let deployed = + json!({"type": "aidecision", "input_transforms": {}, "default": [], "default_node": 5}); + let val: FlowModuleValue = serde_json::from_value(deployed).unwrap(); + assert!(val.is_branched_ai_decision()); } #[test] diff --git a/backend/windmill-worker/src/ai/tools.rs b/backend/windmill-worker/src/ai/tools.rs index f179e06998..6229559bb8 100644 --- a/backend/windmill-worker/src/ai/tools.rs +++ b/backend/windmill-worker/src/ai/tools.rs @@ -491,9 +491,7 @@ async fn execute_windmill_tool( } } // A tool returns its answers to the calling model, so it has no flow to branch into. - FlowModuleValue::AIDecision { branches, default, .. } - if !branches.is_empty() || !default.is_empty() => - { + ref value @ FlowModuleValue::AIDecision { .. } if value.is_branched_ai_decision() => { return Err(Error::BadRequest(format!( "The AI decision tool {} has branches, which only a flow step can run", tool_call.function.name diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs index 6d824daab8..a3bc7744b1 100644 --- a/backend/windmill-worker/src/ai_executor.rs +++ b/backend/windmill-worker/src/ai_executor.rs @@ -61,7 +61,10 @@ use crate::{ common::{ build_args_map, resolve_job_timeout, transform_json_value, OccupancyMetrics, StreamNotifier, }, - handle_child::{run_future_with_polling_update_job_poller_graceful, GracefulPollOutcome}, + handle_child::{ + run_future_with_polling_update_job_poller, + run_future_with_polling_update_job_poller_graceful, GracefulPollOutcome, + }, }; lazy_static::lazy_static! { @@ -631,6 +634,10 @@ pub async fn handle_ai_agent_job( job, &local_args, direct_parent_job_kind == JobKind::AIAgent, + canceled_by, + mem_peak, + occupancy_metrics, + worker_name, ) .await; } @@ -2116,6 +2123,10 @@ async fn handle_ai_decision( job: &MiniPulledJob, local_args: &HashMap>, as_tool: bool, + canceled_by: &mut Option, + mem_peak: &mut i32, + occupancy_metrics: &mut OccupancyMetrics, + worker_name: &str, ) -> error::Result> { let args = serde_json::from_str::(&serde_json::to_string(local_args)?)?; if args.provider.kind != AIProvider::TypeSafe { @@ -2129,12 +2140,24 @@ async fn handle_ai_decision( let timeout = resolve_job_timeout(conn, &job.workspace_id, job.id, job.timeout) .await .0; - let result = run_systemone( - &credentials, - args.provider.get_model(), - state, - questions, - timeout, + // Under the job poller, so the job keeps its heartbeat and a cancel drops the request. + let result = run_future_with_polling_update_job_poller( + job.id, + job.timeout, + conn, + mem_peak, + canceled_by, + run_systemone( + &credentials, + args.provider.get_model(), + state, + questions, + timeout, + ), + worker_name, + &job.workspace_id, + &mut Some(occupancy_metrics), + Box::pin(futures::stream::once(async { 0 })), ) .await?; windmill_common::feature_usage::log_feature_usage( @@ -2169,7 +2192,12 @@ fn decision_tool_description(input_transforms: &HashMap) Some("score") => { let levels = criteria .and_then(|c| c.as_array()) - .map(|c| c.iter().filter_map(|l| l.as_str()).collect::>().join(" < ")) + .map(|c| { + c.iter() + .filter_map(|l| l.as_str()) + .collect::>() + .join(" < ") + }) .unwrap_or_default(); format!("a score over {levels}") } diff --git a/backend/windmill-worker/src/worker_flow.rs b/backend/windmill-worker/src/worker_flow.rs index 6f71c6a70f..275aa46f53 100644 --- a/backend/windmill-worker/src/worker_flow.rs +++ b/backend/windmill-worker/src/worker_flow.rs @@ -1304,11 +1304,17 @@ pub async fn update_flow_status_after_job_completion_internal( false }; success = true; + // An AI decision whose chosen branch is empty completes with no job of its + // own: its job, and so `results.`, stays the decision's answers. + let job = match module_status.decision_job() { + Some(decision_job) if job_id_for_status.is_nil() => decision_job, + _ => *job_id_for_status, + }; ( true, Some(FlowStatusModule::Success { id: module_status.id(), - job: job_id_for_status.clone(), + job, flow_jobs, flow_jobs_success, flow_jobs_duration, @@ -3913,11 +3919,10 @@ async fn push_next_flow_job( // we get the args from the last failed job, except for a decision with branches: // it may have failed in its branch, whose sub-flow took the flow's args, so it // asks again from its own inputs. - if matches!( - module.get_value(), - Ok(FlowModuleValue::AIDecision { ref branches, ref default, .. }) - if !branches.is_empty() || !default.is_empty() - ) { + if module + .get_value() + .is_ok_and(|value| value.is_branched_ai_decision()) + { None } else { status.retry.failed_jobs.last() @@ -5417,18 +5422,18 @@ async fn compute_next_flow_transform( } // The decision job runs as an AI agent job; the arm below takes over once a decision with // branches has answered. - FlowModuleValue::AIDecision { ref branches, ref default, tag, .. } + ref value @ FlowModuleValue::AIDecision { ref tag, .. } if !matches!( status_module, FlowStatusModule::InProgress { decision_job: Some(_), branch_chosen: None, .. } ) => { - let has_branches = !branches.is_empty() || !default.is_empty(); + let has_branches = value.is_branched_ai_decision(); let path = get_path(flow_job, status, module); Ok(NextFlowTransform::Continue( ContinuePayload::SingleJob(JobPayloadWithTag { payload: JobPayload::AIAgent { path }, - tag: tag.filter(|t| !t.trim().is_empty()), + tag: tag.clone().filter(|t| !t.trim().is_empty()), delete_after_use, delete_after_secs, timeout: None, @@ -5728,7 +5733,16 @@ async fn compute_next_flow_transform( .. } => { let mut branch_chosen = BranchChosen::Default; - let idcontext = get_transform_context(&flow_job, previous_id, &status); + // `results.` reads `previous_result`, which for an answered + // decision is its own answers, so the decision is the previous step here. + let predicate_previous_id = match status_module { + FlowStatusModule::InProgress { decision_job: Some(_), .. } => { + module.id.as_str() + } + _ => previous_id, + }; + let idcontext = + get_transform_context(&flow_job, predicate_previous_id, &status); let mut predicate_err: Option = None; for (i, b) in branches.iter().enumerate() { let pred_res = compute_bool_from_expr( diff --git a/frontend/src/lib/components/copilot/lib.ts b/frontend/src/lib/components/copilot/lib.ts index a672e648f4..899751b108 100644 --- a/frontend/src/lib/components/copilot/lib.ts +++ b/frontend/src/lib/components/copilot/lib.ts @@ -148,8 +148,8 @@ export const AI_PROVIDERS: Record = { } } -/** TypeSafe's Jev answers typed questions instead of messages, so only an agent step with decision - * output offers it. The pinned version is there for flows tuned against its probabilities. */ +/** TypeSafe's Jev answers typed questions instead of messages, so only an AI decision offers it. + * The pinned version is there for flows tuned against its probabilities. */ export const DECISION_AI_PROVIDERS: Record< Exclude, AIProviderDetails