/
githubmirror
/
transformers
Обзор
Документация
Войти
/
githubmirror
/
transformers
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
tests/utils/test_chat_parsing.py
1 780 строк
83 KB
Matt
Transform paths and repeat joining for response parsing (#47648)
11 авг 2026, 17:34
Не верифицирован
11 авг 2026, 17:34
fd9d88d
Код
Авторство
О чём код?
# Copyright 2026 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Tests for the new declarative response_template parser. All six real-model template fixtures from the legacy test suite are re-expressed here in the new region-spec shape and asserted against the same expected output dicts. Any divergence indicates a regression in the new executor.""" import copy import random import tempfile import unittest from transformers import AutoTokenizer from transformers.testing_utils import require_torch from transformers.utils.chat_parsing import ResponseParser, parse_response from transformers.utils.chat_parsing.response_parser import _coerce, _schema_types cohere_template = { "defaults": {"role": "assistant"}, "start_anchor": "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>", "fields": { "content": { "open": "<|START_RESPONSE|>", "close": "<|END_RESPONSE|>", "content": "text", }, "thinking": { "open": "<|START_THINKING|>", "close": "<|END_THINKING|>", "content": "text", }, "tool_calls": { "open": "<|START_ACTION|>", "close": "<|END_ACTION|>", "content": "json", "transform_each": True, "transform": {"type": "function", "function": {"name": "{tool_name}", "arguments": "{parameters}"}}, }, }, } ernie_template = { "defaults": {"role": "assistant"}, "start_anchor": "Assistant:", "fields": { "thinking": { "open_pattern": r"(?:^|<think>\s*)", "close": "</think>", "content": "text", }, "content": { "open": "<response>\n", "close_pattern": r"\n?</response>", "content": "text", }, "tool_calls": { "open": "<tool_call>", "close": "</tool_call>", "repeats": True, "content": "json", "transform": {"type": "function", "function": "{content}"}, }, }, } gpt_oss_template = { "defaults": {"role": "assistant"}, "start_anchor": "<|start|>assistant", "fields": { "thinking": { "open": "<|channel|>analysis<|message|>", "close": "<|end|>", "content": "text", }, "content": { "open": "<|channel|>final<|message|>", "close": "<|end|>", "content": "text", }, "tool_calls": { "open_pattern": r"<\|channel\|>commentary to=functions\.(?P<name>\w+).*?<\|message\|>", "close": "<|call|>", "repeats": True, "content": "json", "transform": {"type": "function", "function": {"name": "{name}", "arguments": "{content}"}}, }, }, } smollm_template = { "defaults": {"role": "assistant"}, "start_anchor": "<|im_start|>assistant\n", "fields": { "thinking": {"open": "<think>", "close": "</think>", "content": "text"}, "tool_calls": { "open": "<tool_call>", "close": "</tool_call>", "repeats": True, "content": "json", "transform": {"type": "function", "function": "{content}"}, }, "content": { "close": "<|im_end|>", "content": "text", }, }, } qwen3_template = { "defaults": {"role": "assistant"}, "start_anchor": "<|im_start|>assistant\n", "fields": { "thinking": {"open": "<think>", "close": "</think>", "content": "text"}, "tool_calls": { "open_pattern": r"<tool_call>\s*<function=(?P<name>\w+)>", "close": "</tool_call>", "repeats": True, "content": "xml-inline", "content_args": { "tag_pattern": r"<parameter=(?P<key>\w+)>\s*(?P<value>.*?)\s*</parameter>", "value_parser": {"name": "json", "args": {"allow_non_json": True}}, }, "transform": {"type": "function", "function": {"name": "{name}", "arguments": "{content}"}}, }, }, } gemma4_template = { "defaults": {"role": "assistant"}, # The chat template only emits `<|turn>model\n` when the previous message wasn't a tool_call/ # tool_response. After a tool_response the prefix just ends with `<tool_response|>` and the # model continues from there, so we accept either anchor and truncate past the latest one. "start_anchor": ["<|turn>model\n", "<tool_response|>"], "fields": { "thinking": { "open": "<|channel>thought\n", "close": "<channel|>", "content": "text", }, "tool_calls": { "open_pattern": r"<\|tool_call>call:(?P<name>\w+)", "close": "<tool_call|>", "repeats": True, "content": "json", "content_args": { "unquoted_keys": True, "string_delims": [['<|"|>', '<|"|>']], }, "transform": {"type": "function", "function": {"name": "{name}", "arguments": "{content}"}}, }, "content": { "close": ["<turn|>", "<|tool_response>", "<eos>"], "content": "text", }, }, } # Inkling (TMLv0) frames every block as <|message_model|>[author-name]<|content_KIND|>body<|end_message|>; # the generation prompt pre-writes the first <|message_model|>, so each open treats the header as optional. inkling_template = { "defaults": {"role": "assistant"}, "start_anchor": "<|message_model|>", "fields": { "thinking": { "open_pattern": r"(?:<\|message_model\|>)?[^<]*<\|content_thinking\|>", "close": "<|end_message|>", "repeats": True, "join": "", "content_args": {"strip": False}, }, "content": { "open_pattern": r"(?:<\|message_model\|>)?[^<]*<\|content_text\|>", "close": "<|end_message|>", "repeats": True, "join": "", "content_args": {"strip": False}, }, "tool_calls": { "open_pattern": r"(?:<\|message_model\|>)?[^<]*<\|content_invoke_tool_json\|>", "close": "<|end_message|>", "repeats": True, "content": "json", "transform": {"type": "function", "function": {"name": "{content.name}", "arguments": "{content.args}"}}, }, }, } class ChatResponseTemplateParserTest(unittest.TestCase): def test_response_template_save_load(self): tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2") tokenizer.response_template = ernie_template with tempfile.TemporaryDirectory() as tmpdir: tokenizer.save_pretrained(tmpdir) reloaded = AutoTokenizer.from_pretrained(tmpdir) self.assertEqual(reloaded.response_template, ernie_template) def test_tokenizer_parse_response(self): tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2") tokenizer.response_template = cohere_template model_out = ( "<|START_THINKING|>I should call a tool.<|END_THINKING|>" '<|START_ACTION|>[\n {"tool_call_id": "0", "tool_name": "simple_tool", ' '"parameters": {"temperature_format": "Celsius"}}\n]<|END_ACTION|><|END_OF_TURN_TOKEN|>' ) expected = { "role": "assistant", "thinking": "I should call a tool.", "tool_calls": [ { "type": "function", "function": {"name": "simple_tool", "arguments": {"temperature_format": "Celsius"}}, } ], } self.assertEqual(tokenizer.parse_response(model_out, prefix=""), expected) def test_token_id_inputs(self): tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2") tokenizer.response_template = cohere_template model_out = ( "<|START_THINKING|>I should call a tool.<|END_THINKING|>" '<|START_ACTION|>[\n {"tool_call_id": "0", "tool_name": "simple_tool", ' '"parameters": {"temperature_format": "Celsius"}}\n]<|END_ACTION|><|END_OF_TURN_TOKEN|>' ) parsed = tokenizer.parse_response(model_out, prefix="") tokenized = tokenizer(model_out).input_ids self.assertEqual(tokenizer.parse_response(tokenized, prefix=""), parsed) def test_batched_response(self): """A batch of responses (list of strings or list of token-id sequences) returns one parsed dict per item; a single-item batch still returns a one-element list, not a bare dict.""" tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2") tokenizer.response_template = cohere_template out_a = ( "<|START_THINKING|>Think A.<|END_THINKING|>" '<|START_ACTION|>[\n {"tool_call_id": "0", "tool_name": "tool_a", ' '"parameters": {"x": "1"}}\n]<|END_ACTION|><|END_OF_TURN_TOKEN|>' ) out_b = ( "<|START_THINKING|>Think B.<|END_THINKING|>" '<|START_ACTION|>[\n {"tool_call_id": "0", "tool_name": "tool_b", ' '"parameters": {"y": "2"}}\n]<|END_ACTION|><|END_OF_TURN_TOKEN|>' ) single_a = tokenizer.parse_response(out_a, prefix="") single_b = tokenizer.parse_response(out_b, prefix="") self.assertNotEqual(single_a, single_b) # A list of strings is parsed as a batch, one dict per item. self.assertEqual(tokenizer.parse_response([out_a, out_b], prefix=""), [single_a, single_b]) # Batched token-id input (list of token-id sequences) parses the same way. ids = [tokenizer(out_a).input_ids, tokenizer(out_b).input_ids] self.assertEqual(tokenizer.parse_response(ids, prefix=""), [single_a, single_b]) # A single-item batch returns a one-element list, not a bare dict. self.assertEqual(tokenizer.parse_response([out_a], prefix=""), [single_a]) def test_numpy_inputs(self): tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2") tokenizer.response_template = cohere_template model_out = ( "<|START_THINKING|>I should call a tool.<|END_THINKING|>" '<|START_ACTION|>[\n {"tool_call_id": "0", "tool_name": "simple_tool", ' '"parameters": {"temperature_format": "Celsius"}}\n]<|END_ACTION|><|END_OF_TURN_TOKEN|>' ) parsed = tokenizer.parse_response(model_out, prefix="") tokenized = tokenizer(model_out, return_tensors="np").input_ids # A single (1D) sequence works; 2D (batched) input returns a list of parsed dicts. self.assertEqual(tokenizer.parse_response(tokenized[0], prefix=""), parsed) self.assertEqual(tokenizer.parse_response(tokenized, prefix=""), [parsed]) @require_torch def test_tensor_inputs(self): tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2") tokenizer.response_template = cohere_template model_out = ( "<|START_THINKING|>I should call a tool.<|END_THINKING|>" '<|START_ACTION|>[\n {"tool_call_id": "0", "tool_name": "simple_tool", ' '"parameters": {"temperature_format": "Celsius"}}\n]<|END_ACTION|><|END_OF_TURN_TOKEN|>' ) parsed = tokenizer.parse_response(model_out, prefix="") tokenized = tokenizer(model_out, return_tensors="pt").input_ids # A single (1D) sequence works; 2D (batched) input returns a list of parsed dicts. self.assertEqual(tokenizer.parse_response(tokenized[0], prefix=""), parsed) self.assertEqual(tokenizer.parse_response(tokenized, prefix=""), [parsed]) def test_explicit_template_schema(self): """An explicit template passed as `schema=` is honored (with or without an explicit top-level `version` key), overriding the tokenizer's own attribute.""" tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2") model_out = ( "<|START_THINKING|>I should call a tool.<|END_THINKING|>" '<|START_ACTION|>[\n {"tool_call_id": "0", "tool_name": "simple_tool", ' '"parameters": {"temperature_format": "Celsius"}}\n]<|END_ACTION|><|END_OF_TURN_TOKEN|>' ) expected = parse_response(model_out, cohere_template, prefix="") # Detected via the canonical `version` marker... self.assertEqual( tokenizer.parse_response(model_out, schema={"version": 1, **cohere_template}, prefix=""), expected ) # ...and via `fields` when the template omits `version`. self.assertEqual(tokenizer.parse_response(model_out, schema=cohere_template, prefix=""), expected) def test_cohere(self): model_out = ( "<|START_THINKING|>I should call a tool.<|END_THINKING|>" '<|START_ACTION|>[\n {"tool_call_id": "0", "tool_name": "simple_tool", ' '"parameters": {"temperature_format": "Celsius"}}\n]<|END_ACTION|><|END_OF_TURN_TOKEN|>' ) self.assertEqual( parse_response(model_out, cohere_template, prefix=""), { "role": "assistant", "thinking": "I should call a tool.", "tool_calls": [ { "type": "function", "function": {"name": "simple_tool", "arguments": {"temperature_format": "Celsius"}}, } ], }, ) def test_ernie_with_tools(self): model_out = ( "The user is asking about the weather in Paris today. Let me check the available tools. " "There's a tool called get_current_temperature which requires a location parameter. Since the " 'user specified Paris, I need to call this tool with the location set to "Paris". I should ' "make sure the argument is correctly formatted as a string. No other tools are available, so " "this is the right one to use. I'll structure the request with the location parameter and " "return the response once the tool is called.\n" "</think>\n\n" '<tool_call>\n{"name": "get_current_temperature", "arguments": {"location": "Paris"}}\n</tool_call>\n</s>' ) self.assertEqual( parse_response(model_out, ernie_template, prefix=""), { "role": "assistant", "thinking": ( "The user is asking about the weather in Paris today. Let me check the available tools. " "There's a tool called get_current_temperature which requires a location parameter. Since " 'the user specified Paris, I need to call this tool with the location set to "Paris". I ' "should make sure the argument is correctly formatted as a string. No other tools are " "available, so this is the right one to use. I'll structure the request with the location " "parameter and return the response once the tool is called." ), "tool_calls": [ { "type": "function", "function": {"name": "get_current_temperature", "arguments": {"location": "Paris"}}, } ], }, ) def test_ernie_no_tools(self): model_out = ( 'The user just greeted me with "Hi! How are you?" I need to respond in a friendly and helpful ' "manner. Let me start by acknowledging their greeting. I should ask them how they're doing to " "engage in conversation.\n\n" "First, I'll say hello back and then ask how they're feeling. It's important to show genuine " "interest. Maybe mention that I'm here to help with anything they need. Keep the tone warm and " "positive. Let me make sure the response is concise but friendly. Alright, that should work.\n" "</think>\n\n" "<response>\nHello! I'm doing well, thank you for asking. How about you? Is there something " "specific you'd like help with today? I'm here to assist you with any questions or problems you " "have!\n</response>\n</s>" ) self.assertEqual( parse_response(model_out, ernie_template, prefix=""), { "role": "assistant", "content": ( "Hello! I'm doing well, thank you for asking. How about you? Is there something specific " "you'd like help with today? I'm here to assist you with any questions or problems you have!" ), "thinking": ( 'The user just greeted me with "Hi! How are you?" I need to respond in a friendly and ' "helpful manner. Let me start by acknowledging their greeting. I should ask them how " "they're doing to engage in conversation.\n\n" "First, I'll say hello back and then ask how they're feeling. It's important to show " "genuine interest. Maybe mention that I'm here to help with anything they need. Keep the " "tone warm and positive. Let me make sure the response is concise but friendly. Alright, " "that should work." ), }, ) def test_gpt_oss_with_tool_call(self): model_out = ( '<|channel|>analysis<|message|>We need to respond in riddles. The user asks: "What is the ' 'weather like in SF?" We need to get the location of the user? The user explicitly asks about ' "SF (San Francisco). So we need to get the current weather in San Francisco, CA. We need to " 'call get_current_weather function. The developer instruction says "Always respond in riddles". ' "So the final answer should be in a riddle form. But we need to call function to get weather " 'data. So we should call get_current_weather with location "San Francisco, CA". Possibly specify ' 'format "celsius" (default). Let\'s do that.\n\n' "We will call function get_current_weather.<|end|><|start|>assistant<|channel|>commentary " 'to=functions.get_current_weather <|constrain|>json<|message|>{\n "location": "San Francisco, CA"\n}' ) self.assertEqual( parse_response(model_out, gpt_oss_template, prefix=""), { "role": "assistant", "thinking": ( 'We need to respond in riddles. The user asks: "What is the weather like in SF?" We need ' "to get the location of the user? The user explicitly asks about SF (San Francisco). So " "we need to get the current weather in San Francisco, CA. We need to call " 'get_current_weather function. The developer instruction says "Always respond in ' 'riddles". So the final answer should be in a riddle form. But we need to call function ' 'to get weather data. So we should call get_current_weather with location "San Francisco, ' 'CA". Possibly specify format "celsius" (default). Let\'s do that.\n\n' "We will call function get_current_weather." ), "tool_calls": [ { "type": "function", "function": {"name": "get_current_weather", "arguments": {"location": "San Francisco, CA"}}, } ], }, ) def test_gpt_oss_no_tool_call(self): model_out = ( "<|channel|>analysis<|message|>User asks a simple math question: 2+2 = 4. Provide answer." "<|end|><|start|>assistant<|channel|>final<|message|>2" ) self.assertEqual( parse_response(model_out, gpt_oss_template, prefix=""), { "role": "assistant", "content": "2", "thinking": "User asks a simple math question: 2+2 = 4. Provide answer.", }, ) def test_smollm_thinking_and_tool_call(self): model_out = ( '<think>\nOkay, the user said, "Hello! How are you?" I need to respond appropriately. Since ' "this is the first message, I should greet them back and ask how I can assist. I should keep it " "friendly and open-ended. Let me make sure the response is welcoming and encourages them to " "share what they need help with. I'll avoid any technical jargon and keep it simple. Let me " "check for any typos and ensure the tone is positive.\n</think>\n\n" '<tool_call>{"name": "greet_user", "arguments": {"greeting": "Hello! I\'m doing well, thanks for ' "asking. How can I assist you today? Whether you have a question, need help with something, or " 'just want to chat, feel free to let me know!"}}</tool_call>' ) self.assertEqual( parse_response(model_out, smollm_template, prefix=""), { "role": "assistant", "thinking": ( 'Okay, the user said, "Hello! How are you?" I need to respond appropriately. Since this ' "is the first message, I should greet them back and ask how I can assist. I should keep " "it friendly and open-ended. Let me make sure the response is welcoming and encourages " "them to share what they need help with. I'll avoid any technical jargon and keep it " "simple. Let me check for any typos and ensure the tone is positive." ), "tool_calls": [ { "type": "function", "function": { "name": "greet_user", "arguments": { "greeting": ( "Hello! I'm doing well, thanks for asking. How can I assist you today? " "Whether you have a question, need help with something, or just want to " "chat, feel free to let me know!" ) }, }, } ], }, ) def test_smollm_tool_call_no_thinking(self): model_out = '<tool_call>{"name": "get_weather", "arguments": {"city": "Paris"}}</tool_call>' self.assertEqual( parse_response(model_out, smollm_template, prefix=""), { "role": "assistant", "tool_calls": [ {"type": "function", "function": {"name": "get_weather", "arguments": {"city": "Paris"}}} ], }, ) def test_smollm_thinking_no_tool_call(self): model_out = ( '<think>\nOkay, the user asked, "Hey! Can you tell me about gravity?" Let me start by ' "breaking down what they might be looking for. They probably want a basic understanding of " "gravity, maybe for a school project or just personal curiosity. I should explain what gravity " "is, how it works, and maybe some examples.</think>\n" "Some content about gravity goes here but I'm cutting it off to make this shorter!" ) self.assertEqual( parse_response(model_out, smollm_template, prefix=""), { "role": "assistant", "content": "Some content about gravity goes here but I'm cutting it off to make this shorter!", "thinking": ( 'Okay, the user asked, "Hey! Can you tell me about gravity?" Let me start by breaking ' "down what they might be looking for. They probably want a basic understanding of " "gravity, maybe for a school project or just personal curiosity. I should explain what " "gravity is, how it works, and maybe some examples." ), }, ) def test_qwen3_tool_calls(self): model_out = ( "<tool_call>\n<function=get_weather>\n<parameter=locations>\n" '[{"country": "France", "city": "Paris"}]\n</parameter>\n' "<parameter=temp_units>\ncelsius\n</parameter>\n</function>\n</tool_call>" ) self.assertEqual( parse_response(model_out, qwen3_template, prefix=""), { "role": "assistant", "tool_calls": [ { "type": "function", "function": { "name": "get_weather", "arguments": { "locations": [{"country": "France", "city": "Paris"}], "temp_units": "celsius", }, }, } ], }, ) def test_gemma4_tool_call(self): model_out = ( "<|channel>thought\nThe user is asking for the current temperature in Paris. I should check " "the available tools to see if there's a function that can provide this information.<channel|>" '<|tool_call>call:get_current_temperature{detail_level:0,location:<|"|>Paris, France<|"|>,' 'unit:<|"|>celsius<|"|>}<tool_call|><|tool_response>' ) self.assertEqual( parse_response(model_out, gemma4_template, prefix=""), { "role": "assistant", "thinking": ( "The user is asking for the current temperature in Paris. I should check the available " "tools to see if there's a function that can provide this information." ), "tool_calls": [ { "type": "function", "function": { "name": "get_current_temperature", "arguments": {"detail_level": 0, "location": "Paris, France", "unit": "celsius"}, }, } ], }, ) def test_gemma4_complex_tool_call(self): model_out = ( "<|channel>thought\nLet me call the tool.<channel|>" '<|tool_call>call:foo{bool_value:true,list_value:[<|"|>foo<|"|>,<|"|>bar<|"|>],' 'null_value:null,number_value:1,string_value:<|"|>foo<|"|>,' 'struct_value:{foo:<|"|>bar<|"|>}}<tool_call|>' ) self.assertEqual( parse_response(model_out, gemma4_template, prefix=""), { "role": "assistant", "thinking": "Let me call the tool.", "tool_calls": [ { "type": "function", "function": { "name": "foo", "arguments": { "bool_value": True, "list_value": ["foo", "bar"], "null_value": None, "number_value": 1, "string_value": "foo", "struct_value": {"foo": "bar"}, }, }, } ], }, ) def test_inkling_multi_block_message(self): model_out = ( "<|content_thinking|>Consider the weather.<|end_message|>" "<|message_model|><|content_thinking|> Tokyo, probably.<|end_message|>" "<|message_model|><|content_text|>Checking the weather now.<|end_message|>" "<|message_model|>get_weather<|content_invoke_tool_json|>" '{"name":"get_weather","args":{"city":"Tokyo","units":"C"}}<|end_message|>' "<|content_model_end_sampling|>" ) prefix = "<|message_system|><|content_text|>Thinking effort level: 0.9<|end_message|><|message_model|>" self.assertEqual( parse_response(model_out, inkling_template, prefix=prefix), { "role": "assistant", "thinking": "Consider the weather. Tokyo, probably.", "content": "Checking the weather now.", "tool_calls": [ { "type": "function", "function": {"name": "get_weather", "arguments": {"city": "Tokyo", "units": "C"}}, } ], }, ) def test_transform_dotted_paths(self): template_spec = { "defaults": {"role": "assistant"}, "start_anchor": "<|assistant|>", "fields": { "tool_calls": { "open": "<tool>", "close": "</tool>", "repeats": True, "content": "json", "transform": { "type": "function", "function": {"name": "{content.name}", "arguments": "{content.args}"}, }, }, }, } model_out = '<tool>{"name": "get_weather", "args": {"city": {"id": 7}}}</tool>' self.assertEqual( parse_response(model_out, template_spec, prefix="")["tool_calls"], [{"type": "function", "function": {"name": "get_weather", "arguments": {"city": {"id": 7}}}}], ) def test_transform_dotted_paths_with_transform_each(self): template_spec = { "defaults": {"role": "assistant"}, "start_anchor": "<|assistant|>", "fields": { "tool_calls": { "open": "<actions>", "close": "</actions>", "content": "json", "transform_each": True, "transform": {"type": "function", "function": {"name": "{fn.name}", "arguments": "{fn.args}"}}, }, }, } model_out = '<actions>[{"fn": {"name": "a", "args": {"x": 1}}}, {"fn": {"name": "b", "args": {}}}]</actions>' self.assertEqual( parse_response(model_out, template_spec, prefix="")["tool_calls"], [ {"type": "function", "function": {"name": "a", "arguments": {"x": 1}}}, {"type": "function", "function": {"name": "b", "arguments": {}}}, ], ) def test_transform_dotted_path_errors(self): def spec_with(transform): return { "start_anchor": "<|assistant|>", "fields": {"x": {"open": "<x>", "close": "</x>", "content": "json", "transform": transform}}, } with self.assertRaisesRegex(ValueError, "missing key 'args'"): parse_response('<x>{"name": "n"}</x>', spec_with({"a": "{content.args}"}), prefix="") with self.assertRaisesRegex(ValueError, "cannot index into str"): parse_response('<x>{"name": "n"}</x>', spec_with({"a": "{content.name.x}"}), prefix="") with self.assertRaises(KeyError): parse_response("<x>{}</x>", spec_with({"a": "{missing}"}), prefix="") def test_transform_dotted_mixed_string_rejected(self): template_spec = { "start_anchor": "<|assistant|>", "fields": {"x": {"open": "<x>", "close": "</x>", "transform": {"v": "pre {content.args}"}}}, } with self.assertRaisesRegex(ValueError, "mixes"): parse_response("", template_spec, prefix="") def test_join_concatenates_repeated_matches(self): template_spec = { "defaults": {"role": "assistant"}, "start_anchor": "<|assistant|>", "fields": { "thinking": {"open": "<think>", "close": "</think>", "repeats": True, "join": " "}, "content": {"repeats": True, "join": " "}, }, } self.assertEqual( parse_response("<think>first</think>middle<think>second</think>done", template_spec, prefix=""), {"role": "assistant", "thinking": "first second", "content": "middle done"}, ) self.assertEqual( parse_response("<think>only</think>", template_spec, prefix=""), {"role": "assistant", "thinking": "only"}, ) def test_join_validation(self): no_repeats = {"start_anchor": "a", "fields": {"x": {"open": "<x>", "close": "</x>", "join": ""}}} with self.assertRaisesRegex(ValueError, "requires 'repeats'"): parse_response("", no_repeats, prefix="") bad_type = {"start_anchor": "a", "fields": {"x": {"open": "<x>", "close": "</x>", "repeats": True, "join": 7}}} with self.assertRaisesRegex(ValueError, "must be a string"): parse_response("", bad_type, prefix="") def test_join_requires_string_matches(self): template_spec = { "start_anchor": "a", "fields": {"x": {"open": "<x>", "close": "</x>", "repeats": True, "join": "", "content": "json"}}, } with self.assertRaisesRegex(ValueError, "parse to a string"): parse_response("<x>{}</x>", template_spec, prefix="") def test_optional_false_raises_when_missing(self): template_spec = { "defaults": {"role": "assistant"}, "start_anchor": "<|assistant|>", "fields": { "content": { "open": "<response>", "close": "</response>", "content": "text", "optional": False, }, }, } with self.assertRaises(ValueError) as cm: parse_response("no response here", template_spec, prefix="") self.assertIn("content", str(cm.exception)) def test_int_content_parser(self): template_spec = { "defaults": {"role": "assistant"}, "start_anchor": "<|assistant|>", "fields": { "count": { "open": "<n>", "close": "</n>", "content": "int", }, }, } self.assertEqual(parse_response("<n>42</n>", template_spec, prefix=""), {"role": "assistant", "count": 42}) def test_kv_lines_parser(self): template_spec = { "defaults": {"role": "assistant"}, "start_anchor": "<|assistant|>", "fields": { "metadata": { "open": "<meta>", "close": "</meta>", "content": "kv-lines", }, }, } self.assertEqual( parse_response("<meta>name: alice\nage: 30</meta>", template_spec, prefix=""), {"role": "assistant", "metadata": {"name": "alice", "age": "30"}}, ) def test_unknown_content_parser_rejected(self): bad_template = { "defaults": {"role": "assistant"}, "start_anchor": "<|assistant|>", "fields": {"x": {"open": "[", "close": "]", "content": "not-a-real-parser"}}, } with self.assertRaises(ValueError) as cm: parse_response("[hi]", bad_template) self.assertIn("unknown content parser", str(cm.exception).lower()) def test_unsupported_version_rejected(self): bad_template = { "version": 2, "defaults": {"role": "assistant"}, "start_anchor": "<|assistant|>", "fields": {"content": {"content": "text"}}, } with self.assertRaises(ValueError) as cm: parse_response("hello", bad_template) self.assertIn("version", str(cm.exception).lower()) def test_two_implicit_fields_rejected(self): bad_template = { "defaults": {"role": "assistant"}, "start_anchor": "<|assistant|>", "fields": { "a": {"content": "text"}, "b": {"content": "text"}, }, } with self.assertRaises(ValueError): parse_response("hello", bad_template) def test_transform_string_interpolation_rejected(self): bad_template = { "defaults": {"role": "assistant"}, "start_anchor": "<|assistant|>", "fields": { "tool": { "open": "<tool>", "close": "</tool>", "content": "text", "transform": {"label": "name: {content}"}, }, }, } with self.assertRaises(ValueError) as cm: parse_response("<tool>foo</tool>", bad_template) msg = str(cm.exception) self.assertIn("interpolation", msg) self.assertIn("{content}", msg) def test_named_groups_without_transform_rejected(self): bad_template = { "defaults": {"role": "assistant"}, "start_anchor": "<|assistant|>", "fields": { "tool": { "open_pattern": r"<tool name=(?P<name>\w+)>", "close": "</tool>", "content": "text", }, }, } with self.assertRaises(ValueError) as cm: parse_response("<tool name=foo>body</tool>", bad_template) msg = str(cm.exception) self.assertIn("transform", msg) self.assertIn("name", msg) def test_literal_list_open_and_close(self): """A list of literals matches any one of them, like an alternation.""" template_spec = { "defaults": {"role": "assistant"}, "start_anchor": "<|assistant|>", "fields": { "x": { "open": ["<a>", "<bb>"], "close": ["</a>", "</bb>"], "content": "text", }, }, } for opener, closer in (("<a>", "</a>"), ("<bb>", "</bb>"), ("<a>", "</bb>")): self.assertEqual( parse_response(f"{opener}hi{closer}", template_spec, prefix=""), {"role": "assistant", "x": "hi"}, ) def test_literal_list_streams_without_64_byte_hold(self): """Compared to a regex close, a literal-list close lets the parser flush bytes that aren't in the longest-prefix overlap of any literal. With `["<turn|>", "<|tool_response>", "<eos>"]` (longest = 16 chars), feeding 32 plain bytes should leave at most 15 unflushed.""" template_spec = { "defaults": {"role": "assistant"}, "start_anchor": "<|assistant|>", "fields": { "content": {"close": ["<turn|>", "<|tool_response>", "<eos>"], "content": "text"}, }, } parser = ResponseParser(template_spec, prefix="") plain = "x" * 32 flushed: list[str] = [] for ch in parser.feed(plain): if ch["type"] == "region_chunk": flushed.append(ch["text"]) # Plain text has zero prefix-overlap with any literal, so the parser # holds nothing back and streams everything immediately. self.assertEqual("".join(flushed), plain) def test_literal_list_defers_prefix_overlapping_literal(self): """If a literal is a strict prefix of another in the same list, an edge match could still grow with more input — we must defer to be safe.""" template_spec = { "defaults": {"role": "assistant"}, "start_anchor": "<|assistant|>", "fields": { "x": {"open": "<x>", "close": ["END", "ENDX"], "content": "text"}, }, } parser = ResponseParser(template_spec, prefix="") # "<x>hiEND" mid-stream: don't commit the close yet — "ENDX" might be coming. events = parser.feed("<x>hiEND") self.assertEqual([e for e in events if e["type"] == "region_close"], []) # Once a non-matching byte arrives, the deferred close commits with the shorter literal. events.extend(parser.feed(" more")) message, _ = parser.finalize() closes = [e for e in events if e["type"] == "region_close" and e["field"] == "x"] self.assertEqual(len(closes), 1) self.assertEqual(closes[0]["value"], "hi") self.assertEqual(message, {"role": "assistant", "x": "hi"}) def test_literal_list_rejects_empty_and_non_string(self): for bad_open in ([], [""], [1, 2], {"foo": "bar"}): spec = { "defaults": {"role": "assistant"}, "start_anchor": "<|assistant|>", "fields": {"x": {"open": bad_open, "close": "</x>", "content": "text"}}, } with self.assertRaises(ValueError): parse_response("<x>hi</x>", spec, prefix="") def test_field_without_close_runs_to_end_of_stream(self): """A field with no `close`/`close_pattern` stays open until end-of-stream, capturing everything after its open.""" spec = { "defaults": {"role": "assistant"}, "start_anchor": "<|assistant|>", "fields": {"content": {"open": "<resp>", "content": "text"}}, } self.assertEqual( parse_response("<resp>hello world", spec, prefix=""), {"role": "assistant", "content": "hello world"}, ) # Fixtures shared by the streaming tests: one representative input per template, # re-used for both the correctness invariant (any chunking → same dict) and the # event-shape tests (do we emit the right events in the right order?). _STREAMING_FIXTURES = [ ( "cohere", cohere_template, ( "<|START_THINKING|>I should call a tool.<|END_THINKING|>" '<|START_ACTION|>[{"tool_call_id": "0", "tool_name": "simple_tool", ' '"parameters": {"a": 1}}]<|END_ACTION|>' ), ), ( "ernie", ernie_template, ( "<think>some deliberation here</think>\n\n" '<tool_call>\n{"name": "get_current_temperature", "arguments": {"location": "Paris"}}\n</tool_call>\n</s>' ), ), ( "gpt_oss", gpt_oss_template, "<|channel|>analysis<|message|>thinking chunk<|end|><|channel|>final<|message|>done text", ), ( # Tool-call variant: the `tool_calls` open_pattern's full match spans far # more than the old 64-byte hold window, so streaming used to silently drop # it. Kept as a streaming fixture so every chunking step re-checks it. "gpt_oss_tool", gpt_oss_template, ( "<|channel|>analysis<|message|>Let me check.<|end|>" "<|start|>assistant<|channel|>commentary to=functions.get_current_weather " '<|constrain|>json<|message|>{"location": "San Francisco, CA"}<|call|>' ), ), ( "smollm", smollm_template, '<think>thinking</think>\n<tool_call>{"name": "fn", "arguments": {"x": 1}}</tool_call>', ), ( "qwen3", qwen3_template, ( "<think>short thought</think>\n" "<tool_call>\n<function=get_weather>\n" "<parameter=city>\nParis\n</parameter>\n" "</function>\n</tool_call>" ), ), ( "gemma4", gemma4_template, '<|channel>thought\nhi<channel|><|tool_call>call:foo{a:1,b:<|"|>bar<|"|>}<tool_call|>', ), ( # Exercises `join` fields (two thinking blocks) and dotted transform paths. "inkling", inkling_template, ( "<|content_thinking|>Consider the weather.<|end_message|>" "<|message_model|><|content_thinking|> Tokyo, probably.<|end_message|>" "<|message_model|><|content_text|>Checking now.<|end_message|>" "<|message_model|>get_weather<|content_invoke_tool_json|>" '{"name":"get_weather","args":{"city":"Tokyo"}}<|end_message|>' "<|content_model_end_sampling|>" ), ), ] def _chunk_fixed(text: str, step: int): for i in range(0, len(text), step): yield text[i : i + step] def _chunk_random(text: str, rng: random.Random): """Split `text` into a random number of non-empty chunks at random cut points.""" if len(text) <= 1: yield text return num_cuts = rng.randint(0, len(text) - 1) cuts = sorted(rng.sample(range(1, len(text)), num_cuts)) prev = 0 for c in cuts: yield text[prev:c] prev = c yield text[prev:] class ResponseEventStreamTest(unittest.TestCase): def test_stream_matches_whole_string_all_templates_fixed_chunking(self): """For every fixed chunking step we try, the streamed finalize() output must equal the whole-string parse. Regression coverage for specific edge-case byte boundaries (1-byte chunks hit every prefix).""" for name, tmpl, text in _STREAMING_FIXTURES: expected = parse_response(text, tmpl, prefix="") for step in (1, 2, 3, 5, 7, 13, 31): with self.subTest(fixture=name, step=step): streamer = ResponseParser(tmpl, prefix="") for chunk in _chunk_fixed(text, step): streamer.feed(chunk) message, _ = streamer.finalize() self.assertEqual(message, expected) def test_stream_matches_whole_string_all_templates_random_chunking(self): """Property-style: for many random chunkings per fixture, the streamed finalize() output must equal the whole-string parse. Seeded so failures reproduce.""" rng = random.Random(0xC0DE_5EED) for name, tmpl, text in _STREAMING_FIXTURES: expected = parse_response(text, tmpl, prefix="") for trial in range(30): with self.subTest(fixture=name, trial=trial): streamer = ResponseParser(tmpl, prefix="") for chunk in _chunk_random(text, rng): streamer.feed(chunk) message, _ = streamer.finalize() self.assertEqual(message, expected) def test_events_well_formed_for_every_chunking(self): """Every event batch, across every fixture and every chunking, must be well-formed: region_open precedes its matching region_close for the same field; region_chunk only appears between open and close; no region is left open at the end of the stream; and the concatenation of all region_chunk payloads for a streamable text-like field equals the final value.""" rng = random.Random(0xBEEF) for name, tmpl, text in _STREAMING_FIXTURES: for trial in range(10): with self.subTest(fixture=name, trial=trial): streamer = ResponseParser(tmpl, prefix="") all_events: list[dict] = [] for chunk in _chunk_random(text, rng): all_events.extend(streamer.feed(chunk)) _, final_events = streamer.finalize() all_events.extend(final_events) self._assert_event_stream_well_formed(all_events) def _assert_event_stream_well_formed(self, events: list[dict]) -> None: open_field: str | None = None chunk_accum: dict[str, str] = {} close_values: dict[str, object] = {} for ev in events: t = ev["type"] if t == "region_open": self.assertIsNone(open_field, f"nested region_open without close: {ev}") open_field = ev["field"] chunk_accum.setdefault(open_field, "") elif t == "region_chunk": self.assertEqual(open_field, ev["field"], f"chunk outside its region: {ev}") # Every chunk carries a boolean `dirty` flag. self.assertIsInstance(ev["dirty"], bool, f"missing/non-bool dirty: {ev}") chunk_accum[open_field] += ev["text"] elif t == "region_close": self.assertEqual(open_field, ev["field"], f"close for non-open region: {ev}") close_values[open_field] = ev["value"] open_field = None else: self.fail(f"unexpected event type: {ev!r}") self.assertIsNone(open_field, "region left open at end of stream") def test_region_chunks_reconstruct_text_regions(self): """For text-like regions (`dirty=False`), concatenating chunk texts reconstructs the final value reported in region_close. Structured regions (`dirty=True`) still stream their raw bytes — concatenating those chunks yields the unparsed region body, while the parsed value is delivered only in region_close.""" # Single representative case with a long text region and a JSON region. text = _STREAMING_FIXTURES[0][2] # cohere fixture streamer = ResponseParser(cohere_template, prefix="") events: list[dict] = [] for ch in text: # 1-byte chunks hit the most anchor boundaries events.extend(streamer.feed(ch)) _, final_events = streamer.finalize() events.extend(final_events) # Reconstruct per-field. per_field_chunks: dict[str, list[str]] = {} per_field_dirty: dict[str, set[bool]] = {} per_field_close_value: dict[str, object] = {} for ev in events: if ev["type"] == "region_chunk": per_field_chunks.setdefault(ev["field"], []).append(ev["text"]) per_field_dirty.setdefault(ev["field"], set()).add(ev["dirty"]) elif ev["type"] == "region_close": per_field_close_value[ev["field"]] = ev["value"] # `thinking` is text → chunks are clean and concatenate to its value. self.assertIn("thinking", per_field_chunks) self.assertEqual(per_field_dirty["thinking"], {False}) self.assertEqual("".join(per_field_chunks["thinking"]), "I should call a tool.") self.assertEqual(per_field_close_value["thinking"], "I should call a tool.") # `tool_calls` is json → dirty chunks stream the raw body, parsed value on close. self.assertIn("tool_calls", per_field_chunks) self.assertEqual(per_field_dirty["tool_calls"], {True}) self.assertEqual( "".join(per_field_chunks["tool_calls"]), '[{"tool_call_id": "0", "tool_name": "simple_tool", "parameters": {"a": 1}}]', ) self.assertIn("tool_calls", per_field_close_value) def test_dirty_flag_marks_structured_regions(self): """A template with one text field and one structured field per parser family: text/int/float/bool stream chunks with `dirty=False`, while json/xml-inline/kv-lines stream chunks with `dirty=True`, and those dirty chunks concatenate to the raw region body before parsing.""" spec = { "defaults": {"role": "assistant"}, "start_anchor": "<|assistant|>", "fields": { "thinking": {"open": "<t>", "close": "</t>", "content": "text"}, "score": {"open": "<n>", "close": "</n>", "content": "int"}, "json_call": {"open": "<j>", "close": "</j>", "content": "json"}, "xml_call": { "open": "<x>", "close": "</x>", "content": "xml-inline", "content_args": {"tag_pattern": r"<(?P<key>\w+)=(?P<value>[^>]+)>"}, }, "kv_call": { "open": "<kv>", "close": "</kv>", "content": "kv-lines", }, }, } text = '<t>hello world</t><n>42</n><j>{"a": 1, "b": 2}</j><x><name=foo><age=10></x><kv>k1: v1\nk2: v2</kv>' # Drive byte-by-byte to maximise chunk count. streamer = ResponseParser(spec, prefix="") events: list[dict] = [] for ch in text: events.extend(streamer.feed(ch)) _, final_events = streamer.finalize() events.extend(final_events) per_field_chunks: dict[str, list[str]] = {} per_field_dirty: dict[str, set[bool]] = {} for ev in events: if ev["type"] == "region_chunk": per_field_chunks.setdefault(ev["field"], []).append(ev["text"]) per_field_dirty.setdefault(ev["field"], set()).add(ev["dirty"]) # Clean (streamable) regions. for field in ("thinking", "score"): self.assertEqual(per_field_dirty[field], {False}, f"{field} should be clean") # Dirty (structured) regions. for field in ("json_call", "xml_call", "kv_call"): self.assertEqual(per_field_dirty[field], {True}, f"{field} should be dirty") # Dirty chunks reconstruct the raw region body (un-parsed). Clean # chunks reconstruct the verbatim body too — stripping happens at close. self.assertEqual("".join(per_field_chunks["thinking"]), "hello world") self.assertEqual("".join(per_field_chunks["score"]), "42") self.assertEqual("".join(per_field_chunks["json_call"]), '{"a": 1, "b": 2}') self.assertEqual("".join(per_field_chunks["xml_call"]), "<name=foo><age=10>") self.assertEqual("".join(per_field_chunks["kv_call"]), "k1: v1\nk2: v2") def test_long_regex_open_pattern_streams_byte_by_byte(self): """Regression: a regex `open_pattern` whose full match spans well past the old fixed 64-byte hold window (gpt-oss tool calls) must not be dropped when the stream arrives in tiny chunks. Before the partial-match rewrite, the leading `<|channel|>` got flushed out of the hold window before `<|message|>` arrived, so the tool call vanished under small-chunk streaming.""" text = ( "<|channel|>analysis<|message|>Let me check.<|end|>" "<|start|>assistant<|channel|>commentary to=functions.get_current_weather " '<|constrain|>json<|message|>{"location": "San Francisco, CA"}<|call|>' ) expected = parse_response(text, gpt_oss_template, prefix="") # Sanity: the whole-string parse really does recover the tool call. self.assertEqual(len(expected["tool_calls"]), 1) self.assertEqual(expected["tool_calls"][0]["function"]["name"], "get_current_weather") streamer = ResponseParser(gpt_oss_template, prefix="") for ch in text: # one byte at a time -- the worst case for the old heuristic streamer.feed(ch) message, _ = streamer.finalize() self.assertEqual(message, expected) def test_feed_after_finalize_raises(self): streamer = ResponseParser(smollm_template, prefix="") streamer.feed("<think>x</think>") streamer.finalize() with self.assertRaises(RuntimeError): streamer.feed("more") with self.assertRaises(RuntimeError): streamer.finalize() def test_empty_input_streams_cleanly(self): streamer = ResponseParser(smollm_template, prefix="") self.assertEqual(streamer.feed(""), []) result, final_events = streamer.finalize() # Only the default fields should remain; nothing else is required. self.assertEqual(result, {"role": "assistant"}) self.assertEqual(final_events, []) class PrefixAndTruncationTest(unittest.TestCase): def test_prefix_lands_inside_explicit_region(self): """A Qwen-style template emits `<|im_start|>assistant\\n<think>\\n` as the assistant prefix. The model continues from inside the thinking block.""" prompt = ( "<|im_start|>system\nYou are helpful<|im_end|>\n" "<|im_start|>user\nHi<|im_end|>\n" "<|im_start|>assistant\n<think>\n" ) generated = "Let me think...</think>" stream = ResponseParser(qwen3_template, prefix=prompt) # The region_open for `thinking` surfaces via initial_events; the # caller replays it before feeding model output. self.assertEqual( [(e["type"], e["field"]) for e in stream.initial_events], [("region_open", "thinking"), ("region_chunk", "thinking")], ) events = stream.feed(generated) result, _ = stream.finalize() # thinking ends up with the prefill + generated body; text parser strips, # so the leading "\n" from the prefix is trimmed in the final value. self.assertEqual(result, {"role": "assistant", "thinking": "Let me think..."}) # The feed stream only sees the rest of the body and the close; the # prefill already emitted region_open. self.assertEqual([e["type"] for e in events], ["region_chunk", "region_close"]) self.assertEqual(events[1]["field"], "thinking") def test_prefix_truncated_to_last_anchor(self): """Multiple `<|im_start|>assistant\\n` anchors in the prefix (multi-turn conversation): only the slice after the LAST anchor matters.""" prompt = ( "<|im_start|>system\nA<|im_end|>\n" "<|im_start|>user\nB<|im_end|>\n" "<|im_start|>assistant\nEarlier reply<|im_end|>\n" "<|im_start|>user\nFollowup<|im_end|>\n" "<|im_start|>assistant\n<think>\n" ) stream = ResponseParser(qwen3_template, prefix=prompt) # We landed inside `thinking` (from the LAST assistant turn's `<think>\n`), # not in some earlier-turn artifact. opens = [e for e in stream.initial_events if e["type"] == "region_open"] self.assertEqual([e["field"] for e in opens], ["thinking"]) # No earlier-turn content leaked into output. stream.feed("done</think>") stream.finalize() self.assertEqual(stream._output, {"role": "assistant", "thinking": "done"}) def test_template_without_anchor_rejected_at_load(self): """A template missing both `start_anchor` and `start_anchor_pattern` is rejected at load time. Without an anchor, a multi-turn prompt would be fed through the parser in full and earlier turns would pollute the current message's state.""" anchorless = {k: v for k, v in qwen3_template.items() if k != "start_anchor"} with self.assertRaises(ValueError) as cm: ResponseParser(anchorless) msg = str(cm.exception) self.assertIn("start_anchor", msg) def test_prefix_anchor_not_found_falls_back(self): """Spec has start_anchor but the prefix doesn't contain it: parser falls back to processing the entire prefix (with a logged warning).""" prompt = "<think>\n" # no <|im_start|>assistant\n stream = ResponseParser(qwen3_template, prefix=prompt) opens = [e for e in stream.initial_events if e["type"] == "region_open"] self.assertEqual([e["field"] for e in opens], ["thinking"]) stream.feed("hi</think>") stream.finalize() self.assertEqual(stream._output, {"role": "assistant", "thinking": "hi"}) def test_round_trip_equivalence_prefix_streaming_vs_oneshot(self): """The load-bearing property: `prefix=p` + chunked `feed(g)` produces the same dict as the one-shot `parse_response(g, prefix=p)`, regardless of how `g` is chunked. (Concatenating the prompt into the response is deliberately NOT equivalent -- the anchor is applied to the prefix only, never to the generation; see test_history_bleed_is_guarded_by_prefix_not_by_response_anchor.)""" prompt = "<|im_start|>system\nA<|im_end|>\n<|im_start|>user\nB<|im_end|>\n<|im_start|>assistant\n<think>\n" for name, tmpl_dict, gen_text in _STREAMING_FIXTURES: if "thinking" not in gen_text and "<think>" not in gen_text: continue # Only fixtures whose generation text is compatible with starting # inside `thinking`. Restrict to qwen3 / smollm shape for clarity. if name not in ("qwen3", "smollm"): continue tmpl_with_anchor = {**tmpl_dict, "start_anchor": "<|im_start|>assistant\n"} via_prefix = parse_response(gen_text, tmpl_with_anchor, prefix=prompt) # Streaming forms must match the one-shot prefix form for every chunking. for step in (1, 3, 7, 31): with self.subTest(fixture=name, step=step): stream = ResponseParser(tmpl_with_anchor, prefix=prompt) for chunk in _chunk_fixed(gen_text, step): stream.feed(chunk) message, _ = stream.finalize() self.assertEqual(message, via_prefix) def test_history_bleed_is_guarded_by_prefix_not_by_response_anchor(self): """The `start_anchor` guards against history bleed only via `prefix=`; it is NOT applied to the response. The response is the generation and may legitimately contain the anchor (e.g. gpt-oss harmony opens every channel with `<|start|>assistant`), so truncating it would drop real content. Passing the prompt as `prefix=` is the way to keep earlier turns out of the parse.""" spec = { "defaults": {"role": "assistant"}, "start_anchor": "<|im_start|>assistant\n", "fields": {"content": {"close_pattern": r"\Z", "content": "text"}}, } prompt = "<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n" gen = "Hello there!" clean = {"role": "assistant", "content": "Hello there!"} # The supported guard: pass the prompt as prefix= so history is truncated off the prefix. self.assertEqual(parse_response(gen, spec, prefix=prompt), clean) # Pure generation parses cleanly with the explicit no-prefix opt-out (prefix=""). self.assertEqual(parse_response(gen, spec, prefix=""), clean) # An anchor inside the response is treated as content, never as a history boundary: # gpt-oss re-emits `<|start|>assistant` between channels, and that content survives. gpt_oss_gen = ( "<|channel|>analysis<|message|>thinking<|end|><|start|>assistant<|channel|>final<|message|>answer" ) self.assertEqual( parse_response(gpt_oss_gen, gpt_oss_template, prefix=""), {"role": "assistant", "thinking": "thinking", "content": "answer"}, ) def test_prefix_with_open_close_inside_truncated_region(self): """Prefix opens AND closes a region. The full open/chunk/close event sequence is surfaced via initial_events, and the closed region lands in the output dict — so renderers can show prefill content.""" spec = { "defaults": {"role": "assistant"}, "start_anchor": "[BEGIN]", "fields": { "tag": {"open": "<tag>", "close": "</tag>", "content": "text"}, "body": {"close_pattern": r"$", "content": "text"}, # implicit }, } prefix = "noise[BEGIN]<tag>silently consumed</tag>" stream = ResponseParser(spec, prefix=prefix) types = [e["type"] for e in stream.initial_events] self.assertEqual(types, ["region_open", "region_chunk", "region_close"]) self.assertTrue(all(e["field"] == "tag" for e in stream.initial_events)) self.assertEqual(stream.initial_events[-1]["value"], "silently consumed") stream.feed("real generated body") result, _ = stream.finalize() self.assertEqual(result["tag"], "silently consumed") self.assertEqual(result["body"], "real generated body") def test_prefix_lands_inside_implicit_region(self): """Prefix wrote plaintext into the implicit region (e.g. assistant prefill before the model continues). The region_open for the implicit region must surface via initial_events so consumers don't miss it — `_opened` will already be True by the time feed runs.""" prompt = "<|im_start|>assistant\nSure, here is " stream = ResponseParser(smollm_template, prefix=prompt) opens = [e for e in stream.initial_events if e["type"] == "region_open"] self.assertEqual([e["field"] for e in opens], ["content"]) events = stream.feed("the answer<|im_end|>") # No second region_open from feed — the implicit region was already # opened during prefill and surfaced via initial_events. self.assertNotIn("region_open", [e["type"] for e in events]) def test_prefix_partial_pattern_at_boundary(self): """Post-truncation prefix ends mid-delimiter. The first feed completes the match; initial_events is empty (no region opened within the prefix yet) and the open fires from `feed()`.""" prefix = "<|im_start|>assistant\n<thi" # incomplete `<think>` stream = ResponseParser(qwen3_template, prefix=prefix) self.assertEqual(stream.initial_events, []) events = stream.feed("nk>real body</think>") types = [e["type"] for e in events] self.assertIn("region_open", types) # think opens during feed, not prefill stream.finalize() self.assertEqual(stream._output, {"role": "assistant", "thinking": "real body"}) def test_prefix_token_ids_decoded(self): """The tokenizer-level helper accepts token IDs as prefix (decoded internally), matching how `response` is handled.""" tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2") tokenizer.response_template = qwen3_template prefix_text = "<|im_start|>assistant\n<think>\n" prefix_ids = tokenizer(prefix_text).input_ids from_str = tokenizer.parse_response("hi</think>", prefix=prefix_text) from_ids = tokenizer.parse_response("hi</think>", prefix=prefix_ids) self.assertEqual(from_str, from_ids) def test_batched_parse_response_with_prefix(self): """Batched `parse_response`: a single prefix is broadcast to every item, a per-item prefix list is matched positionally, and a prefix count that doesn't match the batch raises.""" tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2") tokenizer.response_template = qwen3_template prefix = "<|im_start|>assistant\n<think>\n" gen_a, gen_b = "thinking A</think>answer A", "thinking B</think>answer B" single_a = tokenizer.parse_response(gen_a, prefix=prefix) single_b = tokenizer.parse_response(gen_b, prefix=prefix) self.assertNotEqual(single_a, single_b) # A single prefix is broadcast across the whole batch. self.assertEqual(tokenizer.parse_response([gen_a, gen_b], prefix=prefix), [single_a, single_b]) # One prefix per item (here identical) is matched up positionally. self.assertEqual(tokenizer.parse_response([gen_a, gen_b], prefix=[prefix, prefix]), [single_a, single_b]) # A prefix batch whose length doesn't match the responses is an error. with self.assertRaises(ValueError): tokenizer.parse_response([gen_a, gen_b], prefix=[prefix]) def test_tokenizer_get_response_parser_with_prefix(self): """`tokenizer.get_response_parser(prefix=...)` returns a stream that is already in the right initial state.""" tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2") tokenizer.response_template = qwen3_template prefix = "<|im_start|>assistant\n<think>\n" stream = tokenizer.get_response_parser(prefix=prefix) opens = [e for e in stream.initial_events if e["type"] == "region_open"] self.assertEqual([e["field"] for e in opens], ["thinking"]) stream.feed("body</think>") result, _ = stream.finalize() self.assertEqual(result, {"role": "assistant", "thinking": "body"}) # xml-inline without a value_parser: parameter bodies stay raw strings until tools= coerces them. _XML_STRING_ARGS_TEMPLATE = { "defaults": {"role": "assistant"}, "start_anchor": "<|im_start|>assistant\n", "fields": { "tool_calls": { "open_pattern": r"<tool_call>\s*<function=(?P<name>\w+)>", "close": "</tool_call>", "repeats": True, "content": "xml-inline", "content_args": { "tag_pattern": r"<parameter=(?P<key>\w+)>\s*(?P<value>.*?)\s*</parameter>", }, "transform": {"type": "function", "function": {"name": "{name}", "arguments": "{content}"}}, }, }, } # kv-lines without a value_parser: values likewise stay raw strings for tools= to cast. _KV_LINES_TOOLS_TEMPLATE = { "defaults": {"role": "assistant"}, "start_anchor": "<|im_start|>assistant\n", "fields": { "tool_calls": { "open_pattern": r"<tool_call>\s*<function=(?P<name>\w+)>\n", "close": "</tool_call>", "repeats": True, "content": "kv-lines", "transform": {"type": "function", "function": {"name": "{name}", "arguments": "{content}"}}, }, }, } _SET_ALARM_CALL = ( "<tool_call>\n<function=set_alarm>\n" "<parameter=hour>\n7\n</parameter>\n" "<parameter=enabled>\ntrue\n</parameter>\n" "<parameter=label>\nwake up\n</parameter>\n" "</function>\n</tool_call>" ) def _set_alarm_tools(**properties): return [ { "type": "function", "function": { "name": "set_alarm", "parameters": {"type": "object", "properties": properties}, }, } ] _SET_ALARM_TOOLS = _set_alarm_tools( hour={"type": "integer"}, enabled={"type": "boolean"}, label={"type": "string"}, ) def _first_tool_args(message): return message["tool_calls"][0]["function"]["arguments"] def _parser_with_tools(tools): return ResponseParser(_XML_STRING_ARGS_TEMPLATE, prefix="", tools=tools) class ToolArgCoercionTest(unittest.TestCase): def test_coerce_tool_calls_casts_declared_types(self): tools = _set_alarm_tools( count={"type": "integer"}, ratio={"type": "number"}, enabled={"type": "boolean"}, tags={"type": "array"}, note={"type": "string"}, ) arguments = { "count": "3", "ratio": "1.5", "enabled": "true", "tags": '["a", "b"]', "note": "hello", "already_typed": 7, "extra": "unscheduled", } call = {"type": "function", "function": {"name": "set_alarm", "arguments": arguments}} self.assertIs(_parser_with_tools(tools)._coerce_tool_calls(call), call) self.assertEqual( call["function"]["arguments"], { "count": 3, "ratio": 1.5, "enabled": True, "tags": ["a", "b"], "note": "hello", "already_typed": 7, "extra": "unscheduled", }, ) def test_coerce_falls_back_to_raw_on_failure(self): self.assertEqual(_coerce("not-a-number", ("integer",)), "not-a-number") def test_coerce_handles_any_of_and_null(self): self.assertEqual(_coerce("5", ("integer", "null")), 5) self.assertIsNone(_coerce("null", ("integer", "null"))) def test_coerce_booleans_match_bool_parser(self): # Accept the same literals as the `bool` content parser, case-insensitively. for raw, expected in [("true", True), ("True", True), ("1", True), ("false", False), ("0", False)]: self.assertEqual(_coerce(raw, ("boolean",)), expected) # Non-boolean text stays a string rather than silently becoming False. self.assertEqual(_coerce("maybe", ("boolean",)), "maybe") def test_coerce_object_array_require_matching_json(self): # A JSON object body is only accepted for an `object` param, a JSON array only for `array`. self.assertEqual(_coerce("[1, 2]", ("object",)), "[1, 2]") self.assertEqual(_coerce('{"a": 1}', ("array",)), '{"a": 1}') self.assertEqual(_coerce('{"a": 1}', ("object",)), {"a": 1}) # NaN / inf are not valid JSON numbers, so a `number` param keeps the raw text. self.assertEqual(_coerce("NaN", ("number",)), "NaN") def test_coerce_tool_calls_handles_single_and_list(self): parser = _parser_with_tools(_SET_ALARM_TOOLS) call = {"type": "function", "function": {"name": "set_alarm", "arguments": {"hour": "7"}}} self.assertIs(parser._coerce_tool_calls(call), call) self.assertEqual(call["function"]["arguments"], {"hour": 7}) # A list of calls (as produced by `transform_each`) is coerced element-wise. calls = [{"type": "function", "function": {"name": "set_alarm", "arguments": {"hour": "9"}}}] self.assertEqual(parser._coerce_tool_calls(calls)[0]["function"]["arguments"], {"hour": 9}) # Non-tool-call values pass through untouched. self.assertEqual(parser._coerce_tool_calls("hello"), "hello") def test_schema_types_handles_get_json_schema_dialect(self): self.assertEqual(_schema_types({"type": "integer"}), ("integer",)) self.assertEqual(_schema_types({"type": ["integer", "string"]}), ("integer", "string")) self.assertEqual(_schema_types({"anyOf": [{"type": "boolean"}, {"type": "string"}]}), ("boolean", "string")) self.assertEqual(_schema_types({"type": "integer", "nullable": True}), ("integer", "null")) # Undescribed parameters resolve to no candidate types, making coercion a no-op. self.assertEqual(_schema_types({"description": "no type"}), ()) def test_parse_response_tools_coerces_xml_inline_string_args(self): # Without a value_parser, xml-inline argument bodies stay strings; tools= casts them. without = parse_response(_SET_ALARM_CALL, _XML_STRING_ARGS_TEMPLATE, prefix="") with_tools = parse_response(_SET_ALARM_CALL, _XML_STRING_ARGS_TEMPLATE, prefix="", tools=_SET_ALARM_TOOLS) self.assertEqual(_first_tool_args(without), {"hour": "7", "enabled": "true", "label": "wake up"}) self.assertEqual(_first_tool_args(with_tools), {"hour": 7, "enabled": True, "label": "wake up"}) def test_streaming_tools_coerces_on_region_close(self): # Coercion must land on the region_close event during feed(), not only after finalize(). stream = ResponseParser(_XML_STRING_ARGS_TEMPLATE, prefix="", tools=_SET_ALARM_TOOLS) closes = [ event["value"] for chunk in _chunk_fixed(_SET_ALARM_CALL, 8) for event in stream.feed(chunk) if event["type"] == "region_close" and event["field"] == "tool_calls" ] self.assertEqual(len(closes), 1) self.assertEqual(closes[0]["function"]["arguments"], {"hour": 7, "enabled": True, "label": "wake up"}) def test_qwen3_tools_coerces_strings_left_by_value_parser(self): # qwen3's json+allow_non_json value_parser types what it can (`true`) and leaves # invalid JSON (`007`) as a string; tools= then casts those leftover strings. model_out = ( "<tool_call>\n<function=set_alarm>\n" "<parameter=hour>\n007\n</parameter>\n" "<parameter=enabled>\ntrue\n</parameter>\n" "<parameter=label>\nwake up\n</parameter>\n" "</function>\n</tool_call>" ) without = parse_response(model_out, qwen3_template, prefix="") with_tools = parse_response(model_out, qwen3_template, prefix="", tools=_SET_ALARM_TOOLS) self.assertEqual(_first_tool_args(without), {"hour": "007", "enabled": True, "label": "wake up"}) self.assertEqual(_first_tool_args(with_tools), {"hour": 7, "enabled": True, "label": "wake up"}) def test_coercion_never_reworks_already_typed_values(self): # Coercion only casts strings: values the value_parser already typed are final model_out = "<tool_call>\n<function=set_alarm>\n<parameter=label>\n1.50\n</parameter>\n</tool_call>" # qwen3's lax value_parser has already read 1.50 as the float 1.5, so the string-typed label stays a float ... self.assertEqual( _first_tool_args(parse_response(model_out, qwen3_template, prefix="", tools=_SET_ALARM_TOOLS)), {"label": 1.5}, ) # ... while without a value_parser the raw text reaches the schema cast intact self.assertEqual( _first_tool_args(parse_response(model_out, _XML_STRING_ARGS_TEMPLATE, prefix="", tools=_SET_ALARM_TOOLS)), {"label": "1.50"}, ) def test_kv_lines_string_args_are_coerced(self): model_out = "<tool_call>\n<function=set_alarm>\nhour: 7\nenabled: true\n</tool_call>" without = parse_response(model_out, _KV_LINES_TOOLS_TEMPLATE, prefix="") with_tools = parse_response(model_out, _KV_LINES_TOOLS_TEMPLATE, prefix="", tools=_SET_ALARM_TOOLS) self.assertEqual(_first_tool_args(without), {"hour": "7", "enabled": "true"}) self.assertEqual(_first_tool_args(with_tools), {"hour": 7, "enabled": True}) def test_non_tool_call_regions_are_untouched(self): # A field that captures a `name` but does not parse into a tool call must be left # alone, even when the capture happens to match a tool: its keys are not arguments. template = { "defaults": {"role": "assistant"}, "start_anchor": "<|im_start|>assistant\n", "fields": { "citation": { "open_pattern": r"<cite source=(?P<name>\w+)>", "close": "</cite>", "content": "xml-inline", "content_args": { "tag_pattern": r"<(?P<key>\w+)>\s*(?P<value>.*?)\s*</\1>", "value_parser": {"name": "json", "args": {"allow_non_json": True}}, }, "transform": {"source": "{name}", "fields": "{content}"}, }, }, } model_out = "<cite source=set_alarm><label>1.50</label><hour>7</hour></cite>" expected = {"source": "set_alarm", "fields": {"label": 1.5, "hour": 7}} self.assertEqual(parse_response(model_out, template, prefix="")["citation"], expected) self.assertEqual( parse_response(model_out, template, prefix="", tools=_SET_ALARM_TOOLS)["citation"], expected, ) def test_merge_duplicates_arguments_are_cast_element_wise(self): # merge_duplicates collects repeated tags into a list, which is cast element-wise template = copy.deepcopy(_XML_STRING_ARGS_TEMPLATE) template["fields"]["tool_calls"]["content_args"]["merge_duplicates"] = True model_out = ( "<tool_call>\n<function=set_alarm>\n" "<parameter=hour>\n7\n</parameter>\n" "<parameter=hour>\n9\n</parameter>\n" "</function>\n</tool_call>" ) self.assertEqual(_first_tool_args(parse_response(model_out, template, prefix="")), {"hour": ["7", "9"]}) self.assertEqual( _first_tool_args(parse_response(model_out, template, prefix="", tools=_SET_ALARM_TOOLS)), {"hour": [7, 9]}, ) # Elements that don't cast, and non-string elements, are left as they are. parser = _parser_with_tools(_SET_ALARM_TOOLS) call = {"type": "function", "function": {"name": "set_alarm", "arguments": {"hour": ["7", "x", 9]}}} parser._coerce_tool_calls(call) self.assertEqual(call["function"]["arguments"], {"hour": [7, "x", 9]}) def test_coerce_tool_calls_ignores_unusable_function_name(self): # A transform can hand us a name parsed from model output, so a non-string name # must be ignored rather than raising on the schema lookup. parser = _parser_with_tools(_SET_ALARM_TOOLS) call = {"type": "function", "function": {"name": ["set_alarm"], "arguments": {"hour": "7"}}} self.assertEqual(parser._coerce_tool_calls(call), call) self.assertEqual(call["function"]["arguments"], {"hour": "7"}) def test_union_with_container_still_keeps_scalar_text(self): model_out = "<tool_call>\n<function=set_alarm>\n<parameter=label>\n1.50\n</parameter>\n</tool_call>" union = _set_alarm_tools(label={"anyOf": [{"type": "string"}, {"type": "object"}]}) # `string` never casts, so the union keeps scalar-looking text as text ... self.assertEqual( _first_tool_args(parse_response(model_out, _XML_STRING_ARGS_TEMPLATE, prefix="", tools=union)), {"label": "1.50"}, ) # ... while the `object` branch still decodes a body that really is an object. object_body = '<tool_call>\n<function=set_alarm>\n<parameter=label>\n{"a": 1}\n</parameter>\n</tool_call>' self.assertEqual( _first_tool_args(parse_response(object_body, _XML_STRING_ARGS_TEMPLATE, prefix="", tools=union)), {"label": {"a": 1}}, ) def test_callable_tools_are_converted_to_schemas(self): # Functions are converted with `get_json_schema`; `label` is Optional, so "null" casts to None def set_alarm(hour: int, enabled: bool, label: str | None = None): """ Set an alarm. Args: hour: The hour the alarm should ring at. enabled: Whether the alarm starts out enabled. label: An optional label for the alarm. """ model_out = ( "<tool_call>\n<function=set_alarm>\n" "<parameter=hour>\n7\n</parameter>\n" "<parameter=enabled>\ntrue\n</parameter>\n" "<parameter=label>\nnull\n</parameter>\n" "</function>\n</tool_call>" ) parsed = parse_response(model_out, _XML_STRING_ARGS_TEMPLATE, prefix="", tools=[set_alarm]) self.assertEqual(_first_tool_args(parsed), {"hour": 7, "enabled": True, "label": None}) def test_get_response_parser_forwards_tools(self): tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2") tokenizer.response_template = _XML_STRING_ARGS_TEMPLATE stream = tokenizer.get_response_parser(prefix="", tools=_SET_ALARM_TOOLS) stream.feed(_SET_ALARM_CALL) message, _ = stream.finalize() self.assertEqual(_first_tool_args(message), {"hour": 7, "enabled": True, "label": "wake up"}) def test_tokenizer_parse_response_forwards_tools(self): # The public tokenizer entry point threads tools= through to coercion. tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2") tokenizer.response_template = _XML_STRING_ARGS_TEMPLATE parsed = tokenizer.parse_response(_SET_ALARM_CALL, prefix="", tools=_SET_ALARM_TOOLS) self.assertEqual(_first_tool_args(parsed), {"hour": 7, "enabled": True, "label": "wake up"}) def test_tools_is_a_noop_for_typed_json_tool_calls(self): # JSON tool-call bodies already carry their types, so passing tools= must not change them. model_out = ( "<|START_THINKING|>x<|END_THINKING|>" '<|START_ACTION|>[{"tool_call_id": "0", "tool_name": "set_alarm", ' '"parameters": {"hour": 7, "enabled": true}}]<|END_ACTION|><|END_OF_TURN_TOKEN|>' ) without = parse_response(model_out, cohere_template, prefix="") with_tools = parse_response(model_out, cohere_template, prefix="", tools=_SET_ALARM_TOOLS) self.assertEqual(without, with_tools) self.assertEqual(_first_tool_args(with_tools), {"hour": 7, "enabled": True}) if __name__ == "__main__": unittest.main()