How oMLX Implements Tool Calling with Structured Output for Different Model Families

oMLX centralizes tool-calling logic in omlx/api/tool_calling.py to automatically adapt to diverse model families—including OpenAI-style JSON, Gemma 4's unique syntax, and various XML formats—while providing a robust JSON-Schema structured output pipeline.

The oMLX library (jundot/omlx) unifies tool calling across heterogeneous language models in the MLX ecosystem. By implementing a layered parsing strategy that queries tokenizer capabilities and cascades through family-specific fallbacks, oMLX handles native markers, regex-based parsers, and structured output validation through a single API surface.

Core Architecture in omlx/api/tool_calling.py

All tool-calling logic resides in omlx/api/tool_calling.py, which orchestrates detection, parsing, and cleanup across ten distinct processing layers. The system interrogates the mlx-lm TokenizerWrapper—exposing attributes like has_tool_calling, tool_call_start, tool_call_end, and tool_parser—to determine the appropriate parsing strategy for the loaded model.

Native Parser Integration

When has_tool_calling evaluates to true, the parse_tool_calls function delegates to the tokenizer's tool_parser. This parser receives matched content between tool_call_start and tool_call_end markers, returning a dictionary with name and arguments keys. The helper _serialize_tool_call_arguments ensures JSON-safe serialization before instantiating ToolCall objects (defined in omlx/api/openai_models.py).

Gemma 4 Non-Standard Format Handling

For Gemma 4 models that emit the non-standard call:name{...} syntax, oMLX employs _parse_gemma4_tool_call_fallback. This function activates when the native parser fails and tool_call_start == "<|tool_call>" is detected. It uses robust regex matching coupled with _gemma4_args_to_json_robust to convert sloppy argument strings into valid JSON objects.

XML and Namespaced Parsing Fallbacks

If native parsing fails or no parser is exposed, the system falls back to _parse_xml_tool_calls, which handles three distinct XML flavors: generic JSON-in-XML, Qwen/Llama <function=name> style tags, and GLM‑4's <arg_key>/<arg_value> structure. Additionally, _parse_namespaced_tool_calls processes MiniMax-style <minimax:tool_call> blocks, extracting <invoke name="func"> elements to build standardized ToolCall objects.

Bracket-Style Legacy Support

Older conversational models emitting [Calling tool: name(args)] or [Tool call: name] patterns are handled by _parse_bracket_tool_calls. This regex-based parser identifies bracketed markers and instantiates the same ToolCall objects used throughout the system, ensuring backward compatibility without special-casing the API.

Structured Output and JSON Schema Pipeline

Beyond tool calling, oMLX implements a complete structured output pipeline for JSON-Schema validation.

Schema-Aware Extraction

The parse_json_output function extracts JSON from model text and validates it against supplied schemas via validate_json_schema. For models lacking native JSON mode support, build_json_system_prompt auto-generates system prompts that instruct the model to emit strict JSON conforming to the specified schema.

Gemma 4 Schema Compatibility

Special helpers enrich_tool_params_for_gemma4 and restore_gemma4_param_names handle property name conflicts specific to the Gemma 4 family, rewriting conflicting schema parameters before validation and restoring original names after parsing.

Real-Time Streaming with ToolCallStreamFilter

During incremental generation, ToolCallStreamFilter monitors token streams for start/end tags, namespaced tags, and bracket prefixes. The feed method silently drops control markup while preserving surrounding prose, ensuring clients never receive raw <tool_call> fragments. The finish method flushes any remaining safe text after generation completes.

Practical Implementation Examples

Basic Tool Call Extraction

from omlx.api.tool_calling import parse_tool_calls

# tokenizer is the mlx-lm TokenizerWrapper attached to the model

output = """
Answer: The weather looks nice.
<tool_call>{"name":"get_weather","arguments":{"city":"Paris"}}</tool_call>
"""

clean_text, tool_calls = parse_tool_calls(output, tokenizer)

