Add a harmony tool parser for gpt-oss in the server

mlx_lm/chat_templates/harmony.py

@@ -0,0 +1,69 @@
+# Copyright © 2026 Apple Inc.
+
+"""Response format of gpt-oss harmony.
+
+A turn is a header plus a channel body. The generation prompt already ends at
+``<|start|>assistant``, so parsing starts in the header. Header text is
+dropped.
+
+    <|channel|>analysis<|message|>reasoning<|end|>
+    <|channel|>final<|message|>answer<|return|>
+    <|channel|>commentary<|message|>preamble<|end|>
+    to=functions.<name> ... <|message|>{json}<|call|>
+
+``to=functions.`` is the tool marker for both header orders: the recipient
+may come before ``<|channel|>commentary`` or after it. ``<|call|>`` is also
+an end-of-turn token, so a finished call is closed by ``finish`` when that
+token is not in the text.
+"""
+
+RESPONSE_FORMAT = {
+    "initial_state": "header",
+    "transitions": {
+        "header": [
+            ("<|channel|>analysis<|message|>", "reasoning"),
+            ("<|channel|>final<|message|>", "normal"),
+            ("<|channel|>commentary<|message|>", "normal"),
+            ("to=functions.", "tool"),
+        ],
+        "reasoning": [
+            ("<|end|>", "header"),
+            ("<|return|>", "stop"),
+        ],
+        "normal": [
+            ("<|end|>", "header"),
+            ("<|return|>", "stop"),
+        ],
+        "tool": [
+            ("<|call|>", "header"),
+            ("<|end|>", "header"),
+            ("<|return|>", "stop"),
+        ],
+    },
+}
+
+
+def prepare_messages(messages: list) -> None:
+    """Rewrite assistant turns into the shape the harmony template accepts.
+
+    The template reads analysis from ``thinking`` and rejects a tool-call
+    turn that sets both ``content`` and ``thinking``. Clients that follow
+    the chat completion replay the ``reasoning`` field instead.
+    """
+    for message in messages:
+        if message.get("role") != "assistant":
+            continue
+        if not message.get("thinking"):
+            for key in ("reasoning", "reasoning_content", "reasoning_text"):
+                value = message.get(key)
+                if isinstance(value, str) and value:
+                    message["thinking"] = value
+                    break
+        content = message.get("content")
+        if (
+            message.get("tool_calls")
+            and message.get("thinking")
+            and isinstance(content, str)
+            and content
+        ):
+            message["content"] = ""

mlx_lm/generate.py

@@ -553,7 +553,7 @@ def speculative_generate_step(
     if not can_trim_prompt_cache(model_cache):
         types = {type(c).__name__ for c in model_cache if not c.is_trimmable()}
         raise ValueError(
-            f"Speculative decoding requires a trimmable prompt cache " f"(got {types})."
+            f"Speculative decoding requires a trimmable prompt cache (got {types})."
         )
 
     sampler = sampler or greedy_sampler
@@ -1127,10 +1127,13 @@ def from_tokenizer(cls, tokenizer: TokenizerWrapper) -> "TextStateMachine":
                 )
             if tokenizer.has_thinking:
                 transitions["reasoning"].append((tokenizer.tool_call_start, "tool"))
-            # Back to the state the tool call started in.
-            transitions["tool"] = (
-                [(tokenizer.tool_call_end, None)] if tokenizer.tool_call_end else []
-            )
+            # A response format with its own tool edges names the endings
+            # (harmony closes on <|call|> and <|end|>). Otherwise one end
+            # marker returns to the state the call started in.
+            if "tool" not in tokenizer.response_transitions:
+                transitions["tool"] = (
+                    [(tokenizer.tool_call_end, None)] if tokenizer.tool_call_end else []
+                )
 
         if tokenizer.structural_markers:
             for w in tokenizer.structural_markers:

mlx_lm/models/gpt_oss.py

@@ -14,6 +14,13 @@
 from .rope_utils import initialize_rope
 from .switch_layers import SwitchGLU
 
+# The checkpoint already carries the harmony chat template. These select the
+# response format and the tool parser for checkpoints that do not name them.
+TOKENIZER_DEFAULTS = {
+    "chat_template_type": "harmony",
+    "tool_parser_type": "harmony",
+}
+
 
 @dataclass
 class ModelArgs(BaseModelArgs):

mlx_lm/tokenizer_utils.py

@@ -399,6 +399,7 @@ def __init__(
             or has_custom_renderer
         )
         self._tool_parser = tool_parser
+        self._prepare_messages = None
         self._tool_call_start = tool_call_start
         self._tool_call_end = tool_call_end
         self._tool_call_start_tokens = None
@@ -417,6 +418,9 @@ def apply_chat_template(self, *args, tokenize=True, **kwargs):
         if self._thinking_kwarg not in kwargs:
             kwargs[self._thinking_kwarg] = self.has_thinking
 
+        if self._prepare_messages is not None and args and isinstance(args[0], list):
+            self._prepare_messages(args[0])
+
         if self._chat_template is not None:
             out = self._chat_template(*args, **kwargs)
             if tokenize:
@@ -683,6 +687,11 @@ def _infer_tool_parser(tokenizer):
             return "muse_spark"
         elif "<tool_call>" in chat_template and "tool_call.name" in chat_template:
             return "json_tools"
+        elif (
+            "Valid channels: analysis, commentary, final" in chat_template
+            and "<|channel|>" in tokenizer.get_vocab()
+        ):
+            return "harmony"
 
     # No template match, so fall back to the vocab. LFM2.5 conversions for
     # example have the tool-call tokens but do not name them in the template.
@@ -738,6 +747,7 @@ def load(
 
     chat_template = None
     response_format = None
+    prepare_messages = None
 
     tokenizer = AutoTokenizer.from_pretrained(
         model_path, **(tokenizer_config_extra or {})
@@ -748,18 +758,23 @@ def load(
     chat_template_type = tokenizer_config.get(
         "chat_template_type", defaults.get("chat_template_type")
     )
-    if chat_template_type:
-        module = importlib.import_module(f"mlx_lm.chat_templates.{chat_template_type}")
-        # A template in the checkpoint replaces only the default template.
-        if "chat_template_type" in tokenizer_config or tokenizer.chat_template is None:
-            chat_template = module.apply_chat_template
-        response_format = getattr(module, "RESPONSE_FORMAT", None)
-
     tool_parser_type = tokenizer_config.get(
         "tool_parser_type", defaults.get("tool_parser_type")
     )
     if tool_parser_type is None:
         tool_parser_type = _infer_tool_parser(tokenizer)
+    # Harmony's response format and history rewrite live next to the parser.
+    # The checkpoint keeps its own chat template.
+    if chat_template_type is None and tool_parser_type == "harmony":
+        chat_template_type = "harmony"
+
+    if chat_template_type:
+        module = importlib.import_module(f"mlx_lm.chat_templates.{chat_template_type}")
+        # A template in the checkpoint replaces only the default template.
+        if "chat_template_type" in tokenizer_config or tokenizer.chat_template is None:
+            chat_template = getattr(module, "apply_chat_template", None)
+        response_format = getattr(module, "RESPONSE_FORMAT", None)
+        prepare_messages = getattr(module, "prepare_messages", None)
 
     if tool_parser_type is not None:
         tool_module = importlib.import_module(f"mlx_lm.tool_parsers.{tool_parser_type}")
@@ -772,7 +787,16 @@ def load(
         tool_call_start = None
         tool_call_end = None
 
-    return TokenizerWrapper(
+    # <|call|> ends a tool turn. <|return|> is already the tokenizer eos.
+    # <|end|> only ends a channel, so generation continues into the next one.
+    if tool_parser_type == "harmony":
+        call_id = tokenizer.convert_tokens_to_ids("<|call|>")
+        if isinstance(call_id, int) and call_id != tokenizer.unk_token_id:
+            eos_token_ids = list(eos_token_ids or [])
+            if call_id not in eos_token_ids:
+                eos_token_ids.append(call_id)
+
+    wrapper = TokenizerWrapper(
         tokenizer,
         detokenizer_class,
         eos_token_ids=eos_token_ids,
@@ -782,6 +806,8 @@ def load(
         tool_call_end=tool_call_end,
         response_format=response_format,
     )
+    wrapper._prepare_messages = prepare_messages
+    return wrapper
 
 
 def no_bos_or_eos(sequence: List, bos: int, eos: int) -> List:

mlx_lm/tool_parsers/harmony.py

@@ -0,0 +1,47 @@
+# Copyright © 2026 Apple Inc.
+
+"""Tool calls in gpt-oss harmony messages.
+
+The recipient is ``functions.<name>``. The checkpoint renders it before the
+channel (``to=functions.<name><|channel|>commentary json<|message|>``). The
+model also emits the channel first
+(``<|channel|>commentary to=functions.<name> <|constrain|>json<|message|>``).
+The response format starts a call at ``to=functions.``, so this parser sees
+the name and the JSON payload.
+"""
+
+import json
+import re
+from typing import Any
+
+tool_call_start = "to=functions."
+tool_call_end = "<|call|>"
+
+# Name, then whichever of the channel / constrain / content-type the render
+# left in front of the payload.
+_call_regex = re.compile(
+    r"^(?P<name>[A-Za-z_][A-Za-z0-9_]*)"
+    r"(?:<\|channel\|>(?:commentary|analysis))?"
+    r"(?:\s*<\|constrain\|>)?"
+    # ``json`` may sit against <|constrain|> or after a space, as the template writes it.
+    r"(?:\s*[A-Za-z][A-Za-z0-9_]*)?" r"\s*<\|message\|>\s*(?P<args>.*)$",
+    re.DOTALL,
+)
+
+
+def _arguments(text: str) -> dict[str, Any]:
+    text = text.strip()
+    try:
+        value = json.loads(text)
+    except json.JSONDecodeError:
+        value, _ = json.JSONDecoder().raw_decode(text)
+    if not isinstance(value, dict):
+        raise ValueError("Tool call arguments must be a JSON object.")
+    return value
+
+
+def parse_tool_call(text: str, tools: Any | None = None) -> dict[str, Any]:
+    match = _call_regex.match(text.strip())
+    if match is None:
+        raise ValueError("No tool call found.")
+    return {"name": match.group("name"), "arguments": _arguments(match.group("args"))}

tests/test_tool_parsing.py

@@ -6,6 +6,7 @@
     function_gemma,
     gemma4,
     glm47,
+    harmony,
     json_tools,
     kimi_k2,
     kimi_k3,
@@ -69,6 +70,14 @@ def test_parsers(self):
                 '<atem:invoke name="multiply">\n<atem:parameter name="a">12234585</atem:parameter><atem:parameter name="b">48838483920</atem:parameter>\n</atem:invoke>',
                 muse_spark,
             ),
+            (
+                'multiply<|channel|>commentary json<|message|>{"a": 12234585, "b": 48838483920}',
+                harmony,
+            ),
+            (
+                'multiply <|constrain|>json<|message|>{"a": 12234585, "b": 48838483920}',
+                harmony,
+            ),
         ]
 
         tools = [
@@ -204,7 +213,7 @@ def test_glm47_string_typed_args(self):
 
     def test_pythonic_single_quoted_args_with_commas(self):
         # LFM2.5 emits single-quoted strings; embedded commas must not truncate
-        test_case = "[write(filePath='/tmp/hello.py', " "content='# Hello, world!')]"
+        test_case = "[write(filePath='/tmp/hello.py', content='# Hello, world!')]"
         tool_call = pythonic.parse_tool_call(test_case, None)
         self.assertEqual(tool_call["name"], "write")
         self.assertEqual(tool_call["arguments"]["filePath"], "/tmp/hello.py")
@@ -675,7 +684,7 @@ def test_minicpm5(self):
             ('"f"><param name="city">None</param>', {"city": "None"}),
             # History renders dicts in Python style.
             (
-                '"f"><param name="opts">' "{'unit': 'c'}</param>",
+                "\"f\"><param name=\"opts\">{'unit': 'c'}</param>",
                 {"opts": {"unit": "c"}},
             ),
         ]
@@ -900,15 +909,74 @@ def test_qwen3_coder_missing_parameter_tag_close(self):
 
         # A ">" later in the value must not be mistaken for the name's
         test_case = (
-            "<function=read_file>\n"
-            "<parameter=path\n"
-            "a>b\n"
-            "</parameter>\n"
-            "</function>"
+            "<function=read_file>\n<parameter=path\na>b\n</parameter>\n</function>"
         )
         tool_call = qwen3_coder.parse_tool_call(test_case, tools)
         self.assertEqual(tool_call["arguments"], {"path": "a>b"})
 
+    def test_harmony_channels(self):
+        # Channel-first (what gpt-oss emits) and recipient-first (what the
+        # chat template writes back). The call text is what the state machine
+        # keeps after ``to=functions.``.
+        from types import SimpleNamespace
+
+        from mlx_lm.chat_templates.harmony import RESPONSE_FORMAT, prepare_messages
+        from mlx_lm.generate import TextStateMachine
+
+        tokenizer = SimpleNamespace(
+            has_thinking=False,
+            has_tool_calling=True,
+            tool_call_start=harmony.tool_call_start,
+            tool_call_end=harmony.tool_call_end,
+            structural_markers=(),
+            response_transitions=RESPONSE_FORMAT["transitions"],
+        )
+
+        def split(text):
+            reader = TextStateMachine.from_tokenizer(tokenizer).reader(
+                state=RESPONSE_FORMAT["initial_state"]
+            )
+            return reader.feed(text) + reader.finish()
+
+        delta = split(
+            "<|channel|>analysis<|message|>Need the file.<|end|>"
+            "<|start|>assistant<|channel|>commentary to=functions.read"
+            '<|constrain|>json<|message|>{"path": "a"}<|call|>'
+            "<|start|>assistant<|channel|>final<|message|>Done.<|return|>"
+        )
+        self.assertEqual(delta.reasoning, "Need the file.")
+        self.assertEqual(delta.content, "Done.")
+        self.assertEqual(
+            harmony.parse_tool_call(delta.tool_calls[0]),
+            {"name": "read", "arguments": {"path": "a"}},
+        )
+
+        delta = split(
+            "<|channel|>analysis<|message|>Run it.<|end|>"
+            "<|start|>assistant to=functions.bash<|channel|>commentary "
+            'json<|message|>{"command": "echo TOOL_OK"}<|call|>'
+        )
+        self.assertEqual(delta.reasoning, "Run it.")
+        self.assertEqual(delta.content, "")
+        self.assertEqual(
+            harmony.parse_tool_call(delta.tool_calls[0]),
+            {"name": "bash", "arguments": {"command": "echo TOOL_OK"}},
+        )
+
+        messages = [
+            {
+                "role": "assistant",
+                "content": "preamble",
+                "reasoning": "Need the file.",
+                "tool_calls": [
+                    {"function": {"name": "read", "arguments": {"path": "a"}}}
+                ],
+            }
+        ]
+        prepare_messages(messages)
+        self.assertEqual(messages[0]["thinking"], "Need the file.")
+        self.assertEqual(messages[0]["content"], "")
+
 
 if __name__ == "__main__":
     unittest.main()