From cc2308b286d8893d8365c3478046099018c19447 Mon Sep 17 00:00:00 2001 From: Andrew Chen <48723787+chuenchen309@users.noreply.github.com> Date: Thu, 16 Jul 2026 02:10:22 +0800 Subject: [PATCH] fix(protocol): check_function_call_usage raises raw KeyError on malformed tool_choice dict ChatCompletionRequest.tool_choice is typed as Optional[Union[Literal["none","auto"], Dict]] -- the Dict case has no nested field-level schema, so pydantic accepts any dict shape. When tool_choice is a dict missing "type" or with a "function" sub-object missing "name" (both schema-valid per the type hint, just semantically incomplete), check_function_call_usage indexed those keys directly and raised a raw KeyError instead of the documented BadRequestError. Every public MLCEngine/AsyncMLCEngine chat-completion method's docstring documents "Raises: e : BadRequestError -- BadRequestError is raised when the request is invalid," and check_function_call_usage is called directly (no try/except) from engine_base.process_chat_completion_request. A caller written to the documented contract -- catching BadRequestError to produce a graceful error response -- would let this KeyError propagate uncaught instead. Fix: use .get() for the "type" lookup and validate the "function" sub-object (must be a dict containing "name") before indexing into it, raising the existing BadRequestError type with a descriptive message in both cases. No change to the valid-request code path. Added tests/python/protocol/test_openai_api_protocol.py (pytest.mark.unittest, no GPU/TVM build needed) covering: missing "name", missing "type", and the existing valid-tool-choice path. Locally verified red/green with an importlib-based harness that loads the protocol module directly (pydantic/shortuuid/fastapi only) since this environment doesn't have TVM built to import the full mlc_llm package -- both crash cases raised a raw KeyError against unfixed code and the documented BadRequestError after the fix; the valid-tool-choice path is unaffected. ruff check + ruff format: clean. Duplicate-work check: open PR #3190 (long-standing, approved, unmerged) touches this same function but only removes two blank lines around it -- it does not add the missing-key guards this PR adds, so there's no functional overlap, just a possible trivial rebase if #3190 lands first. Co-Authored-By: Claude Opus 4.8 --- .../mlc_llm/protocol/openai_api_protocol.py | 12 +++-- .../protocol/test_openai_api_protocol.py | 49 +++++++++++++++++++ 2 files changed, 57 insertions(+), 4 deletions(-) create mode 100644 tests/python/protocol/test_openai_api_protocol.py diff --git a/python/mlc_llm/protocol/openai_api_protocol.py b/python/mlc_llm/protocol/openai_api_protocol.py index 9471a7bfbe..1749b0a880 100644 --- a/python/mlc_llm/protocol/openai_api_protocol.py +++ b/python/mlc_llm/protocol/openai_api_protocol.py @@ -375,20 +375,24 @@ def check_function_call_usage(self, conv_template: Conversation) -> None: # select the tool based on the tool_choice if specified if isinstance(self.tool_choice, dict): - if self.tool_choice["type"] != "function": + if self.tool_choice.get("type") != "function": raise BadRequestError("Only 'function' tool choice is supported") - if len(self.tool_choice["function"]) > 1: + tool_choice_function = self.tool_choice.get("function") + if not isinstance(tool_choice_function, dict) or "name" not in tool_choice_function: + raise BadRequestError("tool_choice.function must be an object with a 'name' field") + + if len(tool_choice_function) > 1: raise BadRequestError("Only one tool is supported when tool_choice is specified") for tool in self.tools: - if tool.function.name == self.tool_choice["function"]["name"]: + if tool.function.name == tool_choice_function["name"]: conv_template.use_function_calling = True conv_template.function_string = tool.function.model_dump_json(by_alias=True) return raise BadRequestError( - f"The tool_choice function {self.tool_choice['function']['name']}" + f"The tool_choice function {tool_choice_function['name']}" " is not found in the tools list" ) diff --git a/tests/python/protocol/test_openai_api_protocol.py b/tests/python/protocol/test_openai_api_protocol.py new file mode 100644 index 0000000000..942723aaf2 --- /dev/null +++ b/tests/python/protocol/test_openai_api_protocol.py @@ -0,0 +1,49 @@ +import pytest + +from mlc_llm.protocol.conversation_protocol import Conversation +from mlc_llm.protocol.error_protocol import BadRequestError +from mlc_llm.protocol.openai_api_protocol import ChatCompletionRequest + +# test category "unittest" +pytestmark = [pytest.mark.unittest] + + +def _make_request(tool_choice): + return ChatCompletionRequest( + messages=[{"role": "user", "content": "hi"}], + model="test", + tools=[ + { + "type": "function", + "function": {"name": "foo", "description": "d", "parameters": {}}, + } + ], + tool_choice=tool_choice, + ) + + +def _make_conv(): + return Conversation(name="test", roles={"user": "user", "assistant": "assistant"}, seps=[" "]) + + +def test_check_function_call_usage_missing_function_name_raises_bad_request(): + """tool_choice is an unstructured Dict, so a request can omit the required + "name" field. This must raise the documented BadRequestError, not a raw + KeyError.""" + request = _make_request({"type": "function", "function": {}}) + with pytest.raises(BadRequestError): + request.check_function_call_usage(_make_conv()) + + +def test_check_function_call_usage_missing_type_raises_bad_request(): + """Same class of issue for a missing "type" key.""" + request = _make_request({"function": {"name": "foo"}}) + with pytest.raises(BadRequestError): + request.check_function_call_usage(_make_conv()) + + +def test_check_function_call_usage_selects_tool_on_valid_choice(): + request = _make_request({"type": "function", "function": {"name": "foo"}}) + conv_template = _make_conv() + request.check_function_call_usage(conv_template) + assert conv_template.use_function_calling is True