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 <noreply@anthropic.com> * test: add $schema field to all output schema integration tests Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * fix: remove $schema field from Google AI tool parameter schemas Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * test: add workspace script tool test for AI agents Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
2026-08-26 00:01:37 +00:00 · 2026-02-03 02:04:36 +01:00
parent 9851b1ee28
commit dd7664b764
5 changed files with 354 additions and 18 deletions
@@ -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<String>,
-    pub parameters: Box<RawValue>,
+    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<GeminiFunctionDeclaration> = 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!(
@@ -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"));
+    }
 }
@@ -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


@@ -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": {
@@ -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,