print(clean_text)  # "Answer: The weather looks nice."

print(tool_calls[0].function.name)      # "get_weather"

print(tool_calls[0].function.arguments) # '{"city":"Paris"}'

Handling Gemma 4 Syntax

from omlx.api.tool_calling import parse_tool_calls

output = "call:my_func{city: Tokyo, temperature: 20}"
clean, calls = parse_tool_calls(output, tokenizer)

print(calls[0].function.arguments)

# '{"city":"Tokyo","temperature":20}'

Streaming Filter Usage

from omlx.api.tool_calling import ToolCallStreamFilter

stream = ToolCallStreamFilter(tokenizer)

chunks = [
    "Here is the result ",
    "<tool_call>{\"name\":\"search\",\"arguments\":{\"query\":\"MLX\"}}",
    "</tool_call> and that was it."
]

for chunk in chunks:
    safe = stream.feed(chunk)
    if safe:
        print("->", safe, end="")

print(stream.finish())

Structured Output with Schema

from omlx.api.tool_calling import parse_json_output, build_json_system_prompt

schema = {
    "type": "object",
    "properties": {"city": {"type": "string"}, "temp": {"type": "number"}},
    "required": ["city", "temp"],
}

response_format = {"type": "json_schema", "json_schema": {"schema": schema}}
prompt = build_json_system_prompt(response_format)

# After model generation

raw = "The forecast is:\n```json\n{\"city\":\"Paris\",\"temp\":18}\n```"
clean, data, ok, err = parse_json_output(raw, response_format)
print(data)  # {'city': 'Paris', 'temp': 18}

print(ok)    # True

Extracting Tool Calls from Thinking Blocks

from omlx.api.tool_calling import extract_tool_calls_with_thinking

thinking = ""
regular = "Here is the answer."
cleaned_text, tool_calls = extract_tool_calls_with_thinking(
    thinking, regular, tokenizer, tools=my_tool_defs
)

Summary

• oMLX consolidates tool-calling logic in omlx/api/tool_calling.py, providing a single entry point for diverse model families including OpenAI-style JSON, MiniMax XML, Qwen-3 Coder XML, and GLM-4.
• The architecture uses tokenizer attributes (has_tool_calling, tool_parser) to detect native support, falling back to regex-based parsers for Gemma 4's call:name{...} syntax, XML variants, namespaced tags, and bracket-style markers.
• Structured output support includes JSON-Schema validation via parse_json_output and automatic system prompt generation via build_json_system_prompt, with special handling for Gemma 4 schema incompatibilities.
• Streaming applications use ToolCallStreamFilter to suppress control markers in real-time while preserving user-visible text.
• Helper functions like extract_tool_calls_with_thinking and format_tool_call_for_message integrate with reasoning models and standardize API responses.

Frequently Asked Questions

Does oMLX support all MLX models out of the box?

oMLX supports any model exposing the mlx-lm TokenizerWrapper interface with standard tool-calling attributes. For models without native support, the library automatically falls back to regex-based parsing for XML, bracket-style, and other common formats, ensuring broad compatibility across model families.

How does oMLX handle Gemma 4's unique tool-calling syntax?

When the standard parser fails for Gemma 4 models, oMLX invokes _parse_gemma4_tool_call_fallback to process the call:name{...} format. This function uses _gemma4_args_to_json_robust to repair malformed argument strings, converting them into valid JSON before constructing ToolCall objects.

Can I use oMLX for streaming applications?

Yes. The ToolCallStreamFilter class processes incremental token streams, suppressing <tool_call> markup and other control sequences while preserving surrounding text. This ensures that end users receive clean output without raw XML or JSON fragments during generation.

What structured output validation does oMLX provide?

oMLX validates JSON output against supplied JSON Schemas through parse_json_output. The build_json_system_prompt function generates appropriate instructions for models lacking native JSON mode, while family-specific helpers like enrich_tool_params_for_gemma4 handle schema incompatibilities for specific architectures.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}