openai server: type tool arguments from schema + implement tool_choice, 9 new tool-calling tests (#118)

Two OpenAI-compat tool-calling bugs found against the real GLM-5.2 (dnnspaul): (1) string-typed args coerced to numbers — declared schema type now decides, string kept verbatim, bool rejected as number, schema-less params keep permissive decoding; (2) tool_choice was ignored — none/auto/required/{function} now honored, invalid returns 400. Python-only (openai_server.py + tests), engine untouched. 36/36 tests pass (verified independently in a clean worktree).

c/openai_server.py

@@ -196,23 +196,62 @@ def _tool_param_order(tools):
     return out
 
 
+def _tool_param_types(tools):
+    """name -> {param: declared JSON-schema type}. The model emits every argument as text;
+    without the schema a string-typed value that happens to look numeric ("12345" for an
+    order id, an SKU, a phone number) would be json.loads()'d into an int and the tool would
+    receive the wrong type."""
+    out = {}
+    for tool in (tools or []):
+        fn = tool.get("function", tool) if isinstance(tool, dict) else {}
+        name = fn.get("name")
+        if not name:
+            continue
+        props = ((fn.get("parameters") or {}).get("properties") or {})
+        types = {}
+        for key, spec in props.items():
+            if isinstance(spec, dict):
+                t = spec.get("type")
+                if isinstance(t, list):          # {"type": ["string", "null"]}
+                    t = next((x for x in t if x != "null"), None)
+                types[key] = t
+        out[name] = types
+    return out
+
+
+def _coerce_arg(value, declared):
+    """Decode a raw <arg_value> according to the declared schema type.
+
+    A string-typed parameter is kept verbatim -- never parsed as JSON. Everything else keeps
+    the previous permissive behaviour (parse if it parses, otherwise leave as text)."""
+    if declared == "string":
+        return value
+    try:
+        parsed = json.loads(value)
+    except (json.JSONDecodeError, TypeError):
+        return value
+    if declared in ("integer", "number") and isinstance(parsed, bool):
+        return value                              # `true` is not a number
+    if declared and declared not in ("integer", "number", "boolean", "object", "array"):
+        return value
+    return parsed
+
+
 def parse_tool_calls(reply, tools=None):
     """Return (content, tool_calls). Strict GLM parse; optional de-mangler (COLI_TOOL_SALVAGE=1)
     rescues malformed int4 output by mapping a lone payload onto the tool's primary parameter."""
     param_order = _tool_param_order(tools)
+    param_types = _tool_param_types(tools)
     calls, salvaged = [], []
     for match in _BOX_RE.finditer(reply):
         inner = match.group(1)
         name_match = _NAME_RE.match(inner)
         name = name_match.group(1) if name_match else inner.strip()
         args = {}
+        types = param_types.get(name, {})
         for arg in _ARG_RE.finditer(inner):
             key, value = arg.group(1), arg.group(2)
-            try:
+            args[key] = _coerce_arg(value, types.get(key))
-                value = json.loads(value)
-            except (json.JSONDecodeError, TypeError):
-                pass
-            args[key] = value
         if not args and _SALVAGE:
             rest = inner[name_match.end():] if name_match else ""
             payload = _TAG_RE.sub("", rest).strip()
@@ -241,7 +280,8 @@ def parse_tool_calls(reply, tools=None):
     return text.strip(), calls
 
 
-def render_chat(messages, enable_thinking=False, reasoning_effort=None, tools=None):
+def render_chat(messages, enable_thinking=False, reasoning_effort=None, tools=None,
+                tool_choice=None):
     """Render the text-only subset of the official GLM-5.2 chat template."""
     if not isinstance(messages, list) or not messages:
         raise APIError(400, "`messages` must be a non-empty array.", "messages")
@@ -249,6 +289,15 @@ def render_chat(messages, enable_thinking=False, reasoning_effort=None, tools=No
     if enable_thinking:
         effort = "High" if reasoning_effort == "high" else "Max"
         prompt.append(f"<|system|>Reasoning Effort: {effort}")
+    forced = None
+    if isinstance(tool_choice, dict):
+        forced = ((tool_choice.get("function") or {}).get("name")
+                  or tool_choice.get("name"))
+        if forced:
+            tools = [t for t in (tools or [])
+                     if ((t.get("function", t) if isinstance(t, dict) else {}).get("name") == forced)]
+    elif tool_choice == "none":
+        tools = None                              # the client forbade tools: do not offer them
     if tools:
         # AUTHORITATIVE GLM-5.2 tool-declaration block (byte-matches chat_template.jinja): the
         # `# Tools` + <tools></tools> XML structure is what the model was trained on. A made-up
@@ -264,6 +313,10 @@ def render_chat(messages, enable_thinking=False, reasoning_effort=None, tools=No
                       "within the following XML format:\n<tool_call>{function-name}"
                       "<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value>"
                       "<arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call>")
+        if forced:
+            prompt.append(f"\n\nYou must call the function `{forced}`. Do not answer directly.")
+        elif tool_choice == "required":
+            prompt.append("\n\nYou must call one of the functions above. Do not answer directly.")
     prev_tool = False
     for index, message in enumerate(messages):
         if not isinstance(message, dict):
@@ -309,6 +362,27 @@ def generation_options(body, limit):
     if body.get("n", 1) != 1:
         raise APIError(400, "Colibri currently supports `n=1` only.", "n", "unsupported_value")
     # `tools`/`functions` are handled by render_chat (declaration) + parse_tool_calls (output).
+    choice = body.get("tool_choice")
+    if choice is not None:
+        if isinstance(choice, str):
+            if choice not in ("auto", "none", "required"):
+                raise APIError(400, "`tool_choice` must be one of \"auto\", \"none\", \"required\", "
+                                    "or a function object.", "tool_choice", "unsupported_value")
+        elif isinstance(choice, dict):
+            name = (choice.get("function") or {}).get("name") or choice.get("name")
+            if not name:
+                raise APIError(400, "`tool_choice` function object must include a name.",
+                               "tool_choice", "invalid_value")
+            declared = [(t.get("function", t) if isinstance(t, dict) else {}).get("name")
+                        for t in (body.get("tools") or body.get("functions") or [])]
+            if name not in declared:
+                raise APIError(400, f"`tool_choice` names {name!r}, which is not in `tools`.",
+                               "tool_choice", "invalid_value")
+        else:
+            raise APIError(400, "`tool_choice` must be a string or a function object.",
+                           "tool_choice", "invalid_value")
+        if choice != "none" and not (body.get("tools") or body.get("functions")):
+            raise APIError(400, "`tool_choice` requires `tools`.", "tool_choice", "invalid_value")
     if body.get("stop") is not None:
         raise APIError(400, "Custom stop sequences are not supported yet.", "stop", "unsupported_parameter")
     if body.get("logprobs"):
@@ -678,6 +752,8 @@ class APIHandler(BaseHTTPRequestHandler):
     def generation(self, body, prompt, request_id, chat):
         maximum, temperature, top_p = generation_options(body, self.server.max_tokens)
         tools = (body.get("tools") or body.get("functions") or None) if chat else None
+        if body.get("tool_choice") == "none":
+            tools = None          # client forbade tools: never surface tool_calls
         cache_slot = body.get("cache_slot")
         if (cache_slot is not None and
                 (isinstance(cache_slot, bool) or not isinstance(cache_slot, int) or
@@ -878,7 +954,8 @@ class APIHandler(BaseHTTPRequestHandler):
         if not isinstance(enable_thinking, bool):
             raise APIError(400, "`enable_thinking` must be a boolean.", "enable_thinking")
         tools = body.get("tools") or body.get("functions") or None
-        prompt = render_chat(body.get("messages"), enable_thinking, reasoning_effort, tools)
+        prompt = render_chat(body.get("messages"), enable_thinking, reasoning_effort, tools,
+                             body.get("tool_choice"))
         self.generation(body, prompt, request_id, True)
 
     def completion(self, body, request_id):

c/tests/test_openai_server.py

@@ -9,8 +9,8 @@ from urllib.error import HTTPError
 from urllib.request import Request, urlopen
 
 from openai_server import (APIError, APIServer, ClientCancelled, END, GenerationScheduler,
-                           READY, Engine, generation_options, read_engine_turn, render_chat,
+                           READY, Engine, generation_options, parse_tool_calls,
-                           serve)
+                           read_engine_turn, render_chat, serve)
 
 
 class FakeEngine:
@@ -539,5 +539,82 @@ class SchedulerHTTPTest(unittest.TestCase):
         self.assertEqual(first_errors, [])
 
 
+
+ORDER_TOOL = [{"type": "function", "function": {
+    "name": "lookup_order",
+    "parameters": {"type": "object", "properties": {
+        "order_id": {"type": "string"},
+        "qty": {"type": "integer"},
+        "express": {"type": "boolean"},
+    }, "required": ["order_id"]}}}]
+
+
+class ToolArgumentTypeTest(unittest.TestCase):
+    """The model emits every argument as text. Without the schema, a string-typed value that
+    happens to look numeric is json.loads()'d into an int and the tool gets the wrong type."""
+
+    def _args(self, reply, tools=ORDER_TOOL):
+        _, calls = parse_tool_calls(reply, tools)
+        self.assertEqual(len(calls), 1)
+        return json.loads(calls[0]["function"]["arguments"])
+
+    def test_string_parameter_holding_digits_stays_a_string(self):
+        args = self._args("<tool_call>lookup_order"
+                          "<arg_key>order_id</arg_key><arg_value>12345</arg_value></tool_call>")
+        self.assertEqual(args["order_id"], "12345")
+        self.assertIsInstance(args["order_id"], str)
+
+    def test_declared_numeric_and_boolean_parameters_are_decoded(self):
+        args = self._args("<tool_call>lookup_order"
+                          "<arg_key>order_id</arg_key><arg_value>A-1</arg_value>"
+                          "<arg_key>qty</arg_key><arg_value>2</arg_value>"
+                          "<arg_key>express</arg_key><arg_value>true</arg_value></tool_call>")
+        self.assertEqual(args, {"order_id": "A-1", "qty": 2, "express": True})
+        self.assertIsInstance(args["qty"], int)
+        self.assertIs(args["express"], True)
+
+    def test_unknown_parameter_keeps_permissive_decoding(self):
+        args = self._args("<tool_call>lookup_order"
+                          "<arg_key>extra</arg_key><arg_value>7</arg_value></tool_call>")
+        self.assertEqual(args["extra"], 7)
+
+
+class ToolChoiceTest(unittest.TestCase):
+    def test_none_does_not_offer_the_tools(self):
+        prompt = render_chat([{"role": "user", "content": "hi"}], tools=ORDER_TOOL,
+                             tool_choice="none")
+        self.assertNotIn("<tools>", prompt)
+
+    def test_auto_offers_the_tools(self):
+        prompt = render_chat([{"role": "user", "content": "hi"}], tools=ORDER_TOOL,
+                             tool_choice="auto")
+        self.assertIn("<tools>", prompt)
+
+    def test_required_instructs_the_model_to_call_one(self):
+        prompt = render_chat([{"role": "user", "content": "hi"}], tools=ORDER_TOOL,
+                             tool_choice="required")
+        self.assertIn("<tools>", prompt)
+        self.assertIn("must call one of the functions", prompt)
+
+    def test_named_function_restricts_to_that_function(self):
+        tools = ORDER_TOOL + [{"type": "function", "function": {"name": "other", "parameters": {}}}]
+        prompt = render_chat([{"role": "user", "content": "hi"}], tools=tools,
+                             tool_choice={"type": "function", "function": {"name": "lookup_order"}})
+        self.assertIn("must call the function `lookup_order`", prompt)
+        self.assertNotIn('"other"', prompt)
+
+    def test_rejects_unknown_string_and_unknown_function(self):
+        with self.assertRaises(APIError):
+            generation_options({"messages": [], "tools": ORDER_TOOL, "tool_choice": "maybe"}, 128)
+        with self.assertRaises(APIError):
+            generation_options({"messages": [], "tools": ORDER_TOOL,
+                                "tool_choice": {"type": "function",
+                                                "function": {"name": "nope"}}}, 128)
+
+    def test_rejects_tool_choice_without_tools(self):
+        with self.assertRaises(APIError):
+            generation_options({"messages": [], "tool_choice": "required"}, 128)
+
+
 if __name__ == "__main__":
     unittest.main()