fix(protocol): check_function_call_usage raises raw KeyError on malformed tool_choice dict

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"
             )
 

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