diff --git a/c/openai_server.py b/c/openai_server.py index 47672d0..8eba583 100644 --- a/c/openai_server.py +++ b/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 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: - value = json.loads(value) - except (json.JSONDecodeError, TypeError): - pass - args[key] = value + args[key] = _coerce_arg(value, types.get(key)) 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` + 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{function-name}" "{arg-key-1}{arg-value-1}" "{arg-key-2}{arg-value-2}...") + 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): diff --git a/c/tests/test_openai_server.py b/c/tests/test_openai_server.py index 9792403..fd184d6 100644 --- a/c/tests/test_openai_server.py +++ b/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, - serve) + READY, Engine, generation_options, parse_tool_calls, + 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("lookup_order" + "order_id12345") + 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("lookup_order" + "order_idA-1" + "qty2" + "expresstrue") + 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("lookup_order" + "extra7") + 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("", prompt) + + def test_auto_offers_the_tools(self): + prompt = render_chat([{"role": "user", "content": "hi"}], tools=ORDER_TOOL, + tool_choice="auto") + self.assertIn("", 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("", 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()