diff --git a/c/openai_server.py b/c/openai_server.py index 60e0c11..0c26b78 100644 --- a/c/openai_server.py +++ b/c/openai_server.py @@ -151,7 +151,87 @@ def content_text(content, param): return "".join(parts) -def render_chat(messages, enable_thinking=False, reasoning_effort=None): +# ---- GLM-5.2 tool calling ----------------------------------------------------------------- +# The model expresses tool calls as ordinary text (from chat_template.jinja): +# {name}{k}{v}... +# and tool results come back as <|observation|>{content}. +# We render those markers into the prompt and parse them back into OpenAI `tool_calls`. +import re + +BOX_START, BOX_END = "", "" +TR_OPEN, TR_CLOSE = "", "" +THINK_OPEN, THINK_CLOSE = "", "" + +_BOX_RE = re.compile(re.escape(BOX_START) + r"(.*?)" + re.escape(BOX_END), re.DOTALL) +_ARG_RE = re.compile(r"([^<]*)(.*?)", re.DOTALL) +_NAME_RE = re.compile(r"\s*([A-Za-z0-9_.\-]+)") +_TAG_RE = re.compile(r"|") + +# De-mangler: opt-in recovery for heavily-quantized models that drop the +# K structure. Default OFF (never rewrites well-formed output). +_SALVAGE = os.environ.get("COLI_TOOL_SALVAGE", "0") == "1" + + +def _tool_param_order(tools): + """name -> ordered param names (required first) from the request schema, for de-mangling.""" + out = {} + for tool in (tools or []): + fn = tool.get("function", tool) if isinstance(tool, dict) else {} + name = fn.get("name") + if not name: + continue + params = ((fn.get("parameters") or {}).get("properties") or {}) + required = list((fn.get("parameters") or {}).get("required") or []) + out[name] = required + [p for p in params if p not in required] + return out + + +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) + 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 = {} + 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 + if not args and _SALVAGE: + rest = inner[name_match.end():] if name_match else "" + payload = _TAG_RE.sub("", rest).strip() + if payload.startswith("(") and payload.endswith(")"): + payload = payload[1:-1].strip() + if payload: + key = (param_order.get(name) or ["input"])[0] + try: + payload = json.loads(payload) + except (json.JSONDecodeError, TypeError, ValueError): + pass + args = {key: payload} + salvaged.append(name) + calls.append({"id": "call_" + uuid.uuid4().hex[:24], "type": "function", + "function": {"name": name, "arguments": json.dumps(args, ensure_ascii=False)}}) + text = _BOX_RE.sub("", reply) + if THINK_CLOSE in text: + text = text.split(THINK_CLOSE, 1)[1] + text = text.replace(THINK_OPEN, "").replace(THINK_CLOSE, "") + if calls: + dm = len(salvaged) + sys.stderr.write("[api] tool-calls: %d total, %d strict, %d de-mangled [%s]%s\n" + % (len(calls), len(calls) - dm, dm, "CLEAN" if dm == 0 else "DE-MANGLED", + (" -> " + ", ".join(salvaged)) if dm else "")) + sys.stderr.flush() + return text.strip(), calls + + +def render_chat(messages, enable_thinking=False, reasoning_effort=None, tools=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") @@ -159,20 +239,57 @@ def render_chat(messages, enable_thinking=False, reasoning_effort=None): if enable_thinking: effort = "High" if reasoning_effort == "high" else "Max" prompt.append(f"<|system|>Reasoning Effort: {effort}") + 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 + # preamble makes it hallucinate other frameworks' syntax (e.g. `end_action`). + prompt.append("<|system|>\n# Tools\n\nYou may call one or more functions to assist with the " + "user query.\n\nYou are provided with function signatures within " + "XML tags:\n\n") + for tool in tools: + fn = tool.get("function", tool) if isinstance(tool, dict) else {} + clean = {k: v for k, v in fn.items() if k not in ("defer_loading", "strict")} + prompt.append(json.dumps(clean, ensure_ascii=False) + "\n") + prompt.append("\n\nFor each function call, output the function name and arguments " + "within the following XML format:\n{function-name}" + "{arg-key-1}{arg-value-1}" + "{arg-key-2}{arg-value-2}...") + prev_tool = False for index, message in enumerate(messages): if not isinstance(message, dict): raise APIError(400, "Each message must be an object.", f"messages.{index}") role = message.get("role") - text = content_text(message.get("content"), f"messages.{index}.content") if role in ("system", "developer"): - prompt.append(f"<|system|>{text}") + prompt.append(f"<|system|>{content_text(message.get('content'), f'messages.{index}.content')}") elif role == "user": - prompt.append(f"<|user|>{text}") + prompt.append(f"<|user|>{content_text(message.get('content'), f'messages.{index}.content')}") elif role == "assistant": + # content may be null when the message is purely tool_calls + raw = message.get("content") + text = content_text(raw, f"messages.{index}.content") if raw is not None else "" prompt.append(f"<|assistant|>{text.strip()}") + for tc in (message.get("tool_calls") or []): + fn = tc.get("function", tc) if isinstance(tc, dict) else {} + args = fn.get("arguments", "{}") + if isinstance(args, str): + try: + args = json.loads(args) + except (json.JSONDecodeError, TypeError): + args = {} + prompt.append(BOX_START + (fn.get("name") or "")) + for key, value in (args or {}).items(): + prompt.append(f"{key}" + + (value if isinstance(value, str) + else json.dumps(value, ensure_ascii=False)) + "") + prompt.append(BOX_END) + elif role == "tool": + if not prev_tool: # one <|observation|> per consecutive tool run + prompt.append("<|observation|>") + prompt.append(TR_OPEN + content_text(message.get("content"), f"messages.{index}.content") + TR_CLOSE) else: raise APIError(400, f"Unsupported message role: {role!r}.", f"messages.{index}.role", "unsupported_role") + prev_tool = (role == "tool") prompt.append("<|assistant|>" if enable_thinking else "<|assistant|>") return "".join(prompt) @@ -181,9 +298,7 @@ def render_chat(messages, enable_thinking=False, reasoning_effort=None): def generation_options(body, limit): if body.get("n", 1) != 1: raise APIError(400, "Colibri currently supports `n=1` only.", "n", "unsupported_value") - for name in ("tools", "functions"): - if body.get(name): - raise APIError(400, f"`{name}` is not supported yet.", name, "unsupported_parameter") + # `tools`/`functions` are handled by render_chat (declaration) + parse_tool_calls (output). if body.get("stop") is not None: raise APIError(400, "Custom stop sequences are not supported yet.", "stop", "unsupported_parameter") if body.get("logprobs"): @@ -427,6 +542,7 @@ 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 cache_slot = body.get("cache_slot", 0) if isinstance(cache_slot, bool) or not isinstance(cache_slot, int) or not 0 <= cache_slot < self.server.kv_slots: raise APIError(400, f"`cache_slot` must be an integer between 0 and {self.server.kv_slots - 1}.", @@ -450,10 +566,18 @@ class APIHandler(BaseHTTPRequestHandler): stats = self.server.engine.generate( prompt, maximum, temperature, top_p, output.append, cache_slot) text = "".join(output) - finish = "length" if stats["length_limited"] else "stop" - choice = ({"index": 0, "message": {"role": "assistant", "content": text, - "refusal": None}, "logprobs": None, "finish_reason": finish} if chat else - {"index": 0, "text": text, "logprobs": None, "finish_reason": finish}) + length_finish = "length" if stats["length_limited"] else "stop" + if chat and tools: + content, calls = parse_tool_calls(text, tools) + message = {"role": "assistant", "content": content or None, "refusal": None} + if calls: + message["tool_calls"] = calls + finish = "tool_calls" if calls else length_finish + choice = {"index": 0, "message": message, "logprobs": None, "finish_reason": finish} + else: + choice = ({"index": 0, "message": {"role": "assistant", "content": text, + "refusal": None}, "logprobs": None, "finish_reason": length_finish} if chat else + {"index": 0, "text": text, "logprobs": None, "finish_reason": length_finish}) self.send_json(200, {"id": completion_id, "object": object_name, "created": created, "model": self.server.model_id, "choices": [choice], "usage": self.usage(stats)}, request_id, queue_headers) @@ -469,6 +593,18 @@ class APIHandler(BaseHTTPRequestHandler): self.send_cors_headers() self.end_headers() connected = True + # KEEPALIVE: engine.generate() blocks SILENTLY during the (minutes-long) cold + # prefill, and the client drops the socket after its idle timeout. A background pump + # emits a reasoning_content "." delta (the channel that reliably resets the client's + # timer and lands in the thinking panel, so answer content stays clean) whenever no + # event has been written for KA_GAP seconds. All wfile writes share ka_lock so the + # pump and event() never interleave; last_write gates the pump so it stays quiet + # while real tokens are flowing (e.g. during decode). + ka_lock = threading.Lock() + last_write = [time.time()] + ka_stop = threading.Event() + KA_GAP = 10.0 + dbg_echo = os.environ.get("COLI_DEBUG", "0") == "1" # tee decoded tokens to stderr def event(choices, usage_marker=False): nonlocal connected @@ -478,12 +614,23 @@ class APIHandler(BaseHTTPRequestHandler): "model": self.server.model_id, "choices": choices} if include_usage: event_body["usage"] = None if not usage_marker else usage_marker - try: - data = json.dumps(event_body, ensure_ascii=False, separators=(",", ":")) - self.wfile.write(f"data: {data}\n\n".encode()) - self.wfile.flush() - except OSError: - connected = False + data = json.dumps(event_body, ensure_ascii=False, separators=(",", ":")) + with ka_lock: + try: + self.wfile.write(f"data: {data}\n\n".encode()) + self.wfile.flush() + last_write[0] = time.time() + except OSError: + connected = False + + def _keepalive(): + ping = [{"index": 0, "delta": ({"reasoning_content": "."} if chat else {"content": ""}), + "logprobs": None, "finish_reason": None}] + while not ka_stop.wait(1.0): + if not connected: + return + if time.time() - last_write[0] >= KA_GAP: + event(ping) if chat: event([{"index": 0, "delta": {"role": "assistant", "content": ""}, @@ -495,9 +642,54 @@ class APIHandler(BaseHTTPRequestHandler): {"index": 0, "text": text, "logprobs": None, "finish_reason": None}) event([choice]) - stats = self.server.engine.generate( - prompt, maximum, temperature, top_p, emit, cache_slot) - finish = "length" if stats["length_limited"] else "stop" + ka_thread = threading.Thread(target=_keepalive, daemon=True) + ka_thread.start() + if chat and tools: + # Suppress tool-call markers from the streamed content and parse the authoritative + # calls from the FULL reply after generation. Hold back a marker-length tail so a + # split across engine chunks is still caught. + sp = {"buf": "", "tool": False} + hold = len(BOX_START) - 1 + raw = [] + def emit_tools(chunk): + raw.append(chunk) + if dbg_echo: + sys.stderr.write(chunk); sys.stderr.flush() + if sp["tool"]: + return + sp["buf"] += chunk + cut = sp["buf"].find(BOX_START) + if cut >= 0: + if cut: + emit(sp["buf"][:cut]) + sp["buf"] = "" + sp["tool"] = True + return + flush = max(0, len(sp["buf"]) - hold) + if flush: + emit(sp["buf"][:flush]) + sp["buf"] = sp["buf"][flush:] + stats = self.server.engine.generate( + prompt, maximum, temperature, top_p, emit_tools, cache_slot) + if not sp["tool"] and sp["buf"]: + emit(sp["buf"]) # no tool call happened: flush held tail + _content, calls = parse_tool_calls("".join(raw), tools) + for i, tc in enumerate(calls): + event([{"index": 0, "delta": {"tool_calls": [{"index": i, "id": tc["id"], + "type": "function", "function": {"name": tc["function"]["name"], + "arguments": tc["function"]["arguments"]}}]}, + "logprobs": None, "finish_reason": None}]) + finish = "tool_calls" if calls else ("length" if stats["length_limited"] else "stop") + else: + def emit_plain(chunk): + if dbg_echo: + sys.stderr.write(chunk); sys.stderr.flush() + emit(chunk) + stats = self.server.engine.generate( + prompt, maximum, temperature, top_p, emit_plain, cache_slot) + finish = "length" if stats["length_limited"] else "stop" + ka_stop.set() # generation done: stop the keepalive pump + ka_thread.join(timeout=2) final_choice = ({"index": 0, "delta": {}, "logprobs": None, "finish_reason": finish} if chat else {"index": 0, "text": "", "logprobs": None, "finish_reason": finish}) @@ -535,10 +727,17 @@ class APIHandler(BaseHTTPRequestHandler): if reasoning_effort not in efforts: raise APIError(400, "`reasoning_effort` must be none, minimal, low, medium, high, or xhigh.", "reasoning_effort") + # COLI_THINK=1 makes thinking the default when the client sends NEITHER reasoning_effort + # nor enable_thinking (a global switch, like the old server's --think). An explicit + # client value always wins. Default off => exact OpenAI-standard behavior. + if (reasoning_effort is None and "enable_thinking" not in body + and os.environ.get("COLI_THINK", "0") == "1"): + reasoning_effort = "high" enable_thinking = body.get("enable_thinking", reasoning_effort not in (None, "none")) if not isinstance(enable_thinking, bool): raise APIError(400, "`enable_thinking` must be a boolean.", "enable_thinking") - prompt = render_chat(body.get("messages"), enable_thinking, reasoning_effort) + tools = body.get("tools") or body.get("functions") or None + prompt = render_chat(body.get("messages"), enable_thinking, reasoning_effort, tools) self.generation(body, prompt, request_id, True) def completion(self, body, request_id):