From f5ccabda5212da67ae472a53fca8fd7bed3d8f66 Mon Sep 17 00:00:00 2001 From: centdix <40307056+centdix@users.noreply.github.com> Date: Wed, 3 Sep 2025 18:43:53 +0200 Subject: [PATCH] feat(flow): add structured output option for ai agent step (#6515) * add ui for structured output * implement backend logic for openai models * simulate having non required props * cleaning * implement logic for anthropic * cleaning * cleaning * cleaning * avoid name clash * return object * focus new field when adding * fix condition * small opti * use box raw value * avoid unnecessary parsing + return error if parsing fails --- backend/windmill-worker/src/ai_executor.rs | 264 ++++++++++++++++-- frontend/src/lib/components/ArgInput.svelte | 13 + .../lib/components/EditableSchemaForm.svelte | 7 +- .../src/lib/components/flows/flowInfers.ts | 8 +- 4 files changed, 271 insertions(+), 21 deletions(-) diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs index 398d8f876d..11ec4d7069 100644 --- a/backend/windmill-worker/src/ai_executor.rs +++ b/backend/windmill-worker/src/ai_executor.rs @@ -42,13 +42,34 @@ lazy_static::lazy_static! { static ref TOOL_NAME_REGEX: Regex = Regex::new(r"^[a-zA-Z0-9_]+$").unwrap(); } -#[derive(Deserialize, Serialize, Clone)] +/// Find a unique tool name to avoid collisions with user-provided tools +fn find_unique_tool_name(base_name: &str, existing_tools: Option<&[ToolDef]>) -> String { + let Some(tools) = existing_tools else { + return base_name.to_string(); + }; + + if !tools.iter().any(|t| t.function.name == base_name) { + return base_name.to_string(); + } + + for i in 1..100 { + let candidate = format!("{}_{}", base_name, i); + if !tools.iter().any(|t| t.function.name == candidate) { + return candidate; + } + } + + // Fallback with process id if somehow we can't find a unique name + format!("{}_{}_fallback", base_name, std::process::id()) +} + +#[derive(Deserialize, Serialize, Clone, Debug)] struct OpenAIFunction { name: String, arguments: String, } -#[derive(Deserialize, Serialize, Clone)] +#[derive(Deserialize, Serialize, Clone, Debug)] struct OpenAIToolCall { id: String, function: OpenAIFunction, @@ -92,20 +113,37 @@ struct OpenAIRequest<'a> { model: &'a str, messages: &'a Vec, #[serde(skip_serializing_if = "Option::is_none")] - tools: Option<&'a Vec<&'a ToolDef>>, + tools: Option<&'a Vec>, #[serde(skip_serializing_if = "Option::is_none")] temperature: Option, #[serde(skip_serializing_if = "Option::is_none")] max_completion_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + response_format: Option, } -#[derive(Serialize, Clone)] +#[derive(Serialize, Clone, Debug)] +struct ResponseFormat { + r#type: String, + json_schema: JsonSchemaFormat, +} + +#[derive(Serialize, Clone, Debug)] +struct JsonSchemaFormat { + name: String, + schema: OpenAPISchema, + #[serde(skip_serializing_if = "Option::is_none")] + strict: Option, +} + +#[derive(Serialize, Clone, Debug)] struct ToolDefFunction { name: String, + description: Option, parameters: Box, } -#[derive(Serialize, Clone)] +#[derive(Serialize, Clone, Debug)] struct ToolDef { r#type: String, function: ToolDefFunction, @@ -123,6 +161,7 @@ struct AIAgentArgs { user_message: String, temperature: Option, max_completion_tokens: Option, + output_schema: Option, } #[derive(Deserialize, Debug)] @@ -163,13 +202,27 @@ impl Provider { #[derive(Serialize)] struct AIAgentResult<'a> { - output: String, + output: Box, messages: Vec>, } -#[derive(Serialize, Default, Clone, Debug)] +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(untagged)] +enum SchemaType { + Single(String), + Multiple(Vec), +} + +impl Default for SchemaType { + fn default() -> Self { + SchemaType::Single("object".to_string()) + } +} + +#[derive(Serialize, Deserialize, Default, Clone, Debug)] struct OpenAPISchema { - r#type: String, + #[serde(skip_serializing_if = "Option::is_none")] + r#type: Option, #[serde(skip_serializing_if = "Option::is_none")] items: Option>, #[serde(skip_serializing_if = "Option::is_none")] @@ -182,20 +235,29 @@ struct OpenAPISchema { format: Option, #[serde(skip_serializing_if = "Option::is_none")] r#enum: Option>, + #[serde( + skip_serializing_if = "Option::is_none", + rename = "additionalProperties" + )] + additional_properties: Option, } impl OpenAPISchema { fn from_str(typ: &str) -> Self { - OpenAPISchema { r#type: typ.to_string(), ..Default::default() } + OpenAPISchema { r#type: Some(SchemaType::Single(typ.to_string())), ..Default::default() } } fn from_str_with_enum(typ: &str, enu: &Option>) -> Self { - OpenAPISchema { r#type: typ.to_string(), r#enum: enu.clone(), ..Default::default() } + OpenAPISchema { + r#type: Some(SchemaType::Single(typ.to_string())), + r#enum: enu.clone(), + ..Default::default() + } } fn datetime() -> Self { Self { - r#type: "string".to_string(), + r#type: Some(SchemaType::Single("string".to_string())), format: Some("date-time".to_string()), ..Default::default() } @@ -214,12 +276,12 @@ impl OpenAPISchema { Typ::Sql => Self::from_str("string"), Typ::DynSelect(_) => Self::from_str("string"), Typ::List(typ) => OpenAPISchema { - r#type: "array".to_string(), + r#type: Some(SchemaType::Single("array".to_string())), items: Some(Box::new(Self::from_typ(typ))), ..Default::default() }, Typ::Object(typ) => OpenAPISchema { - r#type: "object".to_string(), + r#type: Some(SchemaType::Single("object".to_string())), items: None, properties: typ.props.as_ref().map(|props| { props @@ -234,13 +296,13 @@ impl OpenAPISchema { ..Default::default() }, Typ::OneOf(variants) => OpenAPISchema { - r#type: "object".to_string(), + r#type: Some(SchemaType::Single("object".to_string())), one_of: Some( variants .iter() .map(|variant| { let schema = OpenAPISchema { - r#type: "object".to_string(), + r#type: Some(SchemaType::Single("object".to_string())), properties: Some( variant .properties @@ -280,6 +342,79 @@ impl OpenAPISchema { Typ::Unknown => Self::from_str("object"), } } + + /// Makes this schema compatible with OpenAI's strict mode by: + /// - Adding additionalProperties: false to all object types + /// - Making non-required properties nullable + /// - Ensuring all properties are in the required array + fn make_strict(mut self) -> Self { + // Handle this schema if it's an object type + if let Some(SchemaType::Single(ref type_str)) = self.r#type { + if type_str == "object" { + // Set additionalProperties to false + self.additional_properties = Some(false); + + if let Some(properties) = self.properties.as_mut() { + // Get original required fields + let original_required = self.required.as_ref(); + + if let Some(required) = original_required { + // Update properties to make non-required fields nullable + for (key, prop) in properties.iter_mut() { + let mut new_prop = (**prop).clone(); + // Make non-required fields nullable + if !required.contains(key) { + new_prop = new_prop.make_nullable(); + } + // Recursively make nested schemas strict + new_prop = new_prop.make_strict(); + *prop = Box::new(new_prop); + } + } + + // All properties must be in required array for strict mode + self.required = Some(properties.keys().cloned().collect()); + } + } + } + + // Recursively process nested schemas + if let Some(ref mut items) = self.items { + **items = items.as_ref().clone().make_strict(); + } + + if let Some(ref mut one_of) = self.one_of { + *one_of = one_of + .iter() + .map(|schema| Box::new(schema.as_ref().clone().make_strict())) + .collect(); + } + + self + } + + /// Makes this property nullable by converting its type to a union with null + fn make_nullable(mut self) -> Self { + match self.r#type.take() { + Some(SchemaType::Single(type_str)) => { + if type_str != "null" { + self.r#type = Some(SchemaType::Multiple(vec![type_str, "null".into()])); + } else { + self.r#type = Some(SchemaType::Single("null".into())); + } + } + Some(SchemaType::Multiple(mut types)) => { + if !types.iter().any(|t| t == "null") { + types.push("null".into()); + } + self.r#type = Some(SchemaType::Multiple(types)); + } + None => { + self.r#type = Some(SchemaType::Single("null".into())); + } + } + self + } } async fn update_flow_status_module_with_actions( @@ -350,7 +485,7 @@ fn parse_raw_script_schema(content: &str, language: &ScriptLang) -> Result> = if tools.is_empty() { None } else { - Some(tools.iter().map(|t| &t.def).collect()) + Some(tools.iter().map(|t| t.def.clone()).collect()) }; + let has_output_properties = args + .output_schema + .as_ref() + .and_then(|schema| schema.properties.as_ref()) + .map(|props| !props.is_empty()) + .unwrap_or(false); + let is_anthropic = matches!(args.provider, Provider::Anthropic { .. }); + let mut response_format: Option = None; + let mut used_structured_output_tool = false; + let mut structured_output_tool_name: Option = None; + + if has_output_properties { + let schema = args.output_schema.as_ref().unwrap(); // we know it's some because of the check above + if is_anthropic { + // if output schema is provided, and provider is anthropic, add a structured_output tool in the list of tools + let unique_tool_name = find_unique_tool_name("structured_output", tool_defs.as_deref()); + structured_output_tool_name = Some(unique_tool_name.clone()); + + let output_tool = ToolDef { + r#type: "function".to_string(), + function: ToolDefFunction { + name: unique_tool_name, + description: Some( + "This tool MUST be used last to return a structured JSON object as the final output." + .to_string(), + ), + parameters: to_raw_value(&schema), + }, + }; + if let Some(ref mut existing_tools) = tool_defs { + existing_tools.push(output_tool); + } else { + tool_defs = Some(vec![output_tool]); + } + } else { + // if output schema is provided, and provider is openai, add a response_format with json_schema + let strict_schema = schema.clone().make_strict(); + response_format = Some(ResponseFormat { + r#type: "json_schema".to_string(), + json_schema: JsonSchemaFormat { + name: "structured_output".to_string(), + schema: strict_schema, + strict: Some(true), + }, + }); + } + } + for i in 0..MAX_AGENT_ITERATIONS { + if used_structured_output_tool { + break; + } + let response = { let resp = HTTP_CLIENT .post(format!("{}/chat/completions", base_url)) @@ -650,6 +837,11 @@ async fn run_agent( tools: tool_defs.as_ref(), temperature: args.temperature, max_completion_tokens: args.max_completion_tokens, + response_format: if has_output_properties && !is_anthropic { + response_format.clone() + } else { + None + }, }) .send() .await @@ -717,6 +909,28 @@ async fn run_agent( }); for tool_call in tool_calls.iter() { + // Structured output tool is used, we stop here as this will be the final output + if structured_output_tool_name + .as_ref() + .map_or(false, |name| tool_call.function.name == *name) + { + used_structured_output_tool = true; + messages.push(OpenAIMessage { + role: "tool".to_string(), + content: Some("Successfully ran structured_output tool".to_string()), + tool_call_id: Some(tool_call.id.clone()), + ..Default::default() + }); + messages.push(OpenAIMessage { + role: "assistant".to_string(), + content: Some(tool_call.function.arguments.clone()), + agent_action: Some(AgentAction::Message {}), + ..Default::default() + }); + content = Some(tool_call.function.arguments.clone()); + break; + } + let tool = tools .iter() .find(|t| t.def.function.name == tool_call.function.name); @@ -794,8 +1008,19 @@ async fn run_agent( .map(|m| Message { message: m, agent_action: m.agent_action.as_ref() }) .collect(); + // Parse content as JSON, fallback to string if it fails + let output_value = match content { + Some(content_str) => match has_output_properties { + true => serde_json::from_str::>(&content_str).map_err(|e| { + Error::internal_err(format!("Failed to parse structured output: {}", e)) + })?, + false => to_raw_value(&content_str), + }, + None => to_raw_value(&""), + }; + Ok(to_raw_value(&AIAgentResult { - output: content.unwrap_or_default().clone(), + output: output_value, messages: final_messages, })) } @@ -971,6 +1196,7 @@ pub async fn handle_ai_agent_job( r#type: "function".to_string(), function: ToolDefFunction { name: summary.clone(), + description: None, parameters: schema.unwrap_or_else(|| { to_raw_value(&serde_json::json!({ "type": "object", diff --git a/frontend/src/lib/components/ArgInput.svelte b/frontend/src/lib/components/ArgInput.svelte index 7897832188..9801dbe33c 100644 --- a/frontend/src/lib/components/ArgInput.svelte +++ b/frontend/src/lib/components/ArgInput.svelte @@ -677,6 +677,19 @@ editTab="inputEditor" noPreview addPropertyInEditorTab + on:delete={(e) => { + // Handle property deletion + if (value && value.properties && value.properties[e.detail]) { + delete value.properties[e.detail] + // Also remove from order array if it exists + if (value.order) { + value.order = value.order.filter(key => key !== e.detail) + } + // Update the value to trigger reactivity + value = { ...value } + dispatch('change') + } + }} /> {/await} {:else if inputCat == 'object' && format?.startsWith('jsonschema-')} diff --git a/frontend/src/lib/components/EditableSchemaForm.svelte b/frontend/src/lib/components/EditableSchemaForm.svelte index edaa506dc4..57a11397e4 100644 --- a/frontend/src/lib/components/EditableSchemaForm.svelte +++ b/frontend/src/lib/components/EditableSchemaForm.svelte @@ -528,7 +528,12 @@ {#if jsonEnabled && customUi?.jsonOnly != true}
{#if addPropertyInEditorTab} - + { + openField(propertyName) + }} + > {#snippet trigger()}