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 - ) - yield response - - in_tool_call = False - tool_call_text_parts = [] - continue 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() + ) + 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 + 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"^(.*?)", re.DOTALL) - _func_arg_regex = re.compile( - r"(.*?)(?:\n|\s)*(.*?)(?:|(?=)|$)", - re.DOTALL, - ) - - tool_call_start = "" - tool_call_end = "" - - 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("") + text = text.removesuffix("") + 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="", + end_parsing="", + 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 "" 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("", "", _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), - "", - "", _dummy_parser, ) ) @@ -50,8 +52,6 @@ class TestParseToolCalls: results = list( parse_tool_calls( _make_responses(texts), - "", - "", _dummy_parser, ) ) @@ -76,9 +76,7 @@ class TestParseToolCalls: results = list( parse_tool_calls( _make_responses(texts, finish_on_last=False), - "", - "", - _failing_parser, + make_mlx_parser("", "", _failing_parser), ) )