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

> Discover how oMLX implements flexible tool calling with structured JSON schema output for OpenAI Gemma and XML models. Centralized logic adapts to model families.

- Repository: [Jun Kim/omlx](https://github.com/jundot/omlx)
- Tags: deep-dive
- Published: 2026-05-11

---

**oMLX centralizes tool-calling logic in [`omlx/api/tool_calling.py`](https://github.com/jundot/omlx/blob/main/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`](https://github.com/jundot/omlx/blob/main/omlx/api/tool_calling.py)

All tool-calling logic resides in [`omlx/api/tool_calling.py`](https://github.com/jundot/omlx/blob/main/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`](https://github.com/jundot/omlx/blob/main/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

```python
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

```python
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

```python
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

```python
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

```python
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`](https://github.com/jundot/omlx/blob/main/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.