diff --git a/c/openai_server.py b/c/openai_server.py index 91373088..e5b87c29 100644 --- a/c/openai_server.py +++ b/c/openai_server.py @@ -614,11 +614,14 @@ class ThinkingStreamSplit: """Split GLM's reasoning marker without leaking markers across stream chunks.""" MARKERS = (THINK_OPEN, THINK_CLOSE) - def __init__(self, on_thinking, on_text, on_thinking_end=None): + def __init__(self, on_thinking, on_text, on_thinking_end=None, initial_thinking=True): self.on_thinking = on_thinking self.on_text = on_text self.on_thinking_end = on_thinking_end - self.thinking = True + # #597: GLM emits reasoning only when the prompt opened (thinking on); + # with thinking off the prompt already closed it, so output is pure answer and + # the splitter must start in text mode or it would file the whole answer as reasoning. + self.thinking = initial_thinking self.buf = "" def _emit(self, text): @@ -654,11 +657,13 @@ def finish(self): self._emit(self.buf) self.buf = "" + close = finish # interface parity with InklingStreamSplit in the streaming path -def split_thinking_reply(text): + +def split_thinking_reply(text, enable_thinking=True): """Return the marker-free (thinking, answer) portions of one GLM reply.""" thinking, answer = [], [] - split = ThinkingStreamSplit(thinking.append, answer.append) + split = ThinkingStreamSplit(thinking.append, answer.append, initial_thinking=enable_thinking) split.feed(text) split.finish() return "".join(thinking), "".join(answer) @@ -1546,7 +1551,8 @@ def error_body(self, error): return error_object(error) return {"type": "error", "error": {"type": error.error_type, "message": error.message}} - def generation(self, body, prompt, request_id, chat, tools=None, tool_choice=None): + def generation(self, body, prompt, request_id, chat, tools=None, tool_choice=None, + enable_thinking=False): # COLI_DEBUG tees the engine transaction to stderr: 1 = decoded output stream only, # 2 = both sides (rendered prompt + output). render_chat already folds prior turns and # tool results into `prompt`, so level 2 is the full conversation the engine saw. @@ -1600,6 +1606,11 @@ def generation(self, body, prompt, request_id, chat, tools=None, tool_choice=Non reasoning = "" if ARCH == "inkling": text, reasoning = split_inkling(text) + elif chat: + # #597 item 4: GLM emits reasoning then then the answer. Route the + # reasoning to reasoning_content instead of dumping it (or the raw ) + # into the visible answer / tool-call parser. + reasoning, text = split_thinking_reply(text, enable_thinking) length_finish = "length" if stats["length_limited"] else "stop" if chat and tools: content, calls = parse_tool_calls(text, tools) @@ -1694,6 +1705,10 @@ def emit_reasoning(text): # thinking → reasoning_content deltas (chat only splitter = (InklingStreamSplit(emit, emit_reasoning if chat else None) if ARCH == "inkling" else None) + # #597 item 4: GLM (chat) streams reasoning then then the answer. Split the + # reasoning into reasoning_content deltas instead of leaking it — and the raw — + # into visible content or the tool-call buffer. + glm_think = chat and ARCH != "inkling" ka_thread = threading.Thread(target=_keepalive, daemon=True) ka_thread.start() @@ -1704,10 +1719,8 @@ def emit_reasoning(text): # thinking → reasoning_content deltas (chat only sp = {"buf": "", "tool": False} hold = len(BOX_START) - 1 raw = [] - def emit_tools(chunk): + def feed_content(chunk): # answer text only (post-) raw.append(chunk) - if dbg_echo: - sys.stderr.write(chunk); sys.stderr.flush() if sp["tool"]: return sp["buf"] += chunk @@ -1722,11 +1735,22 @@ def emit_tools(chunk): if flush: emit(sp["buf"][:flush]) sp["buf"] = sp["buf"][flush:] + # #597: keep GLM reasoning out of the tool-call buffer — a think splitter sends it + # to reasoning_content and passes only the answer text on to feed_content/parser. + think = (ThinkingStreamSplit(emit_reasoning, feed_content, + initial_thinking=enable_thinking) + if glm_think else None) + def emit_tools(chunk): + if dbg_echo: + sys.stderr.write(chunk); sys.stderr.flush() + (think.feed if think else feed_content)(chunk) stop_filter = StopFilter(stop_sequences, emit_tools, ignore_leading_stop) stats = self.server.engine.generate( prompt, maximum, temperature, top_p, stop_filter.feed, cache_slot, lambda: not connected, grammar=grammar, stopped=stop_filter.stopped) stop_filter.finish() + if think: + think.finish() 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) @@ -1737,17 +1761,24 @@ def emit_tools(chunk): "logprobs": None, "finish_reason": None}]) finish = "tool_calls" if calls else ("length" if stats["length_limited"] else "stop") else: + if splitter is not None: # inkling content/marker splitter + content_split = splitter + elif glm_think: # GLM reasoning → reasoning_content + content_split = ThinkingStreamSplit(emit_reasoning, emit, + initial_thinking=enable_thinking) + else: + content_split = None def emit_plain(chunk): if dbg_echo: sys.stderr.write(chunk); sys.stderr.flush() - (splitter.feed if splitter else emit)(chunk) + (content_split.feed if content_split else emit)(chunk) stop_filter = StopFilter(stop_sequences, emit_plain, ignore_leading_stop) stats = self.server.engine.generate( prompt, maximum, temperature, top_p, stop_filter.feed, cache_slot, lambda: not connected, grammar=grammar, stopped=stop_filter.stopped) stop_filter.finish() - if splitter: - splitter.close() + if content_split: + content_split.close() finish = "length" if stats["length_limited"] else "stop" ka_stop.set() # generation done: stop the keepalive pump ka_thread.join(timeout=2) @@ -1803,7 +1834,8 @@ def chat_completion(self, body, request_id): renderer = render_chat_inkling if ARCH == "inkling" else render_chat prompt = renderer(body.get("messages"), enable_thinking, reasoning_effort, tools, tool_choice) - self.generation(body, prompt, request_id, True, tools, tool_choice) + self.generation(body, prompt, request_id, True, tools, tool_choice, + enable_thinking=enable_thinking) # ---- Anthropic /v1/messages (#343) ---------------------------------------------------- ANTHROPIC_STOP = {"stop": "end_turn", "length": "max_tokens", "tool_calls": "tool_use"} diff --git a/c/tests/test_openai_server.py b/c/tests/test_openai_server.py index 3bba88d0..ab92acb4 100644 --- a/c/tests/test_openai_server.py +++ b/c/tests/test_openai_server.py @@ -13,9 +13,9 @@ from openai_server import (APIError, APIHandler, APIServer, ClientCancelled, DEFAULT_CHAT_STOP_SEQUENCES, END, GenerationScheduler, - READY, Engine, InklingStreamSplit, StopFilter, _engine_error, - generation_options, parse_tool_calls, read_engine_turn, render_chat, - serve, stop_policy) + READY, Engine, InklingStreamSplit, StopFilter, ThinkingStreamSplit, + _engine_error, generation_options, parse_tool_calls, read_engine_turn, + render_chat, serve, split_thinking_reply, stop_policy) class FakeEngine: @@ -1005,5 +1005,145 @@ def test_untrusted_host_still_rejected_with_allowlist(self): self.assertEqual(self._get_models(server.server_port, "evil.example.com"), 403) +class ThinkingSplitUnitTest(unittest.TestCase): + """#597 item 4: the GLM reasoning splitter, incl. mkelcb's cross-chunk cases.""" + + def test_single_chunk(self): + self.assertEqual(split_thinking_reply("abcdef"), ("abc", "def")) + + def test_close_tag_split_across_chunks(self): + thinking, answer = [], [] + s = ThinkingStreamSplit(thinking.append, answer.append) + s.feed("abcdef"); s.finish() + self.assertEqual(("".join(thinking), "".join(answer)), ("abc", "def")) + + def test_open_tag_split_and_stray_open_marker(self): + thinking, answer = [], [] + s = ThinkingStreamSplit(thinking.append, answer.append) + s.feed("abcdefghi"); s.finish() + self.assertEqual(("".join(thinking), "".join(answer)), ("abcdef", "ghi")) + + def test_thinking_disabled_is_all_answer(self): + # initial_thinking=False: a pure answer with no markers must not be filed as reasoning + self.assertEqual(split_thinking_reply("plain answer", enable_thinking=False), + ("", "plain answer")) + + def test_missing_close_tag_surfaces_reasoning(self): + self.assertEqual(split_thinking_reply("thought with no end"), + ("thought with no end", "")) + + +class _ChunkEngine(FakeEngine): + """Engine that emits a caller-supplied chunk sequence, to exercise the streaming + reasoning splitter across arbitrary chunk boundaries.""" + def __init__(self, chunks): + super().__init__() + self.chunks = list(chunks) + + def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0, + cancelled=None, grammar=None, stopped=None): + self.calls.append((prompt, maximum, temperature, top_p, cache_slot, grammar)) + for chunk in self.chunks: + on_text(chunk) + if stopped and stopped(): + self.stop_requests += 1 + break + return {"prompt_tokens": 7, "completion_tokens": len(self.chunks), "length_limited": False} + + +class GlmReasoningStreamTest(unittest.TestCase): + """#597 item 4 end-to-end: GLM reasoning streams as reasoning_content, the answer as + content, no / leaks, cross-chunk-safe, and reasoning never contaminates + the tool-call buffer.""" + + def _server(self, chunks): + server = APIServer(("127.0.0.1", 0), _ChunkEngine(chunks), "test-model") + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + self.addCleanup(thread.join, 2) + self.addCleanup(server.server_close) + self.addCleanup(server.shutdown) + self.addCleanup(server.scheduler.close) + return f"http://127.0.0.1:{server.server_port}" + + def _post(self, base, body): + req = Request(base + "/v1/chat/completions", + data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}) + with urlopen(req, timeout=3) as response: + return response.read().decode() + + def _deltas(self, raw): + reasoning, content, tool_calls = [], [], [] + for line in raw.splitlines(): + if not line.startswith("data: ") or line == "data: [DONE]": + continue + for choice in json.loads(line[6:])["choices"]: + delta = choice.get("delta") or {} + if delta.get("reasoning_content"): + reasoning.append(delta["reasoning_content"]) + if delta.get("content"): + content.append(delta["content"]) + if delta.get("tool_calls"): + tool_calls.extend(delta["tool_calls"]) + return "".join(reasoning), "".join(content), tool_calls + + def test_streaming_splits_reasoning_from_answer(self): + base = self._server(["I think ", "42", "", "The answer ", "is 42"]) + raw = self._post(base, {"model": "test-model", "stream": True, "enable_thinking": True, + "messages": [{"role": "user", "content": "2+2?"}]}) + reasoning, content, _ = self._deltas(raw) + self.assertEqual(reasoning, "I think 42") + self.assertEqual(content, "The answer is 42") + self.assertNotIn("", raw) + self.assertNotIn("", raw) + + def test_streaming_close_tag_split_across_chunks(self): + base = self._server(["reasonans", "wer"]) + raw = self._post(base, {"model": "test-model", "stream": True, "enable_thinking": True, + "messages": [{"role": "user", "content": "x"}]}) + reasoning, content, _ = self._deltas(raw) + self.assertEqual(reasoning, "reason") + self.assertEqual(content, "answer") + self.assertNotIn("think>", raw) + + def test_streaming_thinking_off_is_all_content(self): + base = self._server(["Just ", "the answer"]) + raw = self._post(base, {"model": "test-model", "stream": True, "enable_thinking": False, + "messages": [{"role": "user", "content": "x"}]}) + reasoning, content, _ = self._deltas(raw) + self.assertEqual(reasoning, "") + self.assertEqual(content, "Just the answer") + + def test_streaming_reasoning_stays_out_of_tool_call(self): + base = self._server(["deciding to call", "", + "get_weathercity" + "Paris"]) + raw = self._post(base, {"model": "test-model", "stream": True, "enable_thinking": True, + "messages": [{"role": "user", "content": "weather?"}], + "tools": [{"type": "function", "function": { + "name": "get_weather", "parameters": {"type": "object", + "properties": {"city": {"type": "string"}}}}}]}) + reasoning, content, tool_calls = self._deltas(raw) + self.assertEqual(reasoning, "deciding to call") + self.assertTrue(tool_calls, "expected a parsed tool call") + args = tool_calls[0]["function"]["arguments"] + self.assertIn("Paris", args) + self.assertNotIn("deciding", args) # reasoning must not leak into the tool arguments + self.assertNotIn("deciding", content) # nor into the visible answer + + def test_nonstreaming_splits_reasoning(self): + base = self._server(["mulling ", "it over", "", "final ", "answer"]) + req = Request(base + "/v1/chat/completions", + data=json.dumps({"model": "test-model", "enable_thinking": True, + "messages": [{"role": "user", "content": "x"}]}).encode(), + headers={"Content-Type": "application/json"}) + with urlopen(req, timeout=3) as response: + body = json.load(response) + message = body["choices"][0]["message"] + self.assertEqual(message["reasoning_content"], "mulling it over") + self.assertEqual(message["content"], "final answer") + + if __name__ == "__main__": unittest.main()