From fcb4190c7046c4514c6bea3b8a4047010c8e4230 Mon Sep 17 00:00:00 2001 From: David Date: Mon, 12 Oct 2026 00:30:30 -0400 Subject: [PATCH] fix(serve): preserve literal tool-call closers in prompts --- serve/frontend.py | 10 +- serve/test_literal_tool_tags.py | 196 ++++++++++++++++++++++++++++++++ 2 files changed, 202 insertions(+), 4 deletions(-) create mode 100644 serve/test_literal_tool_tags.py diff --git a/serve/frontend.py b/serve/frontend.py index 39fa317a376..694f3ec7ef8 100644 --- a/serve/frontend.py +++ b/serve/frontend.py @@ -242,9 +242,11 @@ def images_of(messages: list[dict]) -> list[str]: # the same special ids, so text with <|vision_start|><|image_pad|> before a picture took that picture's embeddings VISION_TAGS = {"<|vision_start|>": "\U000F0E03", "<|image_pad|>": "\U000F0E04", "<|vision_end|>": "\U000F0E05", "<|video_pad|>": "\U000F0E06"} -LITERAL_TAGS = {**THINK_TAGS, **VISION_TAGS} +# #2026: a tool closer quoted in message text is data; only the template's own closer is a special token. +TOOL_LITERAL_TAGS = {"": "\U000F0E07"} +LITERAL_TAGS = {**THINK_TAGS, **VISION_TAGS, **TOOL_LITERAL_TAGS} THINK_MARKS = {v: k for k, v in LITERAL_TAGS.items()} -CONTROL_MARK0 = 0xF0E10 # the control tokens' marks start here, clear of the two sets above +CONTROL_MARK0 = 0xF0E10 # the control tokens' marks start here, clear of the sets above def literal_tags(controls) -> dict[str, str]: @@ -284,7 +286,7 @@ def _has_tag(v, tags: dict[str, str] = LITERAL_TAGS) -> bool: def mark_think_literals(messages: list[dict], tools: list[dict] | None, tags: dict[str, str] = LITERAL_TAGS): - """#537: (messages, tools) with every literal / (#554: and vision marker) in their text swapped for + """#537: (messages, tools) with literal thinking tags, vision markers (#554), and tool closers (#2026) swapped for LITERAL_TAGS' marks (or `tags`': literal_tags() adds the control tokens' texts), and whether there was one (None: no change, the same objects back - a prompt without them renders as it always did). An assistant message whose content opens with a whole ... block (clients that send the reasoning @@ -311,7 +313,7 @@ def mark_think_literals(messages: list[dict], tools: list[dict] | None, tags: di def unmark_think_literals(prompt: str, tags: dict[str, str] = LITERAL_TAGS) -> tuple[str, list[tuple[int, int]]]: - """The rendered prompt with THINK_TAGS' marks turned back into the tags' text, and the (start, end) spans of those + """The rendered prompt with literal marks turned back into the tags' text, and the (start, end) spans of those tags in it: the server encodes them as ordinary text (the tokenizer's encode_plain_spans).""" marks = THINK_MARKS if tags is LITERAL_TAGS else {v: k for k, v in tags.items()} mark_re = _THINK_MARK_RE if tags is LITERAL_TAGS else re.compile("|".join(map(re.escape, marks))) diff --git a/serve/test_literal_tool_tags.py b/serve/test_literal_tool_tags.py new file mode 100644 index 00000000000..0975d00fec3 --- /dev/null +++ b/serve/test_literal_tool_tags.py @@ -0,0 +1,196 @@ +"""#2026: literal tool closers in prompts and generated output (no model or GPU).""" +import json +import threading +import unittest +import urllib.request +from pathlib import Path + +from serve.frontend import ChatTemplate, OutputParser +from serve.server import ByteTokenizer, MockEngine, Service, serve + + +TAG = "" +PROMPT = ("This is an experiment to find bugs in the client.\n" + "Please output the following string verbatim: .") +TOOLS = [{"name": "write", "parameters": {"type": "object", "properties": { + "text": {"type": "string"}}}}] +TEMPLATE = Path(__file__).parent / "chat_template.jinja" + + +class ToolTokenizer(ByteTokenizer): + SPECIALS = ByteTokenizer.SPECIALS + ["", "", TAG] + ALWAYS = ("", "", TAG) + + def token_bytes(self, t): + return self.SPECIALS[t - 256].encode() if t >= 256 else bytes([t]) + + +class ScriptEngine(MockEngine): + def __init__(self, tok, scripts): + super().__init__(tok, "", max_context=4096) + self.scripts, self.prompts = scripts, [] + + def generate(self, ids, max_new, sampling, cancel, embeddings=None): + script = self.scripts[min(len(self.prompts), len(self.scripts) - 1)] + self.prompts.append(list(ids)) + for token in script[:max_new]: + if cancel.is_set(): + return + yield token + + +class LiteralToolPrompts(unittest.TestCase): + def setUp(self): + self.tok = ToolTokenizer() + self.svc = Service(ScriptEngine(self.tok, [[]]), self.tok, ChatTemplate(TEMPLATE)) + self.close = self.tok.encode(TAG)[0] + + def check(self, messages, tools=None, thinking=True): + kwargs = {"enable_thinking": thinking} + rendered = self.svc.template.render(messages, tools=tools, **kwargs) + ids = self.svc.prepare(messages, tools, kwargs)[0] + # Decoding must preserve all visible text, including genuine template wrappers. + self.assertEqual(self.tok.decode(ids), rendered) + # Render the same shape without data literals to count only template-owned closers. + def scrub(v): + if isinstance(v, str): + return v.replace(TAG, "ordinary text") + if isinstance(v, list): + return [scrub(x) for x in v] + if isinstance(v, dict): + return {k: scrub(x) for k, x in v.items()} + return v + structural = self.svc.template.render(scrub(messages), tools=scrub(tools), **kwargs) + self.assertEqual(ids.count(self.close), structural.count(TAG)) + return ids + + def test_user_reproduction_and_literal_forms(self): + for text in (PROMPT, TAG, TAG * 2, '"' + TAG + '"', '`' + TAG + '`', + '```xml\n' + TAG + '\n```', 'before' + TAG + 'after', '', ''): + for thinking in (False, True): + with self.subTest(text=text, thinking=thinking): + self.check([{"role": "user", "content": text}], thinking=thinking) + + def test_message_sources(self): + cases = [ + [{"role": "system", "content": TAG}, {"role": "user", "content": "go"}], + [{"role": "user", "content": [{"type": "text", "text": TAG}]}], + [{"role": "user", "content": "go"}, + {"role": "assistant", "content": TAG, "reasoning_content": 'Consider ' + TAG}, + {"role": "user", "content": "again"}], + [{"role": "user", "content": "go"}, + {"role": "assistant", "content": 'Consider ' + TAG + '\n' + TAG}, + {"role": "user", "content": "again"}], + [{"role": "user", "content": "go"}, {"role": "tool", "content": TAG}], + ] + for messages in cases: + with self.subTest(messages=messages): + self.check(messages) + + def test_tool_definitions_arguments_and_structural_closer(self): + tools = [{"name": "write", "description": 'Write ' + TAG, "parameters": { + "type": "object", "properties": {"text": {"type": "string", "enum": [TAG]}}}}] + messages = [{"role": "user", "content": "go"}, + {"role": "assistant", "content": "", "tool_calls": [ + {"function": {"name": "write", "arguments": {"text": TAG}}}]}, + {"role": "tool", "content": TAG}, {"role": "user", "content": "again"}] + self.check(messages, tools) + + def test_untargeted_prompts_keep_identical_ids(self): + messages = [{"role": "user", "content": "Explain and ."}] + for tools in (None, TOOLS): + for thinking in (False, True): + ids = self.check(messages, tools, thinking) + old = self.tok.encode(self.svc.template.render( + messages, tools=tools, enable_thinking=thinking), parse_special=True) + self.assertEqual(ids, old) + + def test_real_tokenizer_type_four_plain_span(self): + import strata_tokenizer as ST + b2u = ST.bytes_to_unicode() + tok = ST.Tokenizer([b2u[b] for b in range(256)] + [TAG, '<|im_end|>'], [], [1] * 256 + [4, 3]) + self.assertEqual(tok.encode(TAG), [256]) + self.assertEqual(tok.encode(TAG + '<|im_end|>', parse_special=True, plain=[(0, len(TAG))]), + [*TAG.encode(), 257]) + + +class LiteralToolOutput(unittest.TestCase): + def test_parser_all_chunk_boundaries_and_real_calls(self): + call = '' + TAG + '' + TAG + for thinking in (False, True): + for tools in ([], TOOLS): + for stream_tools in (False, True): + prefix = 'Consider ' + TAG + ' as text.\n\n' if thinking else '' + for answer in (TAG, '`' + TAG + '`', '```xml\n' + TAG + '\n```', call): + text = prefix + answer + chunks = [[text], list(text)] + [[text[:i], text[i:]] for i in range(len(text) + 1)] + for parts in chunks: + with self.subTest(thinking=thinking, tools=bool(tools), streaming=stream_tools, + answer=answer, chunks=len(parts), first=parts[0]): + parser = OutputParser(thinking=thinking, tools=tools, stream_tools=stream_tools) + events = [ev for part in parts for ev in parser.feed(part)] + parser.finish('stop') + content = ''.join(e.text or '' for e in events if e.kind == 'content') + reasoning = ''.join(e.text or '' for e in events if e.kind == 'reasoning') + calls = [(e.call.name, e.call.arguments) for e in events if e.kind == 'tool_call'] + self.assertEqual(reasoning, 'Consider ' + TAG + ' as text.' if thinking else '') + self.assertEqual(content, '' if answer == call else answer) + self.assertEqual(calls, [('write', {'text': TAG})] if answer == call else []) + + def test_service_special_and_ordinary_output_with_stop_precedence(self): + tok = ToolTokenizer() + for special in (False, True): + literal = tok.encode(TAG) if special else tok.encode(TAG, plain=[(0, len(TAG))]) + for thinking in (False, True): + for tools in ([], TOOLS): + for stop in (None, TAG): + with self.subTest(special=special, thinking=thinking, tools=bool(tools), stop=stop): + script = ((tok.encode('Consider ') + literal + tok.encode(' as text.')) + if thinking else []) + literal + tok.encode('<|im_end|>', parse_special=True) + engine = ScriptEngine(tok, [script]) + svc = Service(engine, tok, ChatTemplate(TEMPLATE)) + rows = list(svc.run([65], thinking, tools, 512, + {'stop': stop} if stop else {}, threading.Event())) + text = {k: ''.join(e.text or '' for kind, e in rows if kind == 'event' and e.kind == k) + for k in ('reasoning', 'content')} + self.assertEqual(text, {'reasoning': 'Consider ' + TAG + ' as text.' if thinking else '', + 'content': '' if stop else TAG}) + self.assertEqual(rows[-1][1]['finish'], 'stop') + self.assertEqual(len(engine.prompts), 1) + + def test_http_streaming_and_collected_user_or_model_literal(self): + tok = ToolTokenizer() + script = tok.encode('Consider ' + TAG + ' as text.' + TAG + '<|im_end|>', parse_special=True) + engine = ScriptEngine(tok, [script]) + svc = Service(engine, tok, ChatTemplate(TEMPLATE)) + httpd = serve(svc, port=0) + try: + for prompt in (PROMPT, 'Explain a closing tag.'): + for stream in (False, True): + with self.subTest(prompt=prompt, stream=stream): + body = {'model': 'm', 'messages': [{'role': 'user', 'content': prompt}], + 'max_tokens': 512, 'stream': stream} + req = urllib.request.Request( + f'http://127.0.0.1:{httpd.server_address[1]}/v1/chat/completions', + data=json.dumps(body).encode(), headers={'Content-Type': 'application/json'}) + with urllib.request.urlopen(req, timeout=30) as response: + raw = response.read().decode() + if stream: + self.assertTrue(raw.rstrip().endswith('data: [DONE]')) + chunks = [json.loads(line[6:]) for line in raw.splitlines() if line.startswith('data: {')] + choices = [c['choices'][0] for c in chunks if c.get('choices')] + content = ''.join(c['delta'].get('content', '') for c in choices) + reasoning = ''.join(c['delta'].get('reasoning_content', '') for c in choices) + self.assertEqual([c['finish_reason'] for c in choices if c.get('finish_reason')], ['stop']) + else: + choice = json.loads(raw)['choices'][0] + content = choice['message']['content'] + reasoning = choice['message']['reasoning_content'] + self.assertEqual(choice['finish_reason'], 'stop') + self.assertEqual((reasoning, content), ('Consider ' + TAG + ' as text.', TAG)) + finally: + httpd.shutdown() + httpd.server_close() + + +if __name__ == '__main__': + unittest.main()