From 1b23891e96449c7e4c0e67c4c7b29bd0b012ff71 Mon Sep 17 00:00:00 2001 From: Christopher Albert Date: Thu, 26 Mar 2026 23:21:06 +0100 Subject: [PATCH] fix: enforce tool_choice parameter across all API endpoints Strip tools from chat template when tool_choice="none" and skip tool call parsing on the output side. Inject system hints for "required" and named function modes. Fixes #23. --- tests/test_tool_choice.py | 289 ++++++++++++++++++++++++++++++++++++++ vllm_mlx/server.py | 138 +++++++++++++++--- 2 files changed, 411 insertions(+), 16 deletions(-) create mode 100644 tests/test_tool_choice.py diff --git a/tests/test_tool_choice.py b/tests/test_tool_choice.py new file mode 100644 index 000000000..4f19da294 --- /dev/null +++ b/tests/test_tool_choice.py @@ -0,0 +1,289 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for tool_choice enforcement across API endpoints.""" + +import json + +import pytest +from fastapi.testclient import TestClient + +import vllm_mlx.server as srv +from vllm_mlx.engine.base import GenerationOutput + +TOOL_CALL_MARKUP = ( + '\n{"name": "get_weather", "arguments": {"city": "Paris"}}\n' +) + +SAMPLE_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "get_time", + "description": "Get current time in a timezone", + "parameters": { + "type": "object", + "properties": {"timezone": {"type": "string"}}, + "required": ["timezone"], + }, + }, + }, +] + + +class FakeEngine: + """Fake engine that returns canned output containing tool call markup.""" + + model_name = "test-model" + is_mllm = False + preserve_native_tool_format = False + + def __init__(self, text: str = TOOL_CALL_MARKUP): + self._text = text + self.captured_messages = None + self.captured_kwargs = None + + async def chat(self, messages, **kwargs): + self.captured_messages = messages + self.captured_kwargs = kwargs + return GenerationOutput( + text=self._text, + prompt_tokens=10, + completion_tokens=5, + finish_reason="stop", + ) + + +def _patch_engine(engine): + """Context-manager-like helper to swap the global engine.""" + original_engine = srv._engine + original_model = srv._model_name + srv._engine = engine + srv._model_name = "test-model" + return original_engine, original_model + + +def _restore_engine(original_engine, original_model): + srv._engine = original_engine + srv._model_name = original_model + + +# --------------------------------------------------------------------------- +# Unit tests for _apply_tool_choice +# --------------------------------------------------------------------------- + + +class TestApplyToolChoice: + """Direct unit tests for the _apply_tool_choice helper.""" + + def test_none_strips_tools_and_returns_false(self): + chat_kwargs = {"tools": [{"function": {"name": "f"}}]} + messages = [{"role": "user", "content": "hi"}] + result = srv._apply_tool_choice("none", chat_kwargs, messages) + assert result is False + assert "tools" not in chat_kwargs + assert len(messages) == 1 # no system message added + + def test_required_adds_system_message(self): + chat_kwargs = {"tools": [{"function": {"name": "f"}}]} + messages = [{"role": "user", "content": "hi"}] + result = srv._apply_tool_choice("required", chat_kwargs, messages) + assert result is True + assert len(messages) == 2 + assert messages[-1]["role"] == "system" + assert "MUST call" in messages[-1]["content"] + + def test_dict_filters_tools_and_adds_system_message(self): + chat_kwargs = { + "tools": [ + {"function": {"name": "get_weather"}}, + {"function": {"name": "get_time"}}, + ] + } + messages = [{"role": "user", "content": "hi"}] + result = srv._apply_tool_choice( + {"function": {"name": "get_weather"}}, chat_kwargs, messages + ) + assert result is True + assert len(chat_kwargs["tools"]) == 1 + assert chat_kwargs["tools"][0]["function"]["name"] == "get_weather" + assert len(messages) == 2 + assert "get_weather" in messages[-1]["content"] + + def test_dict_with_no_matching_tool_keeps_all(self): + chat_kwargs = { + "tools": [ + {"function": {"name": "get_weather"}}, + ] + } + messages = [{"role": "user", "content": "hi"}] + srv._apply_tool_choice( + {"function": {"name": "nonexistent"}}, chat_kwargs, messages + ) + assert len(chat_kwargs["tools"]) == 1 # no filter applied + + def test_auto_returns_true_no_changes(self): + chat_kwargs = {"tools": [{"function": {"name": "f"}}]} + messages = [{"role": "user", "content": "hi"}] + result = srv._apply_tool_choice("auto", chat_kwargs, messages) + assert result is True + assert "tools" in chat_kwargs + assert len(messages) == 1 + + def test_none_value_returns_true_no_changes(self): + chat_kwargs = {"tools": [{"function": {"name": "f"}}]} + messages = [{"role": "user", "content": "hi"}] + result = srv._apply_tool_choice(None, chat_kwargs, messages) + assert result is True + assert "tools" in chat_kwargs + assert len(messages) == 1 + + +# --------------------------------------------------------------------------- +# Integration tests via the OpenAI chat endpoint +# --------------------------------------------------------------------------- + + +class TestToolChoiceOpenAIEndpoint: + """Integration tests hitting /v1/chat/completions with tool_choice.""" + + def test_tool_choice_none_strips_tools_and_skips_parsing(self): + engine = FakeEngine(text=TOOL_CALL_MARKUP) + orig_engine, orig_model = _patch_engine(engine) + client = TestClient(srv.app) + try: + response = client.post( + "/v1/chat/completions", + json={ + "model": "test-model", + "messages": [{"role": "user", "content": "weather?"}], + "tools": SAMPLE_TOOLS, + "tool_choice": "none", + "max_tokens": 64, + }, + ) + finally: + _restore_engine(orig_engine, orig_model) + + assert response.status_code == 200 + data = response.json() + msg = data["choices"][0]["message"] + # tool_calls must be absent or None + assert msg.get("tool_calls") is None + # tools should have been stripped from kwargs sent to engine + assert "tools" not in engine.captured_kwargs + # The raw markup should appear as content since parsing was skipped + assert "tool_call" in (msg.get("content") or "") + + def test_tool_choice_required_injects_system_message(self): + engine = FakeEngine(text=TOOL_CALL_MARKUP) + orig_engine, orig_model = _patch_engine(engine) + client = TestClient(srv.app) + try: + response = client.post( + "/v1/chat/completions", + json={ + "model": "test-model", + "messages": [{"role": "user", "content": "weather?"}], + "tools": SAMPLE_TOOLS, + "tool_choice": "required", + "max_tokens": 64, + }, + ) + finally: + _restore_engine(orig_engine, orig_model) + + assert response.status_code == 200 + # Verify system message was injected + sys_msgs = [ + m for m in engine.captured_messages if m.get("role") == "system" + ] + assert any("MUST call" in m["content"] for m in sys_msgs) + + def test_tool_choice_named_filters_tools(self): + engine = FakeEngine(text=TOOL_CALL_MARKUP) + orig_engine, orig_model = _patch_engine(engine) + client = TestClient(srv.app) + try: + response = client.post( + "/v1/chat/completions", + json={ + "model": "test-model", + "messages": [{"role": "user", "content": "weather?"}], + "tools": SAMPLE_TOOLS, + "tool_choice": {"function": {"name": "get_weather"}}, + "max_tokens": 64, + }, + ) + finally: + _restore_engine(orig_engine, orig_model) + + assert response.status_code == 200 + # Verify tools were filtered to only get_weather + template_tools = engine.captured_kwargs.get("tools", []) + assert len(template_tools) == 1 + assert template_tools[0]["function"]["name"] == "get_weather" + # Verify system message mentions the function + sys_msgs = [ + m for m in engine.captured_messages if m.get("role") == "system" + ] + assert any("get_weather" in m["content"] for m in sys_msgs) + + def test_tool_choice_auto_no_changes(self): + engine = FakeEngine(text="Just plain text, no tools.") + orig_engine, orig_model = _patch_engine(engine) + client = TestClient(srv.app) + try: + response = client.post( + "/v1/chat/completions", + json={ + "model": "test-model", + "messages": [{"role": "user", "content": "hello"}], + "tools": SAMPLE_TOOLS, + "tool_choice": "auto", + "max_tokens": 64, + }, + ) + finally: + _restore_engine(orig_engine, orig_model) + + assert response.status_code == 200 + # tools should still be in kwargs + assert "tools" in engine.captured_kwargs + assert len(engine.captured_kwargs["tools"]) == 2 + # No extra system message injected + sys_msgs = [ + m for m in engine.captured_messages if m.get("role") == "system" + ] + assert not any("MUST call" in m.get("content", "") for m in sys_msgs) + + def test_tool_choice_omitted_behaves_as_auto(self): + engine = FakeEngine(text="Plain text response.") + orig_engine, orig_model = _patch_engine(engine) + client = TestClient(srv.app) + try: + response = client.post( + "/v1/chat/completions", + json={ + "model": "test-model", + "messages": [{"role": "user", "content": "hello"}], + "tools": SAMPLE_TOOLS, + "max_tokens": 64, + }, + ) + finally: + _restore_engine(orig_engine, orig_model) + + assert response.status_code == 200 + assert "tools" in engine.captured_kwargs + assert len(engine.captured_kwargs["tools"]) == 2 diff --git a/vllm_mlx/server.py b/vllm_mlx/server.py index ae643b23d..6435dd5da 100644 --- a/vllm_mlx/server.py +++ b/vllm_mlx/server.py @@ -448,6 +448,60 @@ def _parse_tool_calls_with_parser( return parse_tool_calls(output_text, request_dict) +def _apply_tool_choice( + tool_choice: str | dict | None, + chat_kwargs: dict, + messages: list[dict], +) -> bool: + """Apply tool_choice policy to chat kwargs and messages. + + Modifies *chat_kwargs* and *messages* in place so that the chat template + and downstream parsing honour the caller's tool_choice setting. + + Returns ``True`` when the model output should be parsed for tool calls, + ``False`` when tool-call parsing must be skipped (``tool_choice="none"``). + """ + if tool_choice == "none": + chat_kwargs.pop("tools", None) + return False + + if tool_choice == "required": + messages.append( + { + "role": "system", + "content": ( + "You MUST call one of the provided tools. " + "Do not respond with plain text." + ), + } + ) + return True + + if isinstance(tool_choice, dict): + func_info = tool_choice.get("function", {}) + fname = func_info.get("name", "") if isinstance(func_info, dict) else "" + if fname: + messages.append( + { + "role": "system", + "content": f"You MUST call the function: {fname}", + } + ) + template_tools = chat_kwargs.get("tools") + if template_tools: + filtered = [ + t + for t in template_tools + if t.get("function", {}).get("name") == fname + ] + if filtered: + chat_kwargs["tools"] = filtered + return True + + # "auto" or None — no changes needed + return True + + def _new_response_item_id(prefix: str) -> str: """Generate stable OpenAI-style item ids.""" return f"{prefix}_{uuid.uuid4().hex}" @@ -676,6 +730,15 @@ def _responses_request_to_chat_request(request: ResponsesRequest) -> ChatComplet detail="Responses reasoning configuration is not supported on this backend", ) + if isinstance(request.input, list): + for item in request.input: + item_type = item.get("type") if isinstance(item, dict) else getattr(item, "type", None) + if item_type == "reasoning": + raise HTTPException( + status_code=400, + detail="reasoning input items are not supported on this backend", + ) + tools, unsupported_tools = _responses_tools_to_chat_tools(request.tools) messages = _responses_input_to_chat_messages(request) if unsupported_tools: @@ -830,7 +893,7 @@ def _build_response_object( def _prepare_responses_request( request: ResponsesRequest, -) -> tuple[BaseEngine, ChatCompletionRequest, list[dict], dict]: +) -> tuple[BaseEngine, ChatCompletionRequest, list[dict], dict, bool]: """Prepare a Responses request for execution on the chat engine.""" _validate_model_name(request.model) engine = get_engine() @@ -856,12 +919,15 @@ def _prepare_responses_request( } if request.tools: chat_kwargs["tools"] = convert_tools_for_template(chat_request.tools) + should_parse_tools = _apply_tool_choice( + chat_request.tool_choice, chat_kwargs, messages + ) if images: chat_kwargs["images"] = images if videos: chat_kwargs["videos"] = videos - return engine, chat_request, messages, chat_kwargs + return engine, chat_request, messages, chat_kwargs, should_parse_tools async def _run_responses_request( @@ -869,7 +935,9 @@ async def _run_responses_request( raw_request: Request, ) -> tuple[ResponseObject | None, list[dict]]: """Execute a Responses API request against the backend chat engine.""" - engine, chat_request, messages, chat_kwargs = _prepare_responses_request(request) + engine, chat_request, messages, chat_kwargs, should_parse_tools = ( + _prepare_responses_request(request) + ) timeout = _default_timeout output = await _wait_with_disconnect( @@ -880,7 +948,12 @@ async def _run_responses_request( if output is None: return None, [] - cleaned_text, tool_calls = _parse_tool_calls_with_parser(output.text, chat_request) + if should_parse_tools: + cleaned_text, tool_calls = _parse_tool_calls_with_parser( + output.text, chat_request + ) + else: + cleaned_text, tool_calls = output.text, None reasoning_text = None if _reasoning_parser and not tool_calls: reasoning_text, cleaned_text = _reasoning_parser.extract_reasoning( @@ -915,7 +988,9 @@ async def _run_responses_request( async def _stream_responses_request(request: ResponsesRequest) -> AsyncIterator[str]: """Execute a Responses API request and stream SSE events incrementally.""" - engine, chat_request, messages, chat_kwargs = _prepare_responses_request(request) + engine, chat_request, messages, chat_kwargs, should_parse_tools = ( + _prepare_responses_request(request) + ) response_id = _new_response_item_id("resp") sequence = 1 @@ -1037,7 +1112,7 @@ def _start_reasoning_item() -> list[str]: tool_parser = None tool_accumulated_text = "" tool_markup_possible = False - if _enable_auto_tool_choice and _tool_call_parser: + if should_parse_tools and _enable_auto_tool_choice and _tool_call_parser: if _tool_parser_instance is None: try: parser_cls = ToolParserManager.get_tool_parser(_tool_call_parser) @@ -1147,9 +1222,12 @@ def _start_reasoning_item() -> list[str]: ) sequence += 1 - cleaned_text, tool_calls = _parse_tool_calls_with_parser( - raw_accumulated_text, chat_request - ) + if should_parse_tools: + cleaned_text, tool_calls = _parse_tool_calls_with_parser( + raw_accumulated_text, chat_request + ) + else: + cleaned_text, tool_calls = raw_accumulated_text, None final_text = accumulated_text if cleaned_text is not None and not final_text and not tool_calls: final_text = clean_output_text(cleaned_text) @@ -2324,11 +2402,20 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re # Add tools if provided if request.tools: chat_kwargs["tools"] = convert_tools_for_template(request.tools) + should_parse_tools = _apply_tool_choice( + request.tool_choice, chat_kwargs, messages + ) if request.stream: return StreamingResponse( _disconnect_guard( - stream_chat_completion(engine, messages, request, **chat_kwargs), + stream_chat_completion( + engine, + messages, + request, + should_parse_tools=should_parse_tools, + **chat_kwargs, + ), raw_request, ), media_type="text/event-stream", @@ -2353,7 +2440,10 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re ) # Parse tool calls from output using configured parser - cleaned_text, tool_calls = _parse_tool_calls_with_parser(output.text, request) + if should_parse_tools: + cleaned_text, tool_calls = _parse_tool_calls_with_parser(output.text, request) + else: + cleaned_text, tool_calls = output.text, None # Extract reasoning content FIRST (strips channel tokens before JSON extraction) reasoning_text = None @@ -2523,6 +2613,9 @@ async def create_anthropic_message( if openai_request.tools: chat_kwargs["tools"] = convert_tools_for_template(openai_request.tools) + should_parse_tools = _apply_tool_choice( + openai_request.tool_choice, chat_kwargs, messages + ) start_time = time.perf_counter() timeout = _default_timeout @@ -2542,9 +2635,12 @@ async def create_anthropic_message( ) # Parse tool calls - cleaned_text, tool_calls = _parse_tool_calls_with_parser( - output.text, openai_request - ) + if should_parse_tools: + cleaned_text, tool_calls = _parse_tool_calls_with_parser( + output.text, openai_request + ) + else: + cleaned_text, tool_calls = output.text, None # Clean output text final_content = None @@ -2680,6 +2776,9 @@ async def _stream_anthropic_messages( if openai_request.tools: chat_kwargs["tools"] = convert_tools_for_template(openai_request.tools) + should_parse_tools = _apply_tool_choice( + openai_request.tool_choice, chat_kwargs, messages + ) # Emit message_start message_start = { @@ -2733,7 +2832,12 @@ async def _stream_anthropic_messages( yield f"event: content_block_delta\ndata: {json.dumps(delta_event)}\n\n" # Check for tool calls in accumulated text - _, tool_calls = _parse_tool_calls_with_parser(accumulated_text, openai_request) + if should_parse_tools: + _, tool_calls = _parse_tool_calls_with_parser( + accumulated_text, openai_request + ) + else: + tool_calls = None # Emit content_block_stop for text block yield f"event: content_block_stop\ndata: {json.dumps({'type': 'content_block_stop', 'index': 0})}\n\n" @@ -2836,6 +2940,8 @@ async def stream_chat_completion( engine: BaseEngine, messages: list, request: ChatCompletionRequest, + *, + should_parse_tools: bool = True, **kwargs, ) -> AsyncIterator[str]: """Stream chat completion response.""" @@ -2882,7 +2988,7 @@ async def stream_chat_completion( tool_accumulated_text = "" tool_calls_detected = False tool_markup_possible = False # Fast path: skip parsing until '<' seen - if _enable_auto_tool_choice and _tool_call_parser: + if should_parse_tools and _enable_auto_tool_choice and _tool_call_parser: # Initialize parser if needed (same as _parse_tool_calls_with_parser) if _tool_parser_instance is None: try: