Repository navigation
Conversation
atdrendel
commented
Sep 17, 2026
atdrendel
commented
- Fix use-after-free crash by backporting implementation merged in ml-explore/mlx#3662
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.
There was a problem hiding this comment.
🟡 Changes recommended
It introduces a major dependency source/version switch (to a fork pinned with exact) that is not clearly justified by the PR description and materially impacts maintainability/supply-chain risk.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
This PR extends the tool-schema and tool-call parsing pipeline to better handle nullable JSON-schema “type” arrays (e.g., ["string","null"]), including normalizing such schemas for Gemma 4 prompt rendering and improving parsing behavior around null/nil/none literals.
Changes:
- Add Gemma 4 tool schema normalization (recursively converts
type: ["T","null"]intotype: "T", nullable: true) and wire it into Gemma 4 processors and LLM prompt rendering via a newToolSchemaGeneratorhook. - Improve tool argument parsing to (a) treat
nullliterals asNSNull()only when the schema permits null, and (b) allow scalar JSON fragment decoding in the Gemma function parser. - Add unit + integration tests covering nullable tool arguments and Gemma 4 tool schema normalization.
File summaries
| File | Description |
|---|---|
| Tests/MLXLMTests/ToolTests.swift | Adds parser test coverage for nullable unions and null literals across multiple tool-call formats. |
| Tests/MLXLMTests/Gemma4ToolSchemaTests.swift | New test suite validating recursive schema normalization and error behavior for unsupported unions. |
| Package.swift | Changes MLX Swift dependency source/version. |
| Libraries/MLXVLM/Models/Gemma4.swift | Normalizes tool schemas before Gemma4/Gemma4Unified apply chat template. |
| Libraries/MLXLMCommon/Tool/Parsers/ParserUtilities.swift | Adds JSON fragment deserialization; improves schema type resolution and null-literal handling. |
| Libraries/MLXLMCommon/Tool/Parsers/GemmaFunctionParser.swift | Allows scalar JSON fragment parsing when schema type is unknown/ambiguous. |
| Libraries/MLXLMCommon/Models/Gemma4.swift | Introduces recursive Gemma 4 tool schema normalizer + normalization error. |
| Libraries/MLXLLM/Models/Gemma4Text.swift | Provides Gemma 4 tool schema generator hook for text model. |
| Libraries/MLXLLM/Models/Gemma4.swift | Adds Gemma 4 tool schema generator and exposes it via the LLMModel hook. |
| Libraries/MLXLLM/LLMModelFactory.swift | Plumbs ToolSchemaGenerator into prompt rendering (applyChatTemplate). |
| Libraries/MLXLLM/LLMModel.swift | Adds toolSchemaGenerator(tokenizer:) + ToolSchemaGenerator protocol and default implementation. |
| Libraries/IntegrationTestHelpers/IntegrationTestHelpers.swift | Adds integration helpers/tests for nullable tool args across models; adds Gemma 4 scenarios. |
| IntegrationTesting/IntegrationTestingTests/ToolCallIntegrationTests.swift | Adds integration tests for nullable tool-argument schemas across multiple model formats. |
Review details
- Files reviewed: 13/13 changed files
- Comments generated: 1
- Review effort level: Lite (auto)
Note
Copilot is running an experiment and ran this review at Lite.
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| ], | ||
| dependencies: [ | ||
| .package(url: "https://github.com/ml-explore/mlx-swift", .upToNextMinor(from: "0.31.4")), | ||
| .package(url: "https://github.com/shareup/mlx-swift", exact: "0.0.2"), |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.