From 58972e75a528b437661c8fe02acf694f55ff0923 Mon Sep 17 00:00:00 2001
From: "claude[bot]" <209825114+claude[bot]@users.noreply.github.com>
Date: Mon, 13 Oct 2025 12:16:39 +0000
Subject: [PATCH] feat(aichat): handle ai agent steps in flow mode

- Added support for 'aiagent' type in the addstep tool schema
- Updated insertStep function to handle AI agent modules
- Modified getStepInputs and setStepInputs to work with AI agent input_transforms
- Added documentation for AI agent steps in flow system instructions

Fixes #6801

Co-authored-by: windmill-internal-app[bot] <windmill-internal-app[bot]@users.noreply.github.com>
---
 .../copilot/chat/flow/FlowAIChat.svelte         | 17 ++++++++++++++---
 .../lib/components/copilot/chat/flow/core.ts    | 12 +++++++++++-
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte b/frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte
index 5dd76d2ba0..ee562b85be 100644
--- a/frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte
+++ b/frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte
@@ -325,6 +325,17 @@
 					)
 					break
 				}
+				case 'aiagent': {
+					if (location.type === 'preprocessor' || location.type === 'failure') {
+						throw new Error('Cannot insert an AI agent module for preprocessing or error handling')
+					}
+					newModules = await flowModuleSchemaMap?.insertNewModuleAtIndex(
+						modules,
+						indexToInsertAt,
+						step.type
+					)
+					break
+				}
 				default: {
 					throw new Error('Unknown step type')
 				}
@@ -363,7 +374,7 @@
 				throw new Error('Module not found')
 			}
 			const inputs =
-				module.value.type === 'script' || module.value.type === 'rawscript'
+				module.value.type === 'script' || module.value.type === 'rawscript' || module.value.type === 'aiagent'
 					? module.value.input_transforms
 					: {}
 
@@ -397,8 +408,8 @@
 					throw new Error('Module not found')
 				}
 
-				if (module.value.type !== 'script' && module.value.type !== 'rawscript') {
-					throw new Error('Module is not a script or rawscript')
+				if (module.value.type !== 'script' && module.value.type !== 'rawscript' && module.value.type !== 'aiagent') {
+					throw new Error('Module is not a script, rawscript, or aiagent')
 				}
 
 				for (const { input, value } of parsedInputs) {
diff --git a/frontend/src/lib/components/copilot/chat/flow/core.ts b/frontend/src/lib/components/copilot/chat/flow/core.ts
index 1fdd8895ff..6d77a07bbe 100644
--- a/frontend/src/lib/components/copilot/chat/flow/core.ts
+++ b/frontend/src/lib/components/copilot/chat/flow/core.ts
@@ -126,7 +126,13 @@ const newStepSchema = z.union([
 		})
 		.describe(
 			'Add a branch one at the specified location: only the first branch that evaluates to true will be executed'
-		)
+		),
+	z
+		.object({
+			type: z.literal('aiagent'),
+			parallel: z.boolean().optional().describe('Whether to execute tools in parallel')
+		})
+		.describe('Add an AI agent step at the specified location')
 ])
 
 type NewStep = z.infer<typeof newStepSchema>
@@ -942,6 +948,10 @@ For special step types, follow these additional steps:
   - Set advanced options (parallel, parallelism, skip_failures) using set_forloop_options
 - For branchone steps: Set the predicates for each branch using set_branch_predicate
 - For branchall steps: No additional setup needed
+- For aiagent steps:
+  - Set the inputs using set_step_inputs (provider, user_message, system_prompt, etc.)
+  - Tools can be added as nested steps inside the AI agent
+  - Optionally set parallel: true to execute tools in parallel
 
 ### Module Control Options
 For any module type, you can set control flow options using set_module_control_options: