mirror of
https://github.com/windmill-labs/windmill.git
synced 2026-10-13 00:03:00 +00:00
fix: poll decision jobs, keep empty-branch answers and deployed defaults
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
a1a08785ec
commit
1859cfca9c
@@ -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()
|
||||
)
|
||||
|
||||
@@ -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.<id>`, 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(())
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
|
||||
@@ -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<Branch>, Vec<FlowModule>, 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]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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<String, Box<RawValue>>,
|
||||
as_tool: bool,
|
||||
canceled_by: &mut Option<CanceledBy>,
|
||||
mem_peak: &mut i32,
|
||||
occupancy_metrics: &mut OccupancyMetrics,
|
||||
worker_name: &str,
|
||||
) -> error::Result<Box<RawValue>> {
|
||||
let args = serde_json::from_str::<AIDecisionArgs>(&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<String, InputTransform>)
|
||||
Some("score") => {
|
||||
let levels = criteria
|
||||
.and_then(|c| c.as_array())
|
||||
.map(|c| c.iter().filter_map(|l| l.as_str()).collect::<Vec<_>>().join(" < "))
|
||||
.map(|c| {
|
||||
c.iter()
|
||||
.filter_map(|l| l.as_str())
|
||||
.collect::<Vec<_>>()
|
||||
.join(" < ")
|
||||
})
|
||||
.unwrap_or_default();
|
||||
format!("a score over {levels}")
|
||||
}
|
||||
|
||||
@@ -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.<id>`, 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.<previous id>` 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<Error> = None;
|
||||
for (i, b) in branches.iter().enumerate() {
|
||||
let pred_res = compute_bool_from_expr(
|
||||
|
||||
@@ -148,8 +148,8 @@ export const AI_PROVIDERS: Record<ChatAIProvider, AIProviderDetails> = {
|
||||
}
|
||||
}
|
||||
|
||||
/** 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<AIProvider, ChatAIProvider>,
|
||||
AIProviderDetails
|
||||
|
||||
Reference in New Issue
Block a user