feat(aichat): add tool to test specific module in flow mode (#6381)

* draft

* cleaning

* cleaning

* fixes

* fix placeholders

* better fallback

* cleaning

frontend/src/lib/components/copilot/chat/AIChatDisplay.svelte

@@ -247,7 +247,6 @@
 				{availableContext}
 				{disabled}
 				isFirstMessage={messages.length === 0}
-				placeholder={messages.length === 0 ? 'Ask anything' : 'Ask followup'}
 			/>
 			<div
 				class={`flex flex-row ${

frontend/src/lib/components/copilot/chat/AIChatInput.svelte

@@ -45,23 +45,25 @@
 
 	// Generate mode-specific placeholder
 	const modePlaceholder = $derived.by(() => {
+		if (!isFirstMessage) {
+			return 'Ask followup'
+		}
+
 		if (placeholder) {
-			// If a custom placeholder is provided, use it
 			return placeholder
 		}
 		
-		// Generate placeholder based on current AI mode
 		switch (aiChatManager.mode) {
 			case AIMode.SCRIPT:
-				return 'Modify this script, fix errors, or generate new code...'
+				return 'Modify this script...'
 			case AIMode.FLOW:
-				return 'Edit this flow, add steps, or modify workflow logic...'
+				return 'Modify this flow...'
 			case AIMode.NAVIGATOR:
-				return 'Help me navigate Windmill or find features...'
+				return 'Navigate the app...'
 			case AIMode.API:
-				return 'Make API calls to fetch data or manage resources...'
+				return 'Make API calls...'
 			case AIMode.ASK:
-				return 'Ask questions about Windmill features and documentation...'
+				return 'Ask questions about Windmill...'
 			default:
 				return 'Ask anything'
 		}

frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts

@@ -8,7 +8,15 @@ import {
 } from './flow/core'
 import ContextManager from './ContextManager.svelte'
 import HistoryManager from './HistoryManager.svelte'
-import { extractCodeFromMarkdown, getLatestAssistantMessage, processToolCall, type DisplayMessage, type Tool, type ToolCallbacks, type ToolDisplayMessage } from './shared'
+import {
+	extractCodeFromMarkdown,
+	getLatestAssistantMessage,
+	processToolCall,
+	type DisplayMessage,
+	type Tool,
+	type ToolCallbacks,
+	type ToolDisplayMessage
+} from './shared'
 import type {
 	ChatCompletionChunk,
 	ChatCompletionMessageParam,
@@ -80,7 +88,9 @@ class AIChatManager {
 	helpers = $state<any | undefined>(undefined)
 
 	scriptEditorOptions = $state<ScriptOptions | undefined>(undefined)
-	scriptEditorApplyCode = $state<((code: string, applyAll?: boolean) => void) | undefined>(undefined)
+	scriptEditorApplyCode = $state<((code: string, applyAll?: boolean) => void) | undefined>(
+		undefined
+	)
 	scriptEditorShowDiffMode = $state<(() => void) | undefined>(undefined)
 	flowAiChatHelpers = $state<FlowAIChatHelpers | undefined>(undefined)
 	pendingNewCode = $state<string | undefined>(undefined)
@@ -660,7 +670,10 @@ class AIChatManager {
 			}
 			switch (this.mode) {
 				case AIMode.FLOW:
-					userMessage = prepareFlowUserMessage(oldInstructions, this.flowAiChatHelpers!.getFlowAndSelectedId())
+					userMessage = prepareFlowUserMessage(
+						oldInstructions,
+						this.flowAiChatHelpers!.getFlowAndSelectedId()
+					)
 					break
 				case AIMode.NAVIGATOR:
 					userMessage = prepareNavigatorUserMessage(oldInstructions)

frontend/src/lib/components/copilot/chat/flow/core.ts

@@ -12,10 +12,8 @@ import {
 	getLangContext,
 	SUPPORTED_CHAT_SCRIPT_LANGUAGES
 } from '../script/core'
-import { createSearchHubScriptsTool, createToolDef, type Tool, executeTestRun } from '../shared'
+import { createSearchHubScriptsTool, createToolDef, type Tool, executeTestRun, buildSchemaForTool, buildTestRunArgs } from '../shared'
 import type { ExtendedOpenFlow } from '$lib/components/flows/types'
-import { copilotSessionModel } from '$lib/stores'
-import { get } from 'svelte/store'
 
 export type AIModuleAction = 'added' | 'modified' | 'removed'
 
@@ -351,6 +349,21 @@ const testRunFlowToolDef = createToolDef(
 	'Execute a test run of the current flow'
 )
 
+const testRunStepSchema = z.object({
+	stepId: z.string().describe('The id of the step to test'),
+	args: z
+		.object({})
+		.nullable()
+		.optional()
+		.describe('Arguments to pass to the step (optional, uses default step inputs if not provided)')
+})
+
+const testRunStepToolDef = createToolDef(
+	testRunStepSchema,
+	'test_run_step',
+	'Execute a test run of a specific step in the flow'
+)
+
 const workspaceScriptsSearch = new WorkspaceScriptsSearch()
 
 export const flowTools: Tool<FlowAIChatHelpers>[] = [
@@ -364,7 +377,8 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 			const parsedArgs = searchScriptsSchema.parse(args)
 			const scriptResults = await workspaceScriptsSearch.search(parsedArgs.query, workspace)
 			toolCallbacks.setToolStatus(toolId, {
-				content: 'Found ' +
+				content:
+					'Found ' +
 					scriptResults.length +
 					' scripts in the workspace related to "' +
 					args.query +
@@ -378,19 +392,20 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 		fn: async ({ args, helpers, toolId, toolCallbacks }) => {
 			const parsedArgs = addStepSchema.parse(args)
 			toolCallbacks.setToolStatus(toolId, {
-				content: parsedArgs.location.type === 'after'
-					? `Adding a step after step '${parsedArgs.location.afterId}'`
-					: parsedArgs.location.type === 'start'
-						? 'Adding a step at the start'
-						: parsedArgs.location.type === 'start_inside_forloop'
-							? `Adding a step at the start of the forloop step '${parsedArgs.location.inside}'`
-							: parsedArgs.location.type === 'start_inside_branch'
-								? `Adding a step at the start of the branch ${parsedArgs.location.branchIndex + 1} of step '${parsedArgs.location.inside}'`
-								: parsedArgs.location.type === 'preprocessor'
-									? 'Adding a preprocessor step'
-									: parsedArgs.location.type === 'failure'
-										? 'Adding a failure step'
-										: 'Adding a step'
+				content:
+					parsedArgs.location.type === 'after'
+						? `Adding a step after step '${parsedArgs.location.afterId}'`
+						: parsedArgs.location.type === 'start'
+							? 'Adding a step at the start'
+							: parsedArgs.location.type === 'start_inside_forloop'
+								? `Adding a step at the start of the forloop step '${parsedArgs.location.inside}'`
+								: parsedArgs.location.type === 'start_inside_branch'
+									? `Adding a step at the start of the branch ${parsedArgs.location.branchIndex + 1} of step '${parsedArgs.location.inside}'`
+									: parsedArgs.location.type === 'preprocessor'
+										? 'Adding a preprocessor step'
+										: parsedArgs.location.type === 'failure'
+											? 'Adding a failure step'
+											: 'Adding a step'
 			})
 			const id = await helpers.insertStep(parsedArgs.location, parsedArgs.step)
 			helpers.selectStep(id)
@@ -469,7 +484,9 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 		def: setCodeToolDef,
 		fn: async ({ args, helpers, toolId, toolCallbacks }) => {
 			const parsedArgs = setCodeSchema.parse(args)
-			toolCallbacks.setToolStatus(toolId, { content: `Setting code for step '${parsedArgs.id}'...` })
+			toolCallbacks.setToolStatus(toolId, {
+				content: `Setting code for step '${parsedArgs.id}'...`
+			})
 			await helpers.setCode(parsedArgs.id, parsedArgs.code)
 			helpers.selectStep(parsedArgs.id)
 			toolCallbacks.setToolStatus(toolId, { content: `Set code for step '${parsedArgs.id}'` })
@@ -519,7 +536,9 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 			const parsedArgs = setForLoopIteratorExpressionSchema.parse(args)
 			await helpers.setForLoopIteratorExpression(parsedArgs.id, parsedArgs.expression)
 			helpers.selectStep(parsedArgs.id)
-			toolCallbacks.setToolStatus(toolId, { content: `Set forloop '${parsedArgs.id}' iterator expression` })
+			toolCallbacks.setToolStatus(toolId, {
+				content: `Set forloop '${parsedArgs.id}' iterator expression`
+			})
 			return `Forloop '${parsedArgs.id}' iterator expression set`
 		}
 	},
@@ -535,28 +554,33 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 				parsedArgs.query,
 				workspace
 			)
-			toolCallbacks.setToolStatus(toolId, { content: 'Retrieved resource types for "' + parsedArgs.query + '"' })
+			toolCallbacks.setToolStatus(toolId, {
+				content: 'Retrieved resource types for "' + parsedArgs.query + '"'
+			})
 			return formattedResourceTypes
 		}
 	},
 	{
 		def: testRunFlowToolDef,
-		fn: async ({ args, workspace, helpers, toolCallbacks, toolId }) => {
+		fn: async function({ args, workspace, helpers, toolCallbacks, toolId }) {
 			const { flow } = helpers.getFlowAndSelectedId()
 
 			if (!flow || !flow.value) {
-				toolCallbacks.setToolStatus(toolId, { 
+				toolCallbacks.setToolStatus(toolId, {
 					content: 'No flow available to test',
 					error: 'No flow found in current context'
 				})
-				throw new Error('No flow available to test. Please ensure you have a flow open in the editor.')
+				throw new Error(
+					'No flow available to test. Please ensure you have a flow open in the editor.'
+				)
 			}
 
+			const parsedArgs = await buildTestRunArgs(args, this.def)
 			return executeTestRun({
 				jobStarter: () => JobService.runFlowPreview({
 					workspace: workspace,
 					requestBody: {
-						args: args,
+						args: parsedArgs,
 						value: flow.value,
 					}
 				}),
@@ -568,25 +592,120 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 			})
 		},
 		setSchema: async function(helpers: FlowAIChatHelpers) {
-			try {
-				if (this.def?.function?.parameters) {
-					const flowInputsSchema = await helpers.getFlowInputsSchema()
-					this.def.function.parameters = { ...flowInputsSchema, additionalProperties: false }
-					// OPEN AI models don't support strict mode well with schema with complex properties, so we disable it
-					const model = get(copilotSessionModel)?.provider
-					if (model === 'openai' || model === 'azure_openai') {
-						this.def.function.strict = false
-					}
-				}
-			} catch (e) {
-				console.error('Error setting schema for test_run_flow tool', e)
-				// fallback to schema with any properties
-				this.def.function.parameters = {
-					type: 'object',
-					properties: {},
-					additionalProperties: true,
-					strict: false
-				}
+			await buildSchemaForTool(this.def, async () => {
+				const flowInputsSchema = await helpers.getFlowInputsSchema()
+				return flowInputsSchema
+			})
+		},
+		requiresConfirmation: true,
+		showDetails: true
+	},
+	{
+		// set strict to false to avoid issues with open ai models
+		def: { ...testRunStepToolDef, function: { ...testRunStepToolDef.function, strict: false } },
+		fn: async ({ args, workspace, helpers, toolCallbacks, toolId }) => {
+			const { flow } = helpers.getFlowAndSelectedId()
+
+			if (!flow || !flow.value) {
+				toolCallbacks.setToolStatus(toolId, {
+					content: 'No flow available to test step from',
+					error: 'No flow found in current context'
+				})
+				throw new Error(
+					'No flow available to test step from. Please ensure you have a flow open in the editor.'
+				)
+			}
+
+			const stepId = args.stepId
+			const stepArgs = args.args || {}
+
+			// Find the step in the flow
+			const modules = helpers.getModules()
+			let targetModule: FlowModule | undefined = modules.find((m) => m.id === stepId)
+
+			if (!targetModule) {
+				toolCallbacks.setToolStatus(toolId, {
+					content: `Step '${stepId}' not found in flow`,
+					error: `Step with id '${stepId}' does not exist in the current flow`
+				})
+				throw new Error(
+					`Step with id '${stepId}' not found in flow. Available steps: ${modules.map((m) => m.id).join(', ')}`
+				)
+			}
+
+			const module = targetModule
+			const moduleValue = module.value
+
+			if (moduleValue.type === 'rawscript') {
+				// Test raw script step
+
+				return executeTestRun({
+					jobStarter: () =>
+						JobService.runScriptPreview({
+							workspace: workspace,
+							requestBody: {
+								content: moduleValue.content ?? '',
+								language: moduleValue.language,
+								args: module.id === 'preprocessor' ? { _ENTRYPOINT_OVERRIDE: 'preprocessor', ...stepArgs } : stepArgs
+							}
+						}),
+					workspace,
+					toolCallbacks,
+					toolId,
+					startMessage: `Starting test run of step '${stepId}'...`,
+					contextName: 'script'
+				})
+			} else if (moduleValue.type === 'script') {
+				// Test script step - need to get the script content
+				const script = moduleValue.hash
+					? await ScriptService.getScriptByHash({
+							workspace: workspace,
+							hash: moduleValue.hash
+						})
+					: await ScriptService.getScriptByPath({
+							workspace: workspace,
+							path: moduleValue.path
+						})
+
+				return executeTestRun({
+					jobStarter: () =>
+						JobService.runScriptPreview({
+							workspace: workspace,
+							requestBody: {
+								content: script.content,
+								language: script.language,
+								args: module.id === 'preprocessor' ? { _ENTRYPOINT_OVERRIDE: 'preprocessor', ...stepArgs } : stepArgs,
+							}
+						}),
+					workspace,
+					toolCallbacks,
+					toolId,
+					startMessage: `Starting test run of script step '${stepId}'...`,
+					contextName: 'script'
+				})
+			} else if (moduleValue.type === 'flow') {
+				// Test flow step
+				return executeTestRun({
+					jobStarter: () =>
+						JobService.runFlowByPath({
+							workspace: workspace,
+							path: moduleValue.path,
+							requestBody: stepArgs
+						}),
+					workspace,
+					toolCallbacks,
+					toolId,
+					startMessage: `Starting test run of flow step '${stepId}'...`,
+					contextName: 'flow'
+				})
+			} else {
+				toolCallbacks.setToolStatus(toolId, {
+					content: `Step type '${moduleValue.type}' not supported for testing`,
+					error: `Cannot test step of type '${moduleValue.type}'`
+				})
+				throw new Error(
+					`Cannot test step of type '${moduleValue.type}'. Supported types: rawscript, script, flow`
+				)
 			}
 		},
 		requiresConfirmation: true,
@@ -599,7 +718,7 @@ export function prepareFlowSystemMessage(): ChatCompletionSystemMessageParam {
 Follow the user instructions carefully.
 Go step by step, and explain what you're doing as you're doing it.
 DO NOT wait for user confirmation before performing an action. Only do it if the user explicitly asks you to wait in their initial instructions.
-ALWAYS use the \`test_run_flow\` tool to test the flow, and iterate on the flow until it works as expected. If the user cancels the test run, do not try again and wait for the next user instruction.
+ALWAYS test your modifications. You have access to the \`test_run_flow\` and \`test_run_step\` tools to test the flow and steps. If you only modified a single step, use the \`test_run_step\` tool to test it. If you modified the flow, use the \`test_run_flow\` tool to test it. If the user cancels the test run, do not try again and wait for the next user instruction.
 
 ## Understanding User Requests
 

frontend/src/lib/components/copilot/chat/script/core.ts

@@ -13,7 +13,7 @@ import { scriptLangToEditorLang } from '$lib/scripts'
 import { getDbSchemas } from '$lib/components/apps/components/display/dbtable/utils'
 import type { CodePieceElement, ContextElement } from '../context'
 import { PYTHON_PREPROCESSOR_MODULE_CODE, TS_PREPROCESSOR_MODULE_CODE } from '$lib/script_helpers'
-import { createSearchHubScriptsTool, type Tool, executeTestRun } from '../shared'
+import { createSearchHubScriptsTool, type Tool, executeTestRun, buildSchemaForTool, buildTestRunArgs } from '../shared'
 import { setupTypeAcquisition, type DepsToGet } from '$lib/ata'
 import { getModelContextWindow } from '../../lib'
 import { inferArgs } from '$lib/infer'
@@ -855,7 +855,7 @@ const TEST_RUN_SCRIPT_TOOL: ChatCompletionTool = {
 
 export const testRunScriptTool: Tool<ScriptChatHelpers> = {
 	def: TEST_RUN_SCRIPT_TOOL,
-	fn: async ({ args, workspace, helpers, toolCallbacks, toolId }) => {
+	fn: async function({ args, workspace, helpers, toolCallbacks, toolId }) {
 		const scriptOptions = helpers.getScriptOptions()
 		
 		if (!scriptOptions) {
@@ -880,17 +880,16 @@ export const testRunScriptTool: Tool<ScriptChatHelpers> = {
 			toolCallbacks.setToolStatus(toolId, { content: 'Code changes applied, starting test...' })
 		}
 
+		const parsedArgs = await buildTestRunArgs(args, this.def)
+
 		return executeTestRun({
 			jobStarter: () => JobService.runScriptPreview({
 				workspace: workspace,
 				requestBody: {
 					path: scriptOptions.path,
 					content: codeToTest,
-					args,
+					args: parsedArgs,
 					language: scriptOptions.lang as ScriptLang,
-					tag: undefined,
-					lock: undefined,
-					script_hash: undefined
 				}
 			}),
 			workspace,
@@ -901,7 +900,7 @@ export const testRunScriptTool: Tool<ScriptChatHelpers> = {
 		})
 	},
 	setSchema: async function(helpers: ScriptChatHelpers) {
-		try {
+		await buildSchemaForTool(this.def, async () => {
 			const scriptOptions = helpers.getScriptOptions()
 			const code = scriptOptions?.code
 			const lang = scriptOptions?.lang
@@ -911,23 +910,10 @@ export const testRunScriptTool: Tool<ScriptChatHelpers> = {
 			if (codeToTest) {
 				const newSchema = emptySchema()
 				await inferArgs(lang, codeToTest, newSchema)
-				this.def.function.parameters = { ...newSchema, additionalProperties: false }
-				// OPEN AI models don't support strict mode well with schema with complex properties, so we disable it
-				const model = get(copilotSessionModel)?.provider
-				if (model === 'openai' || model === 'azure_openai') {
-					this.def.function.strict = false
-				}
+				return newSchema
 			}
-		} catch (e) {
-			console.error('Error setting schema for test_run_script tool', e)
-			// fallback to schema with any properties
-			this.def.function.parameters = {
-				type: 'object',
-				properties: {},
-				additionalProperties: true,
-				strict: false
-			}
-		}
+			return emptySchema()
+		})
 	},
 	requiresConfirmation: true,
 	showDetails: true,

frontend/src/lib/components/copilot/chat/shared.ts

@@ -5,7 +5,7 @@ import type {
 } from 'openai/resources/chat/completions.mjs'
 import { get } from 'svelte/store'
 import type { ContextElement } from './context'
-import { workspaceStore } from '$lib/stores'
+import { copilotSessionModel, workspaceStore } from '$lib/stores'
 import type { ExtendedOpenFlow } from '$lib/components/flows/types'
 import type { FunctionParameters } from 'openai/resources/shared.mjs'
 import { zodToJsonSchema } from 'zod-to-json-schema'
@@ -253,6 +253,32 @@ export const createSearchHubScriptsTool = (withContent: boolean = false) => ({
 	}
 })
 
+export async function buildSchemaForTool(toolDef: ChatCompletionTool, schemaBuilder: () => Promise<FunctionParameters>): Promise<boolean> {
+	try {
+		const schema = await schemaBuilder()
+
+		// if schema properties contains values different from '^[a-zA-Z0-9_.-]{1,64}$'
+		const invalidProperties = Object.keys(schema.properties ?? {}).filter((key) => !/^[a-zA-Z0-9_.-]{1,64}$/.test(key))
+		if (invalidProperties.length > 0) {
+			console.warn(`Invalid flow inputs schema: ${invalidProperties.join(', ')}`)
+			throw new Error(`Invalid flow inputs schema: ${invalidProperties.join(', ')}`)
+		}
+
+		toolDef.function.parameters = { ...schema, additionalProperties: false }
+		// OPEN AI models don't support strict mode well with schema with complex properties, so we disable it
+		const model = get(copilotSessionModel)?.provider
+		if (model === 'openai' || model === 'azure_openai') {
+			toolDef.function.strict = false
+		}
+		return true
+	} catch (error) {
+		console.error('Error building schema for tool', error)
+		// fallback to schema with args as a JSON string
+		toolDef.function.parameters = { type: 'object', properties: { args: { type: 'string', description: 'JSON string containing the arguments for the tool' } }, additionalProperties: false, strict: false, required: ['args'] }
+		return false
+	}
+}
+
 // Constants for result formatting
 const MAX_RESULT_LENGTH = 12000
 const MAX_LOG_LENGTH = 4000
@@ -362,6 +388,20 @@ function getErrorMessage(result: unknown): string {
 	return 'Unknown error'
 }
 
+// Build test run args based on the tool definition, if it contains a fallback schema
+export async function buildTestRunArgs(args: any, toolDef: ChatCompletionTool): Promise<any> {
+	let parsedArgs = args
+	// if the schema is the fallback schema, parse the args as a JSON string
+	if ((toolDef.function.parameters as any).properties?.args?.description === 'JSON string containing the arguments for the tool') {
+		try {
+			parsedArgs = JSON.parse(args.args)
+		} catch (error) {
+			console.error('Error parsing arguments for tool', error)
+		}
+	}
+	return parsedArgs
+}
+
 // Main execution function for test runs
 export async function executeTestRun(config: TestRunConfig): Promise<string> {
 	try {