[Enhancement] Honor tool_choice (forced tool calls)

[michaelharrigan]

Motivation

ChatCompletionRequest.tool_choice is accepted on the wire (api.py) but is never read on the way to generation, so "required" and {"type": "function", "function": {"name": ...}} are silently ignored.

This means the model always chooses whether and which tool to call and in turn it makes tool use non-deterministic, which is a problem for structured extraction, agent reliability, and any integration test that needs a specific tool call to happen.

Changes

• shared/types/text_generation.py: added a tool_choice: str | dict[str, Any] | None field to TextGenerationTaskParams.
• api/adapters/chat_completions.py: thread tool_choice=request.tool_choice into the task params (it was previously dropped).
• worker/engines/mlx/utils_mlx.py:
  • _forced_tool_call_prefill(...): when tool_choice is "required" or names a function and tools are present, return the model's own tokenizer.tool_call_start marker.
  • In render_chat_template, use that as the assistant prefill (when there isn't already an explicit assistant prefill).
  • detect_tool_call_prompt_suffix(prompt, tokenizer): a twin of the existing detect_thinking_prompt_suffix — true when the prompt ends with the tool-call start marker.
• worker/runner/llm_inference/model_output_parsers.py: parse_tool_calls gains starts_in_tool_call (default False); apply_all_parsers seeds it from detect_tool_call_prompt_suffix, exactly like the thinking path seeds starts_in_thinking.
• worker/tests/unittests/test_mlx/test_tool_choice_prefill.py: unit tests for the prefill + suffix detection.

Why It Works

Forcing a tool call is just assistant prefilling, which the engine already supports: a trailing assistant message is stripped before templating and re-appended after (render_chat_template).

Prefilling the model's own tool_call_start marker forces it to continue in its native tool-call format (e.g. Qwen's <tool_call><function=...> XML, not a hardcoded JSON shape).

The one subtlety: the forced marker lands in the prompt, so it is not present at the start of the output stream, and parse_tool_calls detects a call via response.text.startswith(tool_call_start).

The thinking path already solves this exact problem with detect_thinking_prompt_suffix + starts_in_thinking; this change mirrors it (detect_tool_call_prompt_suffix + starts_in_tool_call), seeding the parser to begin inside the call and pre-loading the start marker so the parsed call is well formed.

No new decoding machinery, no per-model format assumptions.

Test Plan

Manual Testing

• Hardware: 2-node EXO ring (M3 Ultra Mac Studio (96GB) + Mac Mini M4 Pro (24GB)), pipeline parallelism, serving mlx-community/Qwen3.5-9B-8bit (node wired limit ~77 GiB, group size 2).

• What I did:

  • Sent a deliberately tool-irrelevant prompt — "What is your name?" — with a single bash tool and tool_choice: {"type":"function","function":{"name":"bash"}}.
  • Before this change: finish_reason: stop, empty content, tool_calls: null (the model declines the tool; a naive JSON prefill instead stalls after 1 token because the format doesn't match Qwen's XML).
  • After this change: finish_reason: tool_calls with a real {"name": "bash", "arguments": "{\"command\": \"whoami\"}"} — the model was forced into the tool call it would never otherwise make.
  • Sanity checked that tool_choice absent / "auto" is unchanged (model chooses freely).

Automated Testing

• Added test_tool_choice_prefill.py (7 tests): named-function and "required" both return the start marker; "auto"/None/no-tools/no-marker are no-ops; detect_tool_call_prompt_suffix is true only when the prompt ends with the marker.
• Run: uv run pytest src/exo/worker/tests/unittests/test_mlx/test_tool_choice_prefill.py