Skip to content

Remove gemma tool-call post-processor once mlx-swift-lm PR #215 ships #76

Description

@clearloop

crates/mlx/src/tool_parser.rs is a workaround for two bugs in mlx-swift-lm 3.31.3 that together break gemma-3/4 tool calling end-to-end:

  • ToolCallFormat.infer matches model_type == "gemma" exactly, so gemma3 / gemma3n / gemma4 fall through to the default .json parser. (Libraries/MLXLMCommon/Tool/ToolCallFormat.swift:174)
  • GemmaFunctionParser hard-codes <start_function_call> / <end_function_call> as start/end tags, but gemma-3/4 chat templates emit <|tool_call>...<tool_call|> with body call:NAME{key:value,key:<|"|>str<|"|>,...}. (Libraries/MLXLMCommon/Tool/Parsers/GemmaFunctionParser.swift:8-9)

Result without the workaround: the model emits a perfectly valid tool call, mlx-swift-lm's processor doesn't recognize the tags, and the call leaks out as raw text in the assistant content with tool_calls: null. Verified against gemma-4-it-e2b-4bit.

Our workaround in crates/mlx/src/tool_parser.rs post-processes the raw text in both MlxProvider::chat_completion and MlxProvider::chat_completion_stream (crates/mlx/src/provider.rs) when request.tools is non-empty and mlx-swift-lm returned no structured calls. The streaming path buffers all chunks for tool-using requests and runs the parser at end-of-stream, so we lose token-level streaming for those requests but produce correct tool_calls.

Upstream fix: ml-explore/mlx-swift-lm#215. Status as of 2026-04-29: open, mergeable_state: blocked, CHANGES_REQUESTED from the maintainer on 2026-04-24 — review feedback covers Gemma 3 streaming regression risk in the new startTag and an EOS fallback mis-detection.

When upstream merges and a release ships:

  • bump mlx/Package.swift past 3.31.3
  • delete crates/mlx/src/tool_parser.rs, the mod tool_parser; declaration in crates/mlx/src/lib.rs, and the post-processing branches in crates/mlx/src/provider.rs::chat_completion and chat_completion_stream
  • restore unconditional token-level streaming for tool-using requests (drop the buffer_for_tools branch)
  • run cargo run -p crabllm-mlx --example tool_capability to confirm gemma-4-it-e2b-4bit still produces structured tool_calls without our parser

Sibling upstream issue: ml-explore/mlx-swift-lm#215 (issue).

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions