← back to Exo
begins
86735ece782877d8976ebcc8b42f632253c903e9 · 2026-02-16 12:20:58 +0000 · Evan
Files touched
M src/exo/master/adapters/responses.pyM src/exo/worker/engines/mlx/generator/generate.pyM src/exo/worker/engines/mlx/utils_mlx.pyM src/exo/worker/runner/runner.pyA src/exo/worker/runner/tool_parsers.pyM src/exo/worker/tests/unittests/test_runner/test_parse_tool_calls.py
Diff
commit 86735ece782877d8976ebcc8b42f632253c903e9
Author: Evan <evanev7@gmail.com>
Date: Mon Feb 16 12:20:58 2026 +0000
begins
begins
---
src/exo/master/adapters/responses.py | 4 +-
src/exo/worker/engines/mlx/generator/generate.py | 2 +-
src/exo/worker/engines/mlx/utils_mlx.py | 46 +++-
src/exo/worker/runner/runner.py | 291 ++++-----------------
src/exo/worker/runner/tool_parsers.py | 72 +++++
.../unittests/test_runner/test_parse_tool_calls.py | 16 +-
6 files changed, 178 insertions(+), 253 deletions(-)
diff --git a/src/exo/master/adapters/responses.py b/src/exo/master/adapters/responses.py
index b37b7d54..5e059fee 100644
--- a/src/exo/master/adapters/responses.py
+++ b/src/exo/master/adapters/responses.py
@@ -144,8 +144,8 @@ async def collect_responses_response(
for tool in chunk.tool_calls:
function_call_items.append(
ResponseFunctionCallItem(
- id=f"fc_{tool.id}",
- call_id=f"call_{tool.id}",
+ id=tool.id,
+ call_id=tool.id,
name=tool.name,
arguments=tool.arguments,
)
diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index bc38c61c..ffa2f5c0 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -306,7 +306,7 @@ def mlx_generate(
max_stop_len = max((len(s) for s in stop_sequences), default=0)
mx_barrier(group)
- logger.info("Ready to prefill")
+ logger.info("Starting prefill")
# Prefill cache with all tokens except the last one
prefill_tps, prefill_tokens, ssm_snapshots_list = prefill(
diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index 6aceb53c..3ed65ecc 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -353,7 +353,13 @@ def load_tokenizer_for_model_id(
return list(hf_tokenizer.model.encode(text, allowed_special="all")) # pyright: ignore[reportUnknownMemberType,reportUnknownArgumentType]
hf_tokenizer.encode = _patched_encode
- return TokenizerWrapper(hf_tokenizer, eos_token_ids=eos_token_ids)
+ return TokenizerWrapper(
+ hf_tokenizer,
+ eos_token_ids=eos_token_ids,
+ tool_call_start="<|tool_calls_section_begin|>",
+ tool_call_end="<|tool_calls_section_end|>",
+ tool_parser=_parse_kimi_tool_calls,
+ )
tokenizer = load_tokenizer(
model_path,
@@ -585,3 +591,41 @@ def mx_barrier(group: Group | None):
mx.array(1.0), group=group, stream=mx.default_stream(mx.Device(mx.cpu))
)
)
+
+
+def _parse_kimi_tool_calls(text: str):
+ import regex as re
+
+ # kimi has a fixed function naming scheme, with a json formatted arg
+ # functions.multiply:0<|tool_call_argument_begin|>{"a": 2, "b": 3}
+ _func_name_regex = re.compile(
+ r"^\s*((?:functions\.)?(.+?):\d+)\s*<\|tool_call_argument_begin\|>", re.DOTALL
+ )
+ _func_arg_regex = re.compile(r"<\|tool_call_argument_begin\|>\s*(.*)\s*", re.DOTALL)
+ _tool_call_split_regex = re.compile(
+ r"<\|tool_call_begin\|>(.*?)<\|tool_call_end\|>", re.DOTALL
+ )
+
+ def _parse_single_tool(text: str) -> dict[str, Any]:
+ func_name_match = _func_name_regex.search(text)
+ if func_name_match is None:
+ raise ValueError("No tool call found.")
+ tool_call_id = func_name_match.group(1) # e.g. "functions.get_weather:0"
+ func_name = func_name_match.group(2) # e.g. "get_weather"
+
+ func_args_match = _func_arg_regex.search(text)
+ if func_args_match is None:
+ raise ValueError("No tool call arguments found.")
+ func_args = func_args_match.group(1)
+ try:
+ arg_dct = json.loads(func_args) # pyright: ignore[reportAny]
+ except Exception:
+ arg_dct = None
+
+ return dict(id=tool_call_id, name=func_name, arguments=arg_dct)
+
+ tool_matches = _tool_call_split_regex.findall(text)
+ if tool_matches:
+ return [_parse_single_tool(match) for match in tool_matches] # pyright: ignore[reportAny]
+ else:
+ return [_parse_single_tool(text)]
diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py
index ad6c78f6..e55456d3 100644
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -1,11 +1,10 @@
import base64
-import json
import math
import resource
import time
from collections.abc import Generator
from functools import cache
-from typing import Any, Callable, Literal
+from typing import Literal
import mlx.core as mx
from mlx_lm.models.gpt_oss import Model as GptOssModel
@@ -16,7 +15,6 @@ from openai_harmony import ( # pyright: ignore[reportMissingTypeStubs]
StreamableParser,
load_harmony_encoding,
)
-from pydantic import ValidationError
from exo.shared.constants import EXO_MAX_CHUNK_SIZE, EXO_TRACING_ENABLED
from exo.shared.models.model_cards import ModelId, ModelTask
@@ -93,6 +91,8 @@ from exo.worker.engines.mlx.utils_mlx import (
)
from exo.worker.runner.bootstrap import logger
+from .tool_parsers import ToolParser, make_mlx_parser
+
def _is_primary_output_node(shard_metadata: ShardMetadata) -> bool:
"""Check if this node is the primary output node for image generation.
@@ -138,6 +138,7 @@ def main(
inference_model: Model | None = None
image_model: DistributedImageModel | None = None
tokenizer = None
+ tool_parser: ToolParser | None = None
group = None
kv_prefix_cache: KVPrefixCache | None = None
check_for_cancel_every: int | None = None
@@ -203,8 +204,17 @@ def main(
bound_instance, group, on_timeout=on_model_load_timeout
)
logger.info(
- f"model has_tool_calling={tokenizer.has_tool_calling}"
+ f"model has_tool_calling={tokenizer.has_tool_calling} using tokens {tokenizer.tool_call_start}, {tokenizer.tool_call_end}"
)
+ if tokenizer.has_tool_calling:
+ assert tokenizer.tool_call_start
+ assert tokenizer.tool_call_end
+ assert tokenizer.tool_parser # pyright: ignore[reportAny]
+ tool_parser = make_mlx_parser(
+ tokenizer.tool_call_start,
+ tokenizer.tool_call_end,
+ tokenizer.tool_parser, # pyright: ignore[reportAny]
+ )
kv_prefix_cache = KVPrefixCache(group)
elif (
@@ -310,31 +320,11 @@ def main(
mlx_generator, tokenizer
)
- # Kimi-K2 has tool call sections - we don't care about them
- if "kimi" in shard_metadata.model_card.model_id.lower():
- mlx_generator = filter_kimi_tokens(mlx_generator)
- patch_kimi_tokenizer(tokenizer)
-
- # GLM models need patched parser (upstream has bug with None regex match)
- elif "glm" in shard_metadata.model_card.model_id.lower():
- patch_glm_tokenizer(tokenizer)
-
# GPT-OSS specific parsing to match other model formats.
- elif isinstance(inference_model, GptOssModel):
+ if isinstance(inference_model, GptOssModel):
mlx_generator = parse_gpt_oss(mlx_generator)
-
- if tokenizer.has_tool_calling and not isinstance(
- inference_model, GptOssModel
- ):
- assert tokenizer.tool_call_start
- assert tokenizer.tool_call_end
- assert tokenizer.tool_parser # pyright: ignore[reportAny]
- mlx_generator = parse_tool_calls(
- mlx_generator,
- tokenizer.tool_call_start,
- tokenizer.tool_call_end,
- tokenizer.tool_parser, # pyright: ignore[reportAny]
- )
+ elif tool_parser:
+ mlx_generator = parse_tool_calls(mlx_generator, tool_parser)
completion_tokens = 0
tokens_since_last_cancel_check = 0
@@ -587,21 +577,8 @@ def get_gpt_oss_encoding():
return encoding
-def filter_kimi_tokens(
- responses: Generator[GenerationResponse | ToolCallResponse],
-) -> Generator[GenerationResponse]:
- for resp in responses:
- assert isinstance(resp, GenerationResponse)
- if (
- resp.text == "<|tool_calls_section_begin|>"
- or resp.text == "<|tool_calls_section_end|>"
- ):
- continue
- yield resp
-
-
def parse_gpt_oss(
- responses: Generator[GenerationResponse | ToolCallResponse],
+ responses: Generator[GenerationResponse],
) -> Generator[GenerationResponse | ToolCallResponse]:
encoding = get_gpt_oss_encoding()
stream = StreamableParser(encoding, role=Role.ASSISTANT)
@@ -658,9 +635,9 @@ def parse_gpt_oss(
def parse_thinking_models(
- responses: Generator[GenerationResponse | ToolCallResponse],
+ responses: Generator[GenerationResponse],
tokenizer: TokenizerWrapper,
-) -> Generator[GenerationResponse | ToolCallResponse]:
+) -> Generator[GenerationResponse]:
"""
For models that inject thinking tags in the prompt (like GLM-4.7),
prepend the thinking tag to the output stream so the frontend
@@ -781,221 +758,55 @@ def _process_image_response(
def parse_tool_calls(
- responses: Generator[GenerationResponse | ToolCallResponse],
- tool_call_start: str,
- tool_call_end: str,
- tool_parser: Callable[[str], dict[str, Any] | list[dict[str, Any]]],
+ responses: Generator[GenerationResponse], tool_parser: ToolParser
) -> Generator[GenerationResponse | ToolCallResponse]:
in_tool_call = False
tool_call_text_parts: list[str] = []
for response in responses:
- assert isinstance(response, GenerationResponse)
- # assumption: the tool call start is one token
- if response.text == tool_call_start:
+ if response.text.startswith(tool_parser.start_parsing):
in_tool_call = True
- continue
- # assumption: the tool call end is one token
- if in_tool_call and response.text == tool_call_end:
- try:
- # tool_parser returns an arbitrarily nested python dictionary
- # we actually don't want the python dictionary, we just want to
- # parse the top level { function: ..., arguments: ... } structure
- # as we're just gonna hand it back to the api anyway
- parsed = tool_parser("".join(tool_call_text_parts).strip())
- logger.info(f"parsed {tool_call_text_parts=} into {parsed=}")
- if isinstance(parsed, list):
- tools = [_validate_single_tool(tool) for tool in parsed]
- else:
- tools = [_validate_single_tool(parsed)]
- yield ToolCallResponse(
- tool_calls=tools, usage=response.usage, stats=response.stats
- )
- except (
- json.JSONDecodeError,
- ValidationError,
- ValueError,
- AttributeError,
- ) as e:
- # ValueError: our parsers raise this for malformed tool calls
- # AttributeError: upstream parsers (e.g. glm47) may raise this when regex doesn't match
- logger.opt(exception=e).warning("tool call parsing failed")
- # assumption: talking about tool calls, not making a tool call
- response.text = (
- tool_call_start + "".join(tool_call_text_parts) + tool_call_end
+ if in_tool_call:
+ tool_call_text_parts.append(response.text)
+ if response.text.endswith(tool_parser.end_parsing):
+ # parse the actual tool calls from the tool call text
+ parsed = tool_parser.parse_tool_calls(
+ "".join(tool_call_text_parts).strip()
)
- yield response
+ logger.info(f"parsed {tool_call_text_parts=} into {parsed=}")
+ if parsed is not None:
+ yield ToolCallResponse(
+ tool_calls=parsed, usage=response.usage, stats=response.stats
+ )
+ else:
+ logger.warning(
+ f"tool call parsing failed for text {''.join(tool_call_text_parts)}"
+ )
+ response.text = "".join(tool_call_text_parts)
+ yield response
- in_tool_call = False
- tool_call_text_parts = []
- continue
+ in_tool_call = False
+ tool_call_text_parts = []
+ continue
- if in_tool_call:
- tool_call_text_parts.append(response.text)
if response.finish_reason is not None:
logger.info(
- "toll call parsing interrupted, yield partial tool call as text"
+ "tool call parsing interrupted, yield partial tool call as text"
)
- yield GenerationResponse(
- text=tool_call_start + "".join(tool_call_text_parts),
- token=0,
- finish_reason=response.finish_reason,
- usage=response.usage,
- stats=response.stats,
+ response = response.model_copy(
+ update={
+ "text": "".join(tool_call_text_parts),
+ "token": 0,
+ }
)
+ yield response
+
continue
+
# fallthrough
yield response
-def patch_kimi_tokenizer(tokenizer: TokenizerWrapper):
- """
- Version of to-be-upstreamed kimi-k2 tool parser
- """
- import ast
- import json
- from typing import Any
-
- import regex as re
-
- # kimi has a fixed function naming scheme, with a json formatted arg
- # functions.multiply:0 <|tool_call_argument_begin|> {"a": 2, "b": 3}
- # Also needs to handle tools like call_0<|tool_call_argument_begin|>{"filePath": "..."}
- _func_name_regex = re.compile(
- r"^\s*(.+)[:](\d+)\s*<\|tool_call_argument_begin\|>", re.DOTALL
- )
- _func_arg_regex = re.compile(r"<\|tool_call_argument_begin\|>\s*(.*)\s*", re.DOTALL)
-
- # kimi has a tool_calls_section - we're leaving this up to the caller to handle
- tool_call_start = "<|tool_call_begin|>"
- tool_call_end = "<|tool_call_end|>"
-
- def _deserialize(value: str) -> Any: # pyright: ignore[reportAny]
- try:
- return json.loads(value) # pyright: ignore[reportAny]
- except Exception:
- pass
-
- try:
- return ast.literal_eval(value) # pyright: ignore[reportAny]
- except Exception:
- pass
- return value
-
- def parse_tool_call(text: str, tools: Any | None = None):
- func_name_match = _func_name_regex.search(text)
- if func_name_match is None:
- raise ValueError(f"Could not parse function name from tool call: {text!r}")
- original_func_name = func_name_match.group(1)
- tool_id = func_name_match.group(2)
- # strip off the `functions.` prefix, if it exists.
- func_name = original_func_name[original_func_name.find(".") + 1 :]
-
- func_args_match = _func_arg_regex.search(text)
- if func_args_match is None:
- raise ValueError(f"Could not parse function args from tool call: {text!r}")
- func_args = func_args_match.group(1)
- # the args should be valid json - no need to check against our tools to deserialize
- arg_dct = _deserialize(func_args) # pyright: ignore[reportAny]
-
- return dict(
- id=f"{original_func_name}:{tool_id}",
- name=func_name,
- arguments=arg_dct, # pyright: ignore[reportAny]
- )
-
- tokenizer._tool_call_start = tool_call_start
- tokenizer._tool_call_end = tool_call_end
- tokenizer._tool_parser = parse_tool_call
-
-
-def patch_glm_tokenizer(tokenizer: TokenizerWrapper):
- """
- Fixed version of mlx_lm's glm47 tool parser that handles regex match failures.
- """
- import ast
- import json
- from typing import Any
-
- import regex as re
-
- _func_name_regex = re.compile(r"^(.*?)<arg_key>", re.DOTALL)
- _func_arg_regex = re.compile(
- r"<arg_key>(.*?)</arg_key>(?:\n|\s)*<arg_value>(.*?)(?:</arg_value>|(?=<arg_key>)|$)",
- re.DOTALL,
- )
-
- tool_call_start = "<tool_call>"
- tool_call_end = "</tool_call>"
-
- def _is_string_type(
- tool_name: str,
- arg_name: str,
- tools: list[Any] | None,
- ) -> bool:
- if tools is None:
- return False
- for tool in tools: # pyright: ignore[reportAny]
- func = tool["function"] # pyright: ignore[reportAny]
- if func["name"] == tool_name:
- params = func["parameters"] # pyright: ignore[reportAny]
- if params is None:
- return False
- props = params.get("properties", {}) # pyright: ignore[reportAny]
- arg_props = props.get(arg_name, {}) # pyright: ignore[reportAny]
- arg_type = arg_props.get("type", None) # pyright: ignore[reportAny]
- return arg_type == "string" # pyright: ignore[reportAny]
- return False
-
- def _deserialize(value: str) -> Any: # pyright: ignore[reportAny]
- try:
- return json.loads(value) # pyright: ignore[reportAny]
- except Exception:
- pass
- try:
- return ast.literal_eval(value) # pyright: ignore[reportAny]
- except Exception:
- pass
- return value
-
- def parse_tool_call(text: str, tools: list[Any] | None = None):
- func_name_match = _func_name_regex.search(text)
- if func_name_match is None:
- raise ValueError(f"Could not parse function name from tool call: {text!r}")
- func_name = func_name_match.group(1)
-
- pairs = _func_arg_regex.findall(text)
- arg_dct: dict[str, Any] = {}
- for key, value in pairs: # pyright: ignore[reportAny]
- arg_key = key.strip() # pyright: ignore[reportAny]
- arg_val = value.strip() # pyright: ignore[reportAny]
- if not _is_string_type(func_name, arg_key, tools): # pyright: ignore[reportAny]
- arg_val = _deserialize(arg_val) # pyright: ignore[reportAny]
- arg_dct[arg_key] = arg_val
- return dict(name=func_name, arguments=arg_dct)
-
- tokenizer._tool_call_start = tool_call_start
- tokenizer._tool_call_end = tool_call_end
- tokenizer._tool_parser = parse_tool_call
-
-
-def _validate_single_tool(obj: dict[str, Any]) -> ToolCallItem:
- if (
- ((name := obj.get("name")) is not None)
- and ((args := obj.get("arguments")) is not None)
- and isinstance(name, str)
- ):
- raw_id: object = obj.get("id")
- extra = {"id": str(raw_id)} if raw_id is not None else {}
- return ToolCallItem(
- **extra,
- name=name,
- arguments=json.dumps(args),
- )
- else:
- raise ValidationError
-
-
EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"
EXO_RUNNER_MUST_OOM = "EXO RUNNER MUST OOM"
EXO_RUNNER_MUST_TIMEOUT = "EXO RUNNER MUST TIMEOUT"
diff --git a/src/exo/worker/runner/tool_parsers.py b/src/exo/worker/runner/tool_parsers.py
new file mode 100644
index 00000000..88cbf0be
--- /dev/null
+++ b/src/exo/worker/runner/tool_parsers.py
@@ -0,0 +1,72 @@
+import json
+from dataclasses import dataclass
+from typing import Any, Callable
+
+from exo.shared.types.api import ToolCallItem
+
+
+@dataclass
+class ToolParser:
+ start_parsing: str
+ end_parsing: str
+ parse_tool_calls: Callable[[str], list[ToolCallItem] | None]
+
+
+def make_mlx_parser(
+ tool_call_start: str,
+ tool_call_end: str,
+ tool_parser: Callable[[str], dict[str, Any] | list[dict[str, Any]]],
+) -> ToolParser:
+ def parse_tool_calls(text: str) -> list[ToolCallItem] | None:
+ try:
+ text = text.removeprefix(tool_call_start)
+ text = text.removesuffix(tool_call_end)
+ parsed = tool_parser(text)
+ if isinstance(parsed, list):
+ return [ToolCallItem.model_validate(_flatten(p)) for p in parsed]
+ else:
+ return [ToolCallItem.model_validate(_flatten(parsed))]
+
+ except Exception:
+ return None
+
+ return ToolParser(
+ start_parsing=tool_call_start,
+ end_parsing=tool_call_end,
+ parse_tool_calls=parse_tool_calls,
+ )
+
+
+# TODO / example code:
+def _parse_json_calls(text: str) -> list[ToolCallItem] | None:
+ try:
+ text = text.removeprefix("<tool_call>")
+ text = text.removesuffix("</tool_call>")
+ top_level = {
+ k: json.dumps(v) if isinstance(v, (dict, list)) else v
+ for k, v in json.loads(text).items() # pyright: ignore[reportAny]
+ }
+ return [ToolCallItem.model_validate(top_level)]
+ except Exception:
+ return None
+
+
+def _flatten(p: dict[str, Any]) -> dict[str, str]:
+ return {
+ k: json.dumps(v) if isinstance(v, (dict, list)) else str(v) # pyright: ignore[reportAny]
+ for k, v in p.items() # pyright: ignore[reportAny]
+ }
+
+
+json_tool_parser = ToolParser(
+ start_parsing="<tool_call>",
+ end_parsing="</tool_call>",
+ parse_tool_calls=_parse_json_calls,
+)
+
+
+def infer_tool_parser(chat_template: str) -> ToolParser | None:
+ """Attempt to auto-infer a tool parser from the chat template."""
+ if "<tool_call>" in chat_template and "tool_call.name" in chat_template:
+ return json_tool_parser
+ return None
diff --git a/src/exo/worker/tests/unittests/test_runner/test_parse_tool_calls.py b/src/exo/worker/tests/unittests/test_runner/test_parse_tool_calls.py
index 31f4822e..8a23a18c 100644
--- a/src/exo/worker/tests/unittests/test_runner/test_parse_tool_calls.py
+++ b/src/exo/worker/tests/unittests/test_runner/test_parse_tool_calls.py
@@ -5,12 +5,13 @@ from typing import Any
from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse
from exo.worker.runner.runner import parse_tool_calls
+from exo.worker.runner.tool_parsers import make_mlx_parser
def _make_responses(
texts: list[str],
finish_on_last: bool = True,
-) -> Generator[GenerationResponse | ToolCallResponse]:
+) -> Generator[GenerationResponse]:
"""Create a sequence of GenerationResponses from text strings."""
for i, text in enumerate(texts):
is_last = i == len(texts) - 1
@@ -22,10 +23,13 @@ def _make_responses(
)
-def _dummy_parser(text: str) -> dict[str, Any]:
+def _dummier_parser(text: str) -> dict[str, Any]:
return {"name": "test_fn", "arguments": {"arg": text}}
+_dummy_parser = make_mlx_parser("<tool_call>", "</tool_call>", _dummier_parser)
+
+
class TestParseToolCalls:
"""Tests for parse_tool_calls generator."""
@@ -35,8 +39,6 @@ class TestParseToolCalls:
results = list(
parse_tool_calls(
_make_responses(texts, finish_on_last=False),
- "<tool_call>",
- "</tool_call>",
_dummy_parser,
)
)
@@ -50,8 +52,6 @@ class TestParseToolCalls:
results = list(
parse_tool_calls(
_make_responses(texts),
- "<tool_call>",
- "</tool_call>",
_dummy_parser,
)
)
@@ -76,9 +76,7 @@ class TestParseToolCalls:
results = list(
parse_tool_calls(
_make_responses(texts, finish_on_last=False),
- "<tool_call>",
- "</tool_call>",
- _failing_parser,
+ make_mlx_parser("<tool_call>", "</tool_call>", _failing_parser),
)
)
← 2759e923 api cancellation (#1276)
·
back to Exo
·
bench: add spec for automatic canary benchmarks (#1483) 8392e78a →