Fix Gemma 4 nullable arguments in tool schemas by atdrendel · Pull Request #11 · shareup/mlx-swift-lm

atdrendel · 2026-07-06T14:20:44Z

This fixes Gemma 4 tool calling when callers provide nullable JSON Schema types using the standard array form, for example:

{ "type": ["string", "null"] }
or
{ "type": ["object", "null"], "properties": {} }

Gemma 4’s chat template expects this as a single string type plus nullable: true, like this:

{ "type": "string", "nullable": true }
or
{ "type": "object", "nullable": true, "properties": {} }

Without that normalization, otherwise valid tool schemas can fail before generation with this error:

Runtime error: upper filter requires string

This pull request fixes this issue for text models by introducing the concept of ToolSchemaGenerator, which I modeled after MessageGenerator. For all non-Gemma 4 models, LLMModel.toolSchemaGenerator(tokenizer:) returns DefaultToolSchemaGenerator, which just returns UserInput.tools. Gemma 4 models return Gemma4ToolSchemaNormalizer.normalizeNullableTypeUnions(input.tools), which normalizes the tool schemas to match Gemma 4's expected format.

For vision models, this pull request uses Gemma4ToolSchemaNormalizer inside of Gemma4Processor.prepare(input:), which leaves all other vision models unchanged.

This pull request also addresses parsing types from tool calls using the array form of specifying types by modifying ParserUtilities.getParameterType() to handle types like ["string", "null"] or ["null"].

This pull request adds integration tests for Gemma 4 to ensure the normalization of tool schemas works for Gemma 4. It also adds integration tests for all of the other supported models to ensure the standard tool schemas are supported by the non-Gemma 4 models.

This pull request also adds the normal suite of integration tests for Gemma 4 12B ("Unified") to ensure the model works as expected.

Copilot

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated no new comments.

This fixes Gemma 4 tool calling when callers provide nullable JSON Schema types using the standard array form, for example: ``` { "type": ["string", "null"] } or { "type": ["object", "null"], "properties": {} } ``` Gemma 4’s chat template expects this as a single string type plus `nullable: true`, like this: ``` { "type": "string", "nullable": true } or { "type": "object", "nullable": true, "properties": {} } ``` Without that normalization, otherwise valid tool schemas can fail before generation with this error: ``` Runtime error: upper filter requires string ``` This pull request fixes this issue for text models by introducing the concept of `ToolSchemaGenerator`, which I modeled after `MessageGenerator`. For all non-Gemma 4 models, `LLMModel.toolSchemaGenerator(tokenizer:)` returns `DefaultToolSchemaGenerator`, which just returns `UserInput.tools`. Gemma 4 models return `Gemma4ToolSchemaNormalizer.normalizeNullableTypeUnions(input.tools)`, which normalizes the tool schemas to match Gemma 4's expected format. For vision models, this pull request uses `Gemma4ToolSchemaNormalizer` inside of `Gemma4Processor.prepare(input:)`, which leaves all other vision models unchanged. This pull request also addresses parsing types from tool calls using the array form of specifying types by modifying `ParserUtilities.getParameterType()` to handle types like `["string", "null"]` or `["null"]`. This pull request adds integration tests for Gemma 4 to ensure the normalization of tool schemas works for Gemma 4. It also adds integration tests for all of the other supported models to ensure the standard tool schemas are supported by the non-Gemma 4 models. This pull request also adds the normal suite of integration tests for Gemma 4 12B ("Unified") to ensure the model works as expected.

atdrendel · 2026-07-07T15:58:16Z

+        try await ToolCallTests.qwen35SupportsJSONArrayNullableToolArguments(container: container)
+    }
+
+    // MARK: - Gemma 4 Unified


Gemma 4 12B, Gemma4Unified, is the only Gemma 4 model I could get to work reliably. The smaller, ExB models all have bugs that are tracked elsewhere.

atdrendel · 2026-07-07T15:58:54Z

        }
    }
-
-    private static func generateWithTools(


I deleted this because I found the two identically-named functions to be confusing.

atdrendel · 2026-07-07T16:01:00Z

+}
+
+/// Generates model-specific tool schemas for prompt rendering.
+public protocol ToolSchemaGenerator: Sendable {


I wasn't really sure where this should be declared. MessageGenerator is declared inside of Chat.swift. I'd be happy to move this to its own file, if you want, but, since it's so small, I figured living inside of LLMModel.swift wasn't a terrible place to put it.

atdrendel · 2026-07-07T16:01:41Z

                    value, paramName: key, funcName: funcName, tools: tools)
            } else if let data = value.data(using: .utf8),
-                let json = deserializeJSON(data)
+                let json = deserializeJSONFragment(data)


This allows "null" to be deserialized to everyone's old friend NSNull.

atdrendel · 2026-07-07T16:02:54Z

+/// Returns the single concrete schema type for a function parameter.
+/// Handles scalar types and nullable single-type arrays like ["integer", "null"].
+/// Returns nil for missing, all-null, or ambiguous union types.
 /// Reference: https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tool_parsers/qwen3_coder.py


I'm not sure if this reference is still needed. I'd be happy to remove it. I don't think the implementation of this function mirrors the one in qwen3_coder.py anymore, especially since it's used for all sorts of different models now.

Copilot AI review requested due to automatic review settings July 6, 2026 14:20

Copilot started reviewing on behalf of atdrendel July 6, 2026 14:21

View session