From dd7664b7645a6f3f5e58ede8d460fd7801836a4f Mon Sep 17 00:00:00 2001 From: centdix <40307056+centdix@users.noreply.github.com> Date: Tue, 3 Feb 2026 02:04:36 +0100 Subject: [PATCH] fix: remove $schema field from Google AI output schema requests (#7765) * fix: remove $schema field from Google AI output schema requests Co-Authored-By: Claude Opus 4.5 * test: add $schema field to all output schema integration tests Co-Authored-By: Claude Opus 4.5 * fix: remove $schema field from Google AI tool parameter schemas Co-Authored-By: Claude Opus 4.5 * test: add workspace script tool test for AI agents Co-Authored-By: Claude Opus 4.5 --------- Co-authored-by: Claude Opus 4.5 --- .../src/ai/providers/google_ai.rs | 40 ++-- backend/windmill-worker/src/ai/types.rs | 181 ++++++++++++++++++ integration_tests/ai_agent_tests/conftest.py | 99 ++++++++++ .../ai_agent_tests/test_output_schema.py | 8 + .../ai_agent_tests/test_tool_calling.py | 44 ++++- 5 files changed, 354 insertions(+), 18 deletions(-) diff --git a/backend/windmill-worker/src/ai/providers/google_ai.rs b/backend/windmill-worker/src/ai/providers/google_ai.rs index bfd4f7d70d..62e6afff75 100644 --- a/backend/windmill-worker/src/ai/providers/google_ai.rs +++ b/backend/windmill-worker/src/ai/providers/google_ai.rs @@ -1,6 +1,5 @@ use async_trait::async_trait; use serde::{Deserialize, Serialize}; -use serde_json::value::RawValue; use windmill_common::{client::AuthedClient, error::Error}; use crate::ai::{ @@ -105,7 +104,7 @@ pub struct GeminiFunctionDeclaration { pub name: String, #[serde(skip_serializing_if = "Option::is_none")] pub description: Option, - pub parameters: Box, + pub parameters: OpenAPISchema, } /// Tool configuration for controlling function calling behavior @@ -479,11 +478,17 @@ impl GoogleAIQueryBuilder { if let Some(tool_defs) = tools { let declarations: Vec = tool_defs .iter() - .map(|t| GeminiFunctionDeclaration { - name: t.function.name.clone(), - description: t.function.description.clone(), - // Use parameters directly to avoid round-trip serialization - parameters: t.function.parameters.clone(), + .filter_map(|t| { + // Deserialize RawValue into OpenAPISchema, sanitize, then use + let mut schema: OpenAPISchema = + serde_json::from_str(t.function.parameters.get()).ok()?; + schema.sanitize_for_google(); + + Some(GeminiFunctionDeclaration { + name: t.function.name.clone(), + description: t.function.description.clone(), + parameters: schema, + }) }) .collect(); @@ -522,10 +527,11 @@ impl GoogleAIQueryBuilder { .unwrap_or(false); let (response_mime_type, response_schema) = if has_output_schema { - let schema = args.output_schema.unwrap(); + let mut schema = args.output_schema.unwrap().clone(); + schema.sanitize_for_google(); ( Some("application/json".to_string()), - serde_json::to_value(schema).ok(), + serde_json::to_value(&schema).ok(), ) } else { (None, None) @@ -632,8 +638,13 @@ impl QueryBuilder for GoogleAIQueryBuilder { } // Convert Gemini usage metadata to TokenUsage - let usage = gemini_usage - .map(|u| TokenUsage::new(u.prompt_token_count, u.candidates_token_count, u.total_token_count)); + let usage = gemini_usage.map(|u| { + TokenUsage::new( + u.prompt_token_count, + u.candidates_token_count, + u.total_token_count, + ) + }); Ok(ParsedResponse::Text { content: if accumulated_content.is_empty() { @@ -649,12 +660,7 @@ impl QueryBuilder for GoogleAIQueryBuilder { }) } - fn get_endpoint( - &self, - base_url: &str, - model: &str, - output_type: &OutputType, - ) -> String { + fn get_endpoint(&self, base_url: &str, model: &str, output_type: &OutputType) -> String { match output_type { OutputType::Text => { format!( diff --git a/backend/windmill-worker/src/ai/types.rs b/backend/windmill-worker/src/ai/types.rs index 78432fe79e..0955e7c7a1 100644 --- a/backend/windmill-worker/src/ai/types.rs +++ b/backend/windmill-worker/src/ai/types.rs @@ -761,6 +761,54 @@ impl OpenAPISchema { self.pattern = other.pattern.clone(); } } + + /// Sanitizes this schema for Google AI's API by removing unsupported fields. + /// Google's Gemini API does not accept JSON Schema metadata fields like $schema. + pub fn sanitize_for_google(&mut self) { + // Remove $schema field - not supported by Google AI + self.schema_url = None; + + // Recursively sanitize nested schemas + if let Some(ref mut items) = self.items { + items.sanitize_for_google(); + } + + if let Some(ref mut properties) = self.properties { + for prop in properties.values_mut() { + prop.sanitize_for_google(); + } + } + + if let Some(ref mut one_of) = self.one_of { + for schema in one_of.iter_mut() { + schema.sanitize_for_google(); + } + } + + if let Some(ref mut any_of) = self.any_of { + for schema in any_of.iter_mut() { + schema.sanitize_for_google(); + } + } + + if let Some(ref mut all_of) = self.all_of { + for schema in all_of.iter_mut() { + schema.sanitize_for_google(); + } + } + + if let Some(ref mut defs) = self.defs { + for schema in defs.values_mut() { + schema.sanitize_for_google(); + } + } + + if let Some(ref mut definitions) = self.definitions { + for schema in definitions.values_mut() { + schema.sanitize_for_google(); + } + } + } } /// Wrapper for S3Object with type discriminator for conversation storage @@ -1256,4 +1304,137 @@ mod tests { let defs = schema.defs.as_ref().expect("defs should exist"); assert!(defs.contains_key("MyType"), "Should have 'MyType' def"); } + + // ========== Google AI sanitization tests ========== + + #[test] + fn test_sanitize_for_google_removes_schema_url() { + let mut schema = OpenAPISchema { + r#type: Some(SchemaType::Single("object".to_string())), + schema_url: Some("http://json-schema.org/draft-07/schema#".to_string()), + ..Default::default() + }; + + schema.sanitize_for_google(); + + assert!(schema.schema_url.is_none(), "$schema should be removed"); + // Type should be preserved + assert!(matches!(&schema.r#type, Some(SchemaType::Single(t)) if t == "object")); + } + + #[test] + fn test_sanitize_for_google_recursive_properties() { + let nested = OpenAPISchema { + r#type: Some(SchemaType::Single("string".to_string())), + schema_url: Some("http://json-schema.org/draft-07/schema#".to_string()), + ..Default::default() + }; + + let mut schema = OpenAPISchema { + r#type: Some(SchemaType::Single("object".to_string())), + schema_url: Some("http://json-schema.org/draft-07/schema#".to_string()), + properties: Some( + vec![("field".to_string(), Box::new(nested))] + .into_iter() + .collect(), + ), + ..Default::default() + }; + + schema.sanitize_for_google(); + + assert!(schema.schema_url.is_none(), "Root $schema should be removed"); + let field = schema.properties.as_ref().unwrap().get("field").unwrap(); + assert!(field.schema_url.is_none(), "Nested $schema should be removed"); + } + + #[test] + fn test_sanitize_for_google_recursive_items() { + let item_schema = OpenAPISchema { + r#type: Some(SchemaType::Single("string".to_string())), + schema_url: Some("http://json-schema.org/draft-07/schema#".to_string()), + ..Default::default() + }; + + let mut schema = OpenAPISchema { + r#type: Some(SchemaType::Single("array".to_string())), + items: Some(Box::new(item_schema)), + ..Default::default() + }; + + schema.sanitize_for_google(); + + let items = schema.items.as_ref().unwrap(); + assert!(items.schema_url.is_none(), "Array items $schema should be removed"); + } + + #[test] + fn test_sanitize_for_google_recursive_one_of() { + let variant = OpenAPISchema { + r#type: Some(SchemaType::Single("string".to_string())), + schema_url: Some("http://json-schema.org/draft-07/schema#".to_string()), + ..Default::default() + }; + + let mut schema = OpenAPISchema { + one_of: Some(vec![Box::new(variant)]), + ..Default::default() + }; + + schema.sanitize_for_google(); + + let variant = &schema.one_of.as_ref().unwrap()[0]; + assert!(variant.schema_url.is_none(), "oneOf variant $schema should be removed"); + } + + #[test] + fn test_sanitize_for_google_recursive_defs() { + let def_schema = OpenAPISchema { + r#type: Some(SchemaType::Single("object".to_string())), + schema_url: Some("http://json-schema.org/draft-07/schema#".to_string()), + ..Default::default() + }; + let mut defs = HashMap::new(); + defs.insert("MyType".to_string(), Box::new(def_schema)); + + let mut schema = OpenAPISchema { + defs: Some(defs), + schema_url: Some("http://json-schema.org/draft-07/schema#".to_string()), + ..Default::default() + }; + + schema.sanitize_for_google(); + + assert!(schema.schema_url.is_none(), "Root $schema should be removed"); + let my_type = schema.defs.as_ref().unwrap().get("MyType").unwrap(); + assert!(my_type.schema_url.is_none(), "$defs schema $schema should be removed"); + } + + #[test] + fn test_sanitize_for_google_preserves_other_fields() { + let mut schema = OpenAPISchema { + r#type: Some(SchemaType::Single("object".to_string())), + schema_url: Some("http://json-schema.org/draft-07/schema#".to_string()), + title: Some("Test Schema".to_string()), + description: Some("A test schema".to_string()), + properties: Some( + vec![("name".to_string(), Box::new(string_schema()))] + .into_iter() + .collect(), + ), + required: Some(vec!["name".to_string()]), + ..Default::default() + }; + + schema.sanitize_for_google(); + + // $schema should be removed + assert!(schema.schema_url.is_none()); + // Other fields should be preserved + assert_eq!(schema.title, Some("Test Schema".to_string())); + assert_eq!(schema.description, Some("A test schema".to_string())); + assert!(schema.properties.is_some()); + assert!(schema.required.is_some()); + assert!(matches!(&schema.r#type, Some(SchemaType::Single(t)) if t == "object")); + } } diff --git a/integration_tests/ai_agent_tests/conftest.py b/integration_tests/ai_agent_tests/conftest.py index 187b63d4e5..921ebcb73f 100644 --- a/integration_tests/ai_agent_tests/conftest.py +++ b/integration_tests/ai_agent_tests/conftest.py @@ -161,6 +161,48 @@ class AIAgentTestClient: if "already exists" not in error.lower(): raise Exception(f"Failed to create resource: {error}") + def script_exists(self, path: str) -> bool: + """Check if a script exists at the given path.""" + response = self._client.get( + f"/api/w/{self.workspace}/scripts/exists/p/{path}", + ) + return response.status_code == 200 and response.content.decode() == "true" + + def create_script( + self, + path: str, + content: str, + language: str = "bun", + summary: str = "", + description: str = "", + schema: dict[str, Any] | None = None, + ): + """Create a script in the workspace.""" + # Check if script already exists + if self.script_exists(path): + return + + payload = { + "path": path, + "summary": summary, + "description": description, + "content": content, + "schema": schema or {"type": "object", "properties": {}, "required": []}, + "is_template": False, + "language": language, + "kind": "script", + } + + response = self._client.post( + f"/api/w/{self.workspace}/scripts/create", + json=payload, + ) + if response.status_code // 100 != 2: + error = response.content.decode() + # Ignore if script already exists + if "already exists" not in error.lower(): + raise Exception(f"Failed to create script: {error}") + def upload_s3_file(self, s3_key: str, file_content: bytes, content_type: str = "image/png") -> dict: """Upload a file to S3 storage via Windmill API.""" response = self._client.post( @@ -343,6 +385,36 @@ def create_websearch_tool() -> dict[str, Any]: } +def create_script_tool( + tool_id: str, + script_path: str, + params: list[str], +) -> dict[str, Any]: + """ + Create a tool that references an existing script in the workspace. + + Args: + tool_id: Unique ID for the tool + script_path: Path to the script in the workspace (e.g., "u/admin/sum_script") + params: List of parameter names (each gets type: ai so the agent provides values) + + Returns: + A tool definition dictionary + """ + input_transforms = {param: {"type": "ai"} for param in params} + + return { + "id": tool_id, + "summary": tool_id, + "value": { + "tool_type": "flowmodule", + "type": "script", + "path": script_path, + "input_transforms": input_transforms, + }, + } + + @pytest.fixture(scope="session") def client(): """Create and return an AI agent test client.""" @@ -417,6 +489,33 @@ def setup_providers(client): "name": "deepwiki" }) + # Create sum script for workspace script tool tests + client.create_script( + path="u/admin/sum_script", + content="""export function main(a: number, b: number): number { + return a + b; +} +""", + language="bun", + summary="Sum two numbers", + description="A simple script that adds two numbers together", + schema={ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": { + "a": { + "type": "number", + "description": "First number to add", + }, + "b": { + "type": "number", + "description": "Second number to add", + }, + }, + "required": ["a", "b"], + }, + ) + yield diff --git a/integration_tests/ai_agent_tests/test_output_schema.py b/integration_tests/ai_agent_tests/test_output_schema.py index a44804e725..a5b50cd3ab 100644 --- a/integration_tests/ai_agent_tests/test_output_schema.py +++ b/integration_tests/ai_agent_tests/test_output_schema.py @@ -26,6 +26,7 @@ export function main(a: number, b: number): number { # Output schema for structured result RESULT_SCHEMA = { + "$schema": "https://json-schema.org/draft/2020-12/schema", "type": "object", "properties": { "sum": {"type": "number"} @@ -163,6 +164,7 @@ class TestSchemaVariations: ): """Test deeply nested object structure.""" schema = { + "$schema": "https://json-schema.org/draft/2020-12/schema", "type": "object", "properties": { "user": { @@ -205,6 +207,7 @@ class TestSchemaVariations: ): """Test array with object items.""" schema = { + "$schema": "https://json-schema.org/draft/2020-12/schema", "type": "object", "properties": { "items": { @@ -250,6 +253,7 @@ class TestSchemaVariations: ): """Test enum constraints.""" schema = { + "$schema": "https://json-schema.org/draft/2020-12/schema", "type": "object", "properties": { "status": { @@ -290,6 +294,7 @@ class TestSchemaVariations: ): """Test that optional fields are handled correctly (made nullable).""" schema = { + "$schema": "https://json-schema.org/draft/2020-12/schema", "type": "object", "properties": { "name": {"type": "string"}, @@ -327,6 +332,7 @@ class TestSchemaVariations: ): """Test min/max constraints on numbers.""" schema = { + "$schema": "https://json-schema.org/draft/2020-12/schema", "type": "object", "properties": { "rating": { @@ -372,6 +378,7 @@ class TestSchemaVariations: pytest.xfail("Google AI does not support $ref with definitions in output_schema") schema = { + "$schema": "https://json-schema.org/draft/2020-12/schema", "type": "object", "properties": { "primary": {"$ref": "#/definitions/Color"}, @@ -419,6 +426,7 @@ class TestSchemaVariations: ): """Test anyOf for union types.""" schema = { + "$schema": "https://json-schema.org/draft/2020-12/schema", "type": "object", "properties": { "value": { diff --git a/integration_tests/ai_agent_tests/test_tool_calling.py b/integration_tests/ai_agent_tests/test_tool_calling.py index 4d17edd100..5205a57727 100644 --- a/integration_tests/ai_agent_tests/test_tool_calling.py +++ b/integration_tests/ai_agent_tests/test_tool_calling.py @@ -3,13 +3,14 @@ Tool calling tests for AI agents. Tests AI agent tool calling with different tool types: - Rawscript tools (inline Bun/TypeScript) +- Workspace script tools (scripts deployed to the workspace) - MCP tools (external MCP servers) - Websearch tools (built-in web search) """ import pytest -from .conftest import AIAgentTestClient, create_ai_agent_flow, create_rawscript_tool +from .conftest import AIAgentTestClient, create_ai_agent_flow, create_rawscript_tool, create_script_tool from .providers import ALL_PROVIDERS, ANTHROPIC, GOOGLE_AI, OPENAI @@ -68,6 +69,47 @@ class TestToolCalling: assert "12" in result_str, f"Expected '12' in result: {result}" print(f"Sum tool result from {provider_config['name']}: {result}") + @pytest.mark.parametrize( + "provider_config", + ALL_PROVIDERS, + ids=get_provider_ids(ALL_PROVIDERS), + ) + def test_workspace_script_tool( + self, + client: AIAgentTestClient, + setup_providers, + provider_config, + ): + """ + Test that an AI agent can call a workspace script tool to add numbers. + + This test uses a script that was deployed to the workspace (u/admin/sum_script) + rather than an inline rawscript. + """ + tools = [ + create_script_tool( + tool_id="sum_numbers", + script_path="u/admin/sum_script", + params=["a", "b"], + ) + ] + + flow_value = create_ai_agent_flow( + provider_input_transform=provider_config["input_transform"], + system_prompt="You are a helpful assistant. Use the sum_numbers tool to perform arithmetic.", + tools=tools, + ) + + result = client.run_preview_flow( + flow_value=flow_value, + args={"user_message": "What is 8 + 15? Use the sum_numbers tool."}, + ) + + assert result is not None + result_str = str(result) + assert "23" in result_str, f"Expected '23' in result: {result}" + print(f"Workspace script tool result from {provider_config['name']}: {result}") + @pytest.mark.parametrize( "provider_config", ALL_PROVIDERS,