[object Object]

← back to Exo

fix: coerce tool-call argument types from tool schema (#1651)

5777bf3c395a2dfb5cd401b8245bb74d13aa5690 · 2026-03-04 04:22:53 -0800 · Owleksiy

Apply schema-aware coercion to parsed tool-call arguments so
Hermes-style toolcalls can still return typed JSON (e.g. integer ids).

 - pass request tools into parse_tool_calls
 - coerce parsed argument values by function parameters schema
 - add unit tests for coercion and unknown-tool passthrough

## Motivation

Models that use Hermes-based toolcall syntax (Qwen3.5) can't reliably
call tools with non-string parameters
Example tool:
```json
    {
      "type": "function",
      "function": {
        "name": "process",
        "description": "Manage background processes",
        "parameters": {
          "type": "object",
          "properties": {
            "action": {
              "type": "string",
              "enum": ["spawn", "output", "kill", "list"]
            },
            "id": {
              "type": "integer",
              "description": "Process id"
            },
            "command": {
              "type": "string",
              "description": "Command to run for spawn"
            }
          },
          "required": ["action"],
          "additionalProperties": false
        }
      }
    }
```
Model transcript:
```
<tool_call>
<function=process>
<parameter=action>
output
</parameter>
<parameter=id>
0
</parameter>
</function>
</tool_call>
```
And the API returns:

`{"id":"a8f11689-d840-4ca5-ab1d-ead3678a11a9","name":"process","arguments":"{\"action\":
\"output\", \"id\": \"0\"}"}}`

Tool definition declared `id` as `integer`, the model output is
type-agnostic, and the translation layer treats everything as a string.

The same Qwen3.5-27B on OpenRouter and GPT-4.1-mini on openai obey the
function signature and emit correct call:
```
{"name":"process","arguments":"{\"action\": \"output\", \"id\": 0}"}}
```

Steps to reproduce:
```
❯ curl -sS -v http://localhost:52415/v1/chat/completions \
  -H 'Content-Type: application/json' \  -d @- <<'JSON'
{
  "model": "mlx-community/Qwen3.5-27B-4bit",
  "stream": false,
  "temperature": 0,
  "messages": [
    {
      "role": "user",
      "content": "Call the process tool with action=output and id=0. Do not explain anything. Just make the tool call."
    }
  ],
  "tools": [
    {
      "type": "function",
      "function": {
        "name": "process",
        "description": "Manage background processes",
        "parameters": {
          "type": "object",
          "properties": {
            "action": {
              "type": "string",
              "enum": ["spawn", "output", "kill", "list"]
            },
            "id": {
              "type": "integer",
              "description": "Process id"
            },
            "command": {
              "type": "string",
              "description": "Command to run for spawn"
            }
          },
          "required": ["action"],
          "additionalProperties": false
        }
      }
    }
  ],
  "tool_choice": {
    "type": "function",
    "function": {
      "name": "process"
    }
  }
}
JSON
```
Look for type of `id` function call

## Changes

Function call parameters are now converted to the types that the
function declaration has

## Why It Works

We now explicitly convert types where we know it (and skip if we don't)

## Test Plan

### Manual Testing
Create Qwen3.5-<ANY> instance, send the curl command above. Check that
'id' is now serialized as a number

### Automated Testing
Unit tests to cover basic type conversions

---------

Co-authored-by: Evan <evanev7@gmail.com>

Files touched

Diff

commit 5777bf3c395a2dfb5cd401b8245bb74d13aa5690
Author: Owleksiy <sblOWPCKCR@protonmail.com>
Date:   Wed Mar 4 04:22:53 2026 -0800

    fix: coerce tool-call argument types from tool schema (#1651)
    
    Apply schema-aware coercion to parsed tool-call arguments so
    Hermes-style toolcalls can still return typed JSON (e.g. integer ids).
    
     - pass request tools into parse_tool_calls
     - coerce parsed argument values by function parameters schema
     - add unit tests for coercion and unknown-tool passthrough
    
    ## Motivation
    
    Models that use Hermes-based toolcall syntax (Qwen3.5) can't reliably
    call tools with non-string parameters
    Example tool:
    ```json
        {
          "type": "function",
          "function": {
            "name": "process",
            "description": "Manage background processes",
            "parameters": {
              "type": "object",
              "properties": {
                "action": {
                  "type": "string",
                  "enum": ["spawn", "output", "kill", "list"]
                },
                "id": {
                  "type": "integer",
                  "description": "Process id"
                },
                "command": {
                  "type": "string",
                  "description": "Command to run for spawn"
                }
              },
              "required": ["action"],
              "additionalProperties": false
            }
          }
        }
    ```
    Model transcript:
    ```
    <tool_call>
    <function=process>
    <parameter=action>
    output
    </parameter>
    <parameter=id>
    0
    </parameter>
    </function>
    </tool_call>
    ```
    And the API returns:
    
    `{"id":"a8f11689-d840-4ca5-ab1d-ead3678a11a9","name":"process","arguments":"{\"action\":
    \"output\", \"id\": \"0\"}"}}`
    
    Tool definition declared `id` as `integer`, the model output is
    type-agnostic, and the translation layer treats everything as a string.
    
    The same Qwen3.5-27B on OpenRouter and GPT-4.1-mini on openai obey the
    function signature and emit correct call:
    ```
    {"name":"process","arguments":"{\"action\": \"output\", \"id\": 0}"}}
    ```
    
    Steps to reproduce:
    ```
    ❯ curl -sS -v http://localhost:52415/v1/chat/completions \
      -H 'Content-Type: application/json' \  -d @- <<'JSON'
    {
      "model": "mlx-community/Qwen3.5-27B-4bit",
      "stream": false,
      "temperature": 0,
      "messages": [
        {
          "role": "user",
          "content": "Call the process tool with action=output and id=0. Do not explain anything. Just make the tool call."
        }
      ],
      "tools": [
        {
          "type": "function",
          "function": {
            "name": "process",
            "description": "Manage background processes",
            "parameters": {
              "type": "object",
              "properties": {
                "action": {
                  "type": "string",
                  "enum": ["spawn", "output", "kill", "list"]
                },
                "id": {
                  "type": "integer",
                  "description": "Process id"
                },
                "command": {
                  "type": "string",
                  "description": "Command to run for spawn"
                }
              },
              "required": ["action"],
              "additionalProperties": false
            }
          }
        }
      ],
      "tool_choice": {
        "type": "function",
        "function": {
          "name": "process"
        }
      }
    }
    JSON
    ```
    Look for type of `id` function call
    
    ## Changes
    
    Function call parameters are now converted to the types that the
    function declaration has
    
    ## Why It Works
    
    We now explicitly convert types where we know it (and skip if we don't)
    
    ## Test Plan
    
    ### Manual Testing
    Create Qwen3.5-<ANY> instance, send the curl command above. Check that
    'id' is now serialized as a number
    
    ### Automated Testing
    Unit tests to cover basic type conversions
    
    ---------
    
    Co-authored-by: Evan <evanev7@gmail.com>
---
 .../worker/runner/llm_inference/batch_generator.py |   1 +
 .../runner/llm_inference/model_output_parsers.py   |  12 +-
 .../worker/runner/llm_inference/tool_parsers.py    | 177 ++++++++++++++++++++-
 .../unittests/test_runner/test_parse_tool_calls.py |  95 +++++++++++
 4 files changed, 277 insertions(+), 8 deletions(-)

diff --git a/src/exo/worker/runner/llm_inference/batch_generator.py b/src/exo/worker/runner/llm_inference/batch_generator.py
index 189efe5e..1ad7b6a6 100644
--- a/src/exo/worker/runner/llm_inference/batch_generator.py
+++ b/src/exo/worker/runner/llm_inference/batch_generator.py
@@ -218,6 +218,7 @@ class SequentialGenerator(InferenceGenerator):
             self.tokenizer,
             type(self.model),
             self.model_id,
+            task.task_params.tools,
         )
         self._active = (task, mlx_gen, queue, output_generator)
 
diff --git a/src/exo/worker/runner/llm_inference/model_output_parsers.py b/src/exo/worker/runner/llm_inference/model_output_parsers.py
index 7c13d827..99e4b6f2 100644
--- a/src/exo/worker/runner/llm_inference/model_output_parsers.py
+++ b/src/exo/worker/runner/llm_inference/model_output_parsers.py
@@ -1,5 +1,6 @@
 from collections.abc import Generator
 from functools import cache
+from typing import Any
 
 from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model
 from mlx_lm.models.gpt_oss import Model as GptOssModel
@@ -37,6 +38,7 @@ def apply_all_parsers(
     tokenizer: TokenizerWrapper,
     model_type: type[Model],
     model_id: ModelId,
+    tools: list[dict[str, Any]] | None,
 ) -> Generator[GenerationResponse | ToolCallResponse | None]:
     mlx_generator = receiver
 
@@ -55,7 +57,7 @@ def apply_all_parsers(
     ):
         mlx_generator = parse_deepseek_v32(mlx_generator)
     elif tool_parser:
-        mlx_generator = parse_tool_calls(mlx_generator, tool_parser)
+        mlx_generator = parse_tool_calls(mlx_generator, tool_parser, tools)
 
     return mlx_generator
 
@@ -325,7 +327,9 @@ def parse_thinking_models(
 
 
 def parse_tool_calls(
-    responses: Generator[GenerationResponse | None], tool_parser: ToolParser
+    responses: Generator[GenerationResponse | None],
+    tool_parser: ToolParser,
+    tools: list[dict[str, Any]] | None,
 ) -> Generator[GenerationResponse | ToolCallResponse | None]:
     in_tool_call = False
     tool_call_text_parts: list[str] = []
@@ -340,9 +344,7 @@ def parse_tool_calls(
             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()
-                )
+                parsed = tool_parser.parse("".join(tool_call_text_parts).strip(), tools)
                 logger.info(f"parsed {tool_call_text_parts=} into {parsed=}")
                 if parsed is not None:
                     yield ToolCallResponse(
diff --git a/src/exo/worker/runner/llm_inference/tool_parsers.py b/src/exo/worker/runner/llm_inference/tool_parsers.py
index ee6d1bb5..8fd02620 100644
--- a/src/exo/worker/runner/llm_inference/tool_parsers.py
+++ b/src/exo/worker/runner/llm_inference/tool_parsers.py
@@ -1,4 +1,5 @@
 import json
+import math
 from dataclasses import dataclass
 from typing import Any, Callable
 
@@ -9,7 +10,177 @@ from exo.shared.types.api import ToolCallItem
 class ToolParser:
     start_parsing: str
     end_parsing: str
-    parse_tool_calls: Callable[[str], list[ToolCallItem] | None]
+    _inner_parser: Callable[[str], list[ToolCallItem] | None]
+
+    def parse(
+        self, text: str, tools: list[dict[str, Any]] | None
+    ) -> list[ToolCallItem] | None:
+        parsed = self._inner_parser(text)
+        if parsed is None:
+            return None
+        if tools is not None:
+            parsed = _coerce_tool_calls_to_schema(parsed, tools)
+        return parsed
+
+
+def _json_type_matches(value: Any, expected_type: str) -> bool:  # pyright: ignore[reportAny]
+    if expected_type == "object":
+        return isinstance(value, dict)
+    if expected_type == "array":
+        return isinstance(value, list)
+    if expected_type == "string":
+        return isinstance(value, str)
+    if expected_type == "integer":
+        return isinstance(value, int) and not isinstance(value, bool)
+    if expected_type == "number":
+        return (isinstance(value, int) and not isinstance(value, bool)) or isinstance(
+            value, float
+        )
+    if expected_type == "boolean":
+        return isinstance(value, bool)
+    if expected_type == "null":
+        return value is None
+    return False
+
+
+def _coerce_tool_arg_with_schema(value: Any, schema: dict[str, Any]) -> Any:  # pyright: ignore[reportAny]
+    schema_type = schema.get("type")
+
+    if isinstance(schema_type, list):
+        for candidate in schema_type:  # pyright: ignore[reportUnknownVariableType]
+            if not isinstance(candidate, str):
+                continue
+            if candidate == "null" and value is None:
+                return None
+            candidate_schema = {**schema, "type": candidate}
+            coerced = _coerce_tool_arg_with_schema(value, candidate_schema)  # pyright: ignore[reportAny]
+            if _json_type_matches(coerced, candidate):
+                return coerced  # pyright: ignore[reportAny]
+        return value  # pyright: ignore[reportAny]
+
+    if not isinstance(schema_type, str):
+        return value  # pyright: ignore[reportAny]
+
+    if schema_type == "object":
+        parsed = value  # pyright: ignore[reportAny]
+        if isinstance(parsed, str):
+            try:
+                parsed = json.loads(parsed)  # pyright: ignore[reportAny]
+            except Exception:
+                return value  # pyright: ignore[reportAny]
+        if not isinstance(parsed, dict):
+            return value  # pyright: ignore[reportAny]
+        properties = schema.get("properties")
+        if not isinstance(properties, dict):
+            return parsed  # pyright: ignore[reportUnknownVariableType]
+        return {
+            key: (
+                _coerce_tool_arg_with_schema(prop_value, prop_schema)  # pyright: ignore[reportUnknownArgumentType]
+                if isinstance(prop_schema, dict)
+                else prop_value
+            )
+            for key, prop_value in parsed.items()  # pyright: ignore[reportUnknownVariableType]
+            for prop_schema in [properties.get(key)]  # type: ignore
+        }
+
+    if schema_type == "array":
+        parsed = value  # pyright: ignore[reportAny]
+        if isinstance(parsed, str):
+            try:
+                parsed = json.loads(parsed)  # pyright: ignore[reportAny]
+            except Exception:
+                return value  # pyright: ignore[reportAny]
+        if not isinstance(parsed, list):
+            return value  # pyright: ignore[reportAny]
+        item_schema = schema.get("items")
+        if not isinstance(item_schema, dict):
+            return parsed  # pyright: ignore[reportUnknownVariableType]
+        return [_coerce_tool_arg_with_schema(item, item_schema) for item in parsed]  # type: ignore
+
+    if schema_type == "integer":
+        if isinstance(value, bool):
+            return value
+        if isinstance(value, int):
+            return value
+        if isinstance(value, float) and value.is_integer():
+            return int(value)
+        if isinstance(value, str):
+            try:
+                return int(value.strip())
+            except ValueError:
+                return value
+        return value
+
+    if schema_type == "number":
+        if isinstance(value, bool):
+            return value
+        if isinstance(value, (int, float)):
+            return value
+        if isinstance(value, str):
+            try:
+                num = float(value.strip())
+                if math.isfinite(num):
+                    return num
+            except ValueError:
+                return value
+        return value
+
+    if schema_type == "boolean":
+        if isinstance(value, bool):
+            return value
+        if isinstance(value, str):
+            lowered = value.strip().lower()
+            if lowered == "true":
+                return True
+            if lowered == "false":
+                return False
+        return value
+
+    return value  # pyright: ignore[reportAny]
+
+
+def _coerce_tool_calls_to_schema(
+    tool_calls: list[ToolCallItem], tools: list[dict[str, Any]]
+) -> list[ToolCallItem]:
+    schema_by_name: dict[str, dict[str, Any]] = {}
+    for tool in tools:
+        function = tool.get("function")
+        if not isinstance(function, dict):
+            continue
+        name = function.get("name")  # type: ignore
+        parameters = function.get("parameters")  # type: ignore
+        if isinstance(name, str) and isinstance(parameters, dict):
+            schema_by_name[name] = parameters
+
+    if not schema_by_name:
+        return tool_calls
+
+    coerced_calls: list[ToolCallItem] = []
+    for tool_call in tool_calls:
+        schema = schema_by_name.get(tool_call.name)
+        if schema is None:
+            coerced_calls.append(tool_call)
+            continue
+
+        try:
+            parsed_args = json.loads(tool_call.arguments)  # pyright: ignore[reportAny]
+        except Exception:
+            coerced_calls.append(tool_call)
+            continue
+
+        if not isinstance(parsed_args, dict):
+            coerced_calls.append(tool_call)
+            continue
+
+        coerced_args = _coerce_tool_arg_with_schema(parsed_args, schema)  # pyright: ignore[reportAny]
+        if not isinstance(coerced_args, dict):
+            coerced_calls.append(tool_call)
+            continue
+
+        coerced_calls.append(
+            tool_call.model_copy(update={"arguments": json.dumps(coerced_args)})
+        )
+    return coerced_calls
 
 
 def make_mlx_parser(
@@ -33,7 +204,7 @@ def make_mlx_parser(
     return ToolParser(
         start_parsing=tool_call_start,
         end_parsing=tool_call_end,
-        parse_tool_calls=parse_tool_calls,
+        _inner_parser=parse_tool_calls,
     )
 
 
@@ -62,7 +233,7 @@ def make_json_parser() -> ToolParser:
     return ToolParser(
         start_parsing="<tool_call>",
         end_parsing="</tool_call>",
-        parse_tool_calls=_parse_json_calls,
+        _inner_parser=_parse_json_calls,
     )
 
 
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 4c4bc392..03be01bd 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
@@ -1,5 +1,6 @@
 """Tests for parse_tool_calls generator, especially unclosed tool call handling."""
 
+import json
 from collections.abc import Generator
 from typing import Any
 
@@ -40,6 +41,7 @@ class TestParseToolCalls:
             parse_tool_calls(
                 _make_responses(texts, finish_on_last=False),
                 _dummy_parser,
+                tools=None,
             )
         )
 
@@ -53,6 +55,7 @@ class TestParseToolCalls:
             parse_tool_calls(
                 _make_responses(texts),
                 _dummy_parser,
+                tools=None,
             )
         )
 
@@ -77,9 +80,101 @@ class TestParseToolCalls:
             parse_tool_calls(
                 _make_responses(texts, finish_on_last=False),
                 make_mlx_parser("<tool_call>", "</tool_call>", _failing_parser),
+                tools=None,
             )
         )
 
         assert len(results) == 1
         assert isinstance(results[0], GenerationResponse)
         assert results[0].text == "<tool_call>bad content</tool_call>"
+
+    def test_tool_schema_coerces_string_arguments_to_expected_types(self):
+        """Tool argument values should be coerced using provided JSON schema."""
+
+        def _parser_with_string_args(_text: str) -> dict[str, Any]:
+            return {
+                "name": "process",
+                "arguments": {
+                    "action": "output",
+                    "id": "0",
+                    "verbose": "true",
+                    "temperature": "0.75",
+                },
+            }
+
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "process",
+                    "description": "Manage background processes",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "action": {"type": "string"},
+                            "id": {"type": "integer"},
+                            "verbose": {"type": "boolean"},
+                            "temperature": {"type": "number"},
+                        },
+                        "required": ["action"],
+                    },
+                },
+            }
+        ]
+
+        results = list(
+            parse_tool_calls(
+                _make_responses(["<tool_call>", "process", "</tool_call>"]),
+                make_mlx_parser(
+                    "<tool_call>", "</tool_call>", _parser_with_string_args
+                ),
+                tools,
+            )
+        )
+
+        assert len(results) == 1
+        assert isinstance(results[0], ToolCallResponse)
+
+        args = json.loads(results[0].tool_calls[0].arguments)  # pyright: ignore[reportAny]
+        assert args == {
+            "action": "output",
+            "id": 0,
+            "verbose": True,
+            "temperature": 0.75,
+        }
+
+    def test_schema_coercion_skips_unknown_tools(self):
+        """If no matching tool schema exists, arguments should remain unchanged."""
+
+        def _parser_with_string_id(_text: str) -> dict[str, Any]:
+            return {
+                "name": "process",
+                "arguments": {"action": "output", "id": "0"},
+            }
+
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "different_tool",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {"id": {"type": "integer"}},
+                    },
+                },
+            }
+        ]
+
+        results = list(
+            parse_tool_calls(
+                _make_responses(["<tool_call>", "process", "</tool_call>"]),
+                make_mlx_parser("<tool_call>", "</tool_call>", _parser_with_string_id),
+                tools,
+            )
+        )
+
+        assert len(results) == 1
+        assert isinstance(results[0], ToolCallResponse)
+
+        args = json.loads(results[0].tool_calls[0].arguments)  # pyright: ignore[reportAny]
+        assert args == {"action": "output", "id": "0"}

← 886192f1 ignore closed resource errors when trying to cancel a task (  ·  back to Exo  ·  #Add reasoning parms to chat completion and responses APIs ( 4de8f801 →