feat: enforce tool_choice parameter in chat/completions#16
Conversation
Apply tool_choice policy before passing tools to the chat template:
- "none": strip tools entirely (no tool calls possible)
- "auto": pass tools unchanged (default, current behavior)
- "required": keep tools + inject system instruction to force tool use
- {"type":"function","function":{"name":"X"}}: filter to specific tool
New helper: resolve_tool_choice() returns filtered tools and optional
system instruction based on the policy.
Adds tool_choice and tools fields to ChatRequest model.
Adds 9 tests: 7 unit tests for resolve_tool_choice, 2 endpoint tests.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
Adds tools / tool_choice support to the /chat/completions request schema and applies tool_choice policy when building the prompt, with accompanying unit tests.
Changes:
- Extend
ChatRequestwithtoolsandtool_choicefields. - Add
resolve_tool_choice()and inject a system instruction forrequired/ specific-tool modes. - Add tests covering
resolve_tool_choice()behavior and basic endpoint wiring.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
mlx_vlm/server.py |
Adds request fields and applies tool_choice policy before calling apply_chat_template. |
mlx_vlm/tests/test_server.py |
Adds unit tests for tool-choice resolution and endpoint behavior. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| class ChatRequest(GenerationRequest): | ||
| messages: List[ChatMessage] | ||
| tools: Optional[List[dict]] = Field( | ||
| None, description="Tool definitions the model may call.", | ||
| ) | ||
| tool_choice: Optional[Any] = Field( |
There was a problem hiding this comment.
tool_choice is typed as Optional[Any], which means invalid values (e.g. unknown strings or malformed dicts) will silently pass validation and then be ignored by resolve_tool_choice. If the intent is to enforce semantics, consider narrowing the type (e.g. Literal["none","auto","required"] | {type:"function", function:{name:str}}) or adding a Pydantic validator that raises 422/400 for unsupported values.
| class ChatRequest(GenerationRequest): | |
| messages: List[ChatMessage] | |
| tools: Optional[List[dict]] = Field( | |
| None, description="Tool definitions the model may call.", | |
| ) | |
| tool_choice: Optional[Any] = Field( | |
| class ToolChoiceFunctionSpec(TypedDict): | |
| name: str | |
| class ToolChoiceFunction(TypedDict): | |
| type: Required[Literal["function"]] | |
| function: Required[ToolChoiceFunctionSpec] | |
| ToolChoice: TypeAlias = Union[ | |
| Literal["none", "auto", "required"], | |
| ToolChoiceFunction, | |
| ] | |
| class ChatRequest(GenerationRequest): | |
| messages: List[ChatMessage] | |
| tools: Optional[List[dict]] = Field( | |
| None, description="Tool definitions the model may call.", | |
| ) | |
| tool_choice: Optional[ToolChoice] = Field( |
| return ( | ||
| filtered or tools, |
There was a problem hiding this comment.
For a specific-function tool_choice, when the requested tool name is not present in tools, the code falls back to returning all tools but still emits an instruction requiring the missing tool. This creates an impossible instruction and can cause the model to produce unparseable tool calls. Prefer rejecting the request (HTTP 400/422) or changing the instruction to reflect the fallback (e.g. require one of the available tools).
| return ( | |
| filtered or tools, | |
| if not filtered: | |
| raise HTTPException( | |
| status_code=422, | |
| detail=f"Requested tool_choice function '{name}' is not present in tools.", | |
| ) | |
| return ( | |
| filtered, |
| # Apply tool_choice policy | ||
| tool_choice = getattr(request, "tool_choice", None) | ||
| tools, tool_instruction = resolve_tool_choice(tools, tool_choice) | ||
| if tool_instruction: | ||
| processed_messages.insert(0, {"role": "system", "content": tool_instruction}) |
There was a problem hiding this comment.
When tool_choice resolves to tools=None (e.g. tool_choice="none"), the endpoint still infers a tool_parser_type from the chat template and will continue parsing/removing tool-call blocks from the model output. To match none semantics, consider disabling tool parsing (set tool_parser_type=None / skip process_tool_calls) when tools is None.
| def test_resolve_tool_choice_auto_passthrough(): | ||
| """tool_choice='auto' should return tools unchanged.""" | ||
| tools = [{"function": {"name": "search"}}] | ||
| result_tools, instruction = server.resolve_tool_choice(tools, "auto") | ||
| assert result_tools is tools | ||
| assert instruction is None |
There was a problem hiding this comment.
These tests assert object identity (is tools) rather than equality. This makes the tests brittle (a future refactor that returns an equivalent copy would fail while preserving behavior). Prefer asserting equality and/or that the contents are unchanged instead of requiring the same list instance.
| def test_resolve_tool_choice_specific_unknown_falls_back(): | ||
| """Unknown function name should fall back to all tools.""" | ||
| tools = [{"function": {"name": "search"}}] | ||
| choice = {"type": "function", "function": {"name": "nonexistent"}} | ||
| result_tools, instruction = server.resolve_tool_choice(tools, choice) | ||
| assert result_tools is tools | ||
| assert "nonexistent" in instruction |
There was a problem hiding this comment.
The test encodes behavior where an unknown specific tool_choice both (a) falls back to all tools and (b) still instructs the model to call the unknown tool. That combination is self-contradictory and leads to impossible prompts. Consider updating expected behavior to either reject the request (400/422) or emit an instruction consistent with the fallback.
| def test_resolve_tool_choice_specific_unknown_falls_back(): | |
| """Unknown function name should fall back to all tools.""" | |
| tools = [{"function": {"name": "search"}}] | |
| choice = {"type": "function", "function": {"name": "nonexistent"}} | |
| result_tools, instruction = server.resolve_tool_choice(tools, choice) | |
| assert result_tools is tools | |
| assert "nonexistent" in instruction | |
| def test_resolve_tool_choice_specific_unknown_falls_back_without_instruction(): | |
| """Unknown function name should fall back to all tools without impossible instructions.""" | |
| tools = [{"function": {"name": "search"}}] | |
| choice = {"type": "function", "function": {"name": "nonexistent"}} | |
| result_tools, instruction = server.resolve_tool_choice(tools, choice) | |
| assert result_tools is tools | |
| assert instruction is None |
- Validate tool_choice: reject invalid string values with 400 - Unknown tool name now returns 400 instead of falling back - Skip tool parsing when tool_choice="none" - Use == instead of is for equality checks in tests Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Validate tool_choice: reject invalid string values with 400 - Unknown tool name now returns 400 instead of falling back - Skip tool parsing when tool_choice="none" - Use == instead of is for equality checks in tests Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Summary\ntool_choice: none/auto/required/specific function. 9 tests.