| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # Copyright 2026 The HuggingFace Team. All rights reserved. |
| 5 | # ruff: noqa |
| 6 | # fmt: off |
| 7 | # |
| 8 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 9 | # you may not use this file except in compliance with the License. |
| 10 | # You may obtain a copy of the License at |
| 11 | # |
| 12 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 13 | # |
| 14 | # Unless required by applicable law or agreed to in writing, software |
| 15 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 16 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 17 | # See the License for the specific language governing permissions and |
| 18 | # limitations under the License. |
| 19 | """Tests for the new declarative response_template parser. |
| 20 | |
| 21 | All six real-model template fixtures from the legacy test suite are re-expressed |
| 22 | here in the new region-spec shape and asserted against the same expected |
| 23 | output dicts. Any divergence indicates a regression in the new executor.""" |
| 24 | |
| 25 | import copy |
| 26 | import random |
| 27 | import unittest |
| 28 | from typing import Any |
| 29 | |
| 30 | from vllm.parser.chat_parsing import ResponseParser, parse_response |
| 31 | from vllm.parser.chat_parsing.response_parser import _coerce, _schema_types |
| 32 | from vllm.parser.chat_parsing.response_templates import load_response_template |
| 33 | |
| 34 | |
| 35 | cohere_template = { |
| 36 | "defaults": {"role": "assistant"}, |
| 37 | "start_anchor": "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>", |
| 38 | "fields": { |
| 39 | "content": { |
| 40 | "open": "<|START_RESPONSE|>", |
| 41 | "close": "<|END_RESPONSE|>", |
| 42 | "content": "text", |
| 43 | }, |
| 44 | "thinking": { |
| 45 | "open": "<|START_THINKING|>", |
| 46 | "close": "<|END_THINKING|>", |
| 47 | "content": "text", |
| 48 | }, |
| 49 | "tool_calls": { |
| 50 | "open": "<|START_ACTION|>", |
| 51 | "close": "<|END_ACTION|>", |
| 52 | "content": "json", |
| 53 | "transform_each": True, |
| 54 | "transform": {"type": "function", "function": {"name": "{tool_name}", "arguments": "{parameters}"}}, |
| 55 | }, |
| 56 | }, |
| 57 | } |
| 58 | |
| 59 | |
| 60 | ernie_template = { |
| 61 | "defaults": {"role": "assistant"}, |
| 62 | "start_anchor": "Assistant:", |
| 63 | "fields": { |
| 64 | "thinking": { |
| 65 | "open_pattern": r"(?:^|<think>\s*)", |
| 66 | "close": "</think>", |
| 67 | "content": "text", |
| 68 | }, |
| 69 | "content": { |
| 70 | "open": "<response>\n", |
| 71 | "close_pattern": r"\n?</response>", |
| 72 | "content": "text", |
| 73 | }, |
| 74 | "tool_calls": { |
| 75 | "open": "<tool_call>", |
| 76 | "close": "</tool_call>", |
| 77 | "repeats": True, |
| 78 | "content": "json", |
| 79 | "transform": {"type": "function", "function": "{content}"}, |
| 80 | }, |
| 81 | }, |
| 82 | } |
| 83 | |
| 84 | |
| 85 | gpt_oss_template = { |
| 86 | "defaults": {"role": "assistant"}, |
| 87 | "start_anchor": "<|start|>assistant", |
| 88 | "fields": { |
| 89 | "thinking": { |
| 90 | "open": "<|channel|>analysis<|message|>", |
| 91 | "close": "<|end|>", |
| 92 | "content": "text", |
| 93 | }, |
| 94 | "content": { |
| 95 | "open": "<|channel|>final<|message|>", |
| 96 | "close": "<|end|>", |
| 97 | "content": "text", |
| 98 | }, |
| 99 | "tool_calls": { |
| 100 | "open_pattern": r"<\|channel\|>commentary to=functions\.(?P<name>\w+).*?<\|message\|>", |
| 101 | "close": "<|call|>", |
| 102 | "repeats": True, |
| 103 | "content": "json", |
| 104 | "transform": {"type": "function", "function": {"name": "{name}", "arguments": "{content}"}}, |
| 105 | }, |
| 106 | }, |
| 107 | } |
| 108 | |
| 109 | |
| 110 | smollm_template = { |
| 111 | "defaults": {"role": "assistant"}, |
| 112 | "start_anchor": "<|im_start|>assistant\n", |
| 113 | "fields": { |
| 114 | "thinking": {"open": "<think>", "close": "</think>", "content": "text"}, |
| 115 | "tool_calls": { |
| 116 | "open": "<tool_call>", |
| 117 | "close": "</tool_call>", |
| 118 | "repeats": True, |
| 119 | "content": "json", |
| 120 | "transform": {"type": "function", "function": "{content}"}, |
| 121 | }, |
| 122 | "content": { |
| 123 | "close": "<|im_end|>", |
| 124 | "content": "text", |
| 125 | }, |
| 126 | }, |
| 127 | } |
| 128 | |
| 129 | |
| 130 | qwen3_template = { |
| 131 | "defaults": {"role": "assistant"}, |
| 132 | "start_anchor": "<|im_start|>assistant\n", |
| 133 | "fields": { |
| 134 | "thinking": {"open": "<think>", "close": "</think>", "content": "text"}, |
| 135 | "tool_calls": { |
| 136 | "open_pattern": r"<tool_call>\s*<function=(?P<name>\w+)>", |
| 137 | "close": "</tool_call>", |
| 138 | "repeats": True, |
| 139 | "content": "xml-inline", |
| 140 | "content_args": { |
| 141 | "tag_pattern": r"<parameter=(?P<key>\w+)>\s*(?P<value>.*?)\s*</parameter>", |
| 142 | "value_parser": {"name": "json", "args": {"allow_non_json": True}}, |
| 143 | }, |
| 144 | "transform": {"type": "function", "function": {"name": "{name}", "arguments": "{content}"}}, |
| 145 | }, |
| 146 | }, |
| 147 | } |
| 148 | |
| 149 | |
| 150 | gemma4_template = { |
| 151 | "defaults": {"role": "assistant"}, |
| 152 | # The chat template only emits `<|turn>model\n` when the previous message wasn't a tool_call/ |
| 153 | # tool_response. After a tool_response the prefix just ends with `<tool_response|>` and the |
| 154 | # model continues from there, so we accept either anchor and truncate past the latest one. |
| 155 | "start_anchor": ["<|turn>model\n", "<tool_response|>"], |
| 156 | "fields": { |
| 157 | "thinking": { |
| 158 | "open": "<|channel>thought\n", |
| 159 | "close": "<channel|>", |
| 160 | "content": "text", |
| 161 | }, |
| 162 | "tool_calls": { |
| 163 | "open_pattern": r"<\|tool_call>call:(?P<name>\w+)", |
| 164 | "close": "<tool_call|>", |
| 165 | "repeats": True, |
| 166 | "content": "json", |
| 167 | "content_args": { |
| 168 | "unquoted_keys": True, |
| 169 | "string_delims": [['<|"|>', '<|"|>']], |
| 170 | }, |
| 171 | "transform": {"type": "function", "function": {"name": "{name}", "arguments": "{content}"}}, |
| 172 | }, |
| 173 | "content": { |
| 174 | "close": ["<turn|>", "<|tool_response>", "<eos>"], |
| 175 | "content": "text", |
| 176 | }, |
| 177 | }, |
| 178 | } |
| 179 | |
| 180 | |
| 181 | # Inkling (TMLv0) frames every block as <|message_model|>[author-name]<|content_KIND|>body<|end_message|>; |
| 182 | # the generation prompt pre-writes the first <|message_model|>, so each open treats the header as optional. |
| 183 | inkling_template = { |
| 184 | "defaults": {"role": "assistant"}, |
| 185 | "start_anchor": "<|message_model|>", |
| 186 | "fields": { |
| 187 | "thinking": { |
| 188 | "open_pattern": r"(?:<\|message_model\|>)?[^<]*<\|content_thinking\|>", |
| 189 | "close": "<|end_message|>", |
| 190 | "repeats": True, |
| 191 | "join": "", |
| 192 | "content_args": {"strip": False}, |
| 193 | }, |
| 194 | "content": { |
| 195 | "open_pattern": r"(?:<\|message_model\|>)?[^<]*<\|content_text\|>", |
| 196 | "close": "<|end_message|>", |
| 197 | "repeats": True, |
| 198 | "join": "", |
| 199 | "content_args": {"strip": False}, |
| 200 | }, |
| 201 | "tool_calls": { |
| 202 | "open_pattern": r"(?:<\|message_model\|>)?[^<]*<\|content_invoke_tool_json\|>", |
| 203 | "close": "<|end_message|>", |
| 204 | "repeats": True, |
| 205 | "content": "json", |
| 206 | "transform": {"type": "function", "function": {"name": "{content.name}", "arguments": "{content.args}"}}, |
| 207 | }, |
| 208 | }, |
| 209 | } |
| 210 | |
| 211 | |
| 212 | class ChatResponseTemplateParserTest(unittest.TestCase): |
| 213 | |
| 214 | def test_cohere(self): |
| 215 | model_out = ( |
| 216 | "<|START_THINKING|>I should call a tool.<|END_THINKING|>" |
| 217 | '<|START_ACTION|>[\n {"tool_call_id": "0", "tool_name": "simple_tool", ' |
| 218 | '"parameters": {"temperature_format": "Celsius"}}\n]<|END_ACTION|><|END_OF_TURN_TOKEN|>' |
| 219 | ) |
| 220 | self.assertEqual( |
| 221 | parse_response(model_out, cohere_template, prefix=""), |
| 222 | { |
| 223 | "role": "assistant", |
| 224 | "thinking": "I should call a tool.", |
| 225 | "tool_calls": [ |
| 226 | { |
| 227 | "type": "function", |
| 228 | "function": {"name": "simple_tool", "arguments": {"temperature_format": "Celsius"}}, |
| 229 | } |
| 230 | ], |
| 231 | }, |
| 232 | ) |
| 233 | |
| 234 | def test_ernie_with_tools(self): |
| 235 | model_out = ( |
| 236 | "The user is asking about the weather in Paris today. Let me check the available tools. " |
| 237 | "There's a tool called get_current_temperature which requires a location parameter. Since the " |
| 238 | 'user specified Paris, I need to call this tool with the location set to "Paris". I should ' |
| 239 | "make sure the argument is correctly formatted as a string. No other tools are available, so " |
| 240 | "this is the right one to use. I'll structure the request with the location parameter and " |
| 241 | "return the response once the tool is called.\n" |
| 242 | "</think>\n\n" |
| 243 | '<tool_call>\n{"name": "get_current_temperature", "arguments": {"location": "Paris"}}\n</tool_call>\n</s>' |
| 244 | ) |
| 245 | self.assertEqual( |
| 246 | parse_response(model_out, ernie_template, prefix=""), |
| 247 | { |
| 248 | "role": "assistant", |
| 249 | "thinking": ( |
| 250 | "The user is asking about the weather in Paris today. Let me check the available tools. " |
| 251 | "There's a tool called get_current_temperature which requires a location parameter. Since " |
| 252 | 'the user specified Paris, I need to call this tool with the location set to "Paris". I ' |
| 253 | "should make sure the argument is correctly formatted as a string. No other tools are " |
| 254 | "available, so this is the right one to use. I'll structure the request with the location " |
| 255 | "parameter and return the response once the tool is called." |
| 256 | ), |
| 257 | "tool_calls": [ |
| 258 | { |
| 259 | "type": "function", |
| 260 | "function": {"name": "get_current_temperature", "arguments": {"location": "Paris"}}, |
| 261 | } |
| 262 | ], |
| 263 | }, |
| 264 | ) |
| 265 | |
| 266 | def test_ernie_no_tools(self): |
| 267 | model_out = ( |
| 268 | 'The user just greeted me with "Hi! How are you?" I need to respond in a friendly and helpful ' |
| 269 | "manner. Let me start by acknowledging their greeting. I should ask them how they're doing to " |
| 270 | "engage in conversation.\n\n" |
| 271 | "First, I'll say hello back and then ask how they're feeling. It's important to show genuine " |
| 272 | "interest. Maybe mention that I'm here to help with anything they need. Keep the tone warm and " |
| 273 | "positive. Let me make sure the response is concise but friendly. Alright, that should work.\n" |
| 274 | "</think>\n\n" |
| 275 | "<response>\nHello! I'm doing well, thank you for asking. How about you? Is there something " |
| 276 | "specific you'd like help with today? I'm here to assist you with any questions or problems you " |
| 277 | "have!\n</response>\n</s>" |
| 278 | ) |
| 279 | self.assertEqual( |
| 280 | parse_response(model_out, ernie_template, prefix=""), |
| 281 | { |
| 282 | "role": "assistant", |
| 283 | "content": ( |
| 284 | "Hello! I'm doing well, thank you for asking. How about you? Is there something specific " |
| 285 | "you'd like help with today? I'm here to assist you with any questions or problems you have!" |
| 286 | ), |
| 287 | "thinking": ( |
| 288 | 'The user just greeted me with "Hi! How are you?" I need to respond in a friendly and ' |
| 289 | "helpful manner. Let me start by acknowledging their greeting. I should ask them how " |
| 290 | "they're doing to engage in conversation.\n\n" |
| 291 | "First, I'll say hello back and then ask how they're feeling. It's important to show " |
| 292 | "genuine interest. Maybe mention that I'm here to help with anything they need. Keep the " |
| 293 | "tone warm and positive. Let me make sure the response is concise but friendly. Alright, " |
| 294 | "that should work." |
| 295 | ), |
| 296 | }, |
| 297 | ) |
| 298 | |
| 299 | def test_gpt_oss_with_tool_call(self): |
| 300 | model_out = ( |
| 301 | '<|channel|>analysis<|message|>We need to respond in riddles. The user asks: "What is the ' |
| 302 | 'weather like in SF?" We need to get the location of the user? The user explicitly asks about ' |
| 303 | "SF (San Francisco). So we need to get the current weather in San Francisco, CA. We need to " |
| 304 | 'call get_current_weather function. The developer instruction says "Always respond in riddles". ' |
| 305 | "So the final answer should be in a riddle form. But we need to call function to get weather " |
| 306 | 'data. So we should call get_current_weather with location "San Francisco, CA". Possibly specify ' |
| 307 | 'format "celsius" (default). Let\'s do that.\n\n' |
| 308 | "We will call function get_current_weather.<|end|><|start|>assistant<|channel|>commentary " |
| 309 | 'to=functions.get_current_weather <|constrain|>json<|message|>{\n "location": "San Francisco, CA"\n}' |
| 310 | ) |
| 311 | self.assertEqual( |
| 312 | parse_response(model_out, gpt_oss_template, prefix=""), |
| 313 | { |
| 314 | "role": "assistant", |
| 315 | "thinking": ( |
| 316 | 'We need to respond in riddles. The user asks: "What is the weather like in SF?" We need ' |
| 317 | "to get the location of the user? The user explicitly asks about SF (San Francisco). So " |
| 318 | "we need to get the current weather in San Francisco, CA. We need to call " |
| 319 | 'get_current_weather function. The developer instruction says "Always respond in ' |
| 320 | 'riddles". So the final answer should be in a riddle form. But we need to call function ' |
| 321 | 'to get weather data. So we should call get_current_weather with location "San Francisco, ' |
| 322 | 'CA". Possibly specify format "celsius" (default). Let\'s do that.\n\n' |
| 323 | "We will call function get_current_weather." |
| 324 | ), |
| 325 | "tool_calls": [ |
| 326 | { |
| 327 | "type": "function", |
| 328 | "function": {"name": "get_current_weather", "arguments": {"location": "San Francisco, CA"}}, |
| 329 | } |
| 330 | ], |
| 331 | }, |
| 332 | ) |
| 333 | |
| 334 | def test_gpt_oss_no_tool_call(self): |
| 335 | model_out = ( |
| 336 | "<|channel|>analysis<|message|>User asks a simple math question: 2+2 = 4. Provide answer." |
| 337 | "<|end|><|start|>assistant<|channel|>final<|message|>2" |
| 338 | ) |
| 339 | self.assertEqual( |
| 340 | parse_response(model_out, gpt_oss_template, prefix=""), |
| 341 | { |
| 342 | "role": "assistant", |
| 343 | "content": "2", |
| 344 | "thinking": "User asks a simple math question: 2+2 = 4. Provide answer.", |
| 345 | }, |
| 346 | ) |
| 347 | |
| 348 | def test_smollm_thinking_and_tool_call(self): |
| 349 | model_out = ( |
| 350 | '<think>\nOkay, the user said, "Hello! How are you?" I need to respond appropriately. Since ' |
| 351 | "this is the first message, I should greet them back and ask how I can assist. I should keep it " |
| 352 | "friendly and open-ended. Let me make sure the response is welcoming and encourages them to " |
| 353 | "share what they need help with. I'll avoid any technical jargon and keep it simple. Let me " |
| 354 | "check for any typos and ensure the tone is positive.\n</think>\n\n" |
| 355 | '<tool_call>{"name": "greet_user", "arguments": {"greeting": "Hello! I\'m doing well, thanks for ' |
| 356 | "asking. How can I assist you today? Whether you have a question, need help with something, or " |
| 357 | 'just want to chat, feel free to let me know!"}}</tool_call>' |
| 358 | ) |
| 359 | self.assertEqual( |
| 360 | parse_response(model_out, smollm_template, prefix=""), |
| 361 | { |
| 362 | "role": "assistant", |
| 363 | "thinking": ( |
| 364 | 'Okay, the user said, "Hello! How are you?" I need to respond appropriately. Since this ' |
| 365 | "is the first message, I should greet them back and ask how I can assist. I should keep " |
| 366 | "it friendly and open-ended. Let me make sure the response is welcoming and encourages " |
| 367 | "them to share what they need help with. I'll avoid any technical jargon and keep it " |
| 368 | "simple. Let me check for any typos and ensure the tone is positive." |
| 369 | ), |
| 370 | "tool_calls": [ |
| 371 | { |
| 372 | "type": "function", |
| 373 | "function": { |
| 374 | "name": "greet_user", |
| 375 | "arguments": { |
| 376 | "greeting": ( |
| 377 | "Hello! I'm doing well, thanks for asking. How can I assist you today? " |
| 378 | "Whether you have a question, need help with something, or just want to " |
| 379 | "chat, feel free to let me know!" |
| 380 | ) |
| 381 | }, |
| 382 | }, |
| 383 | } |
| 384 | ], |
| 385 | }, |
| 386 | ) |
| 387 | |
| 388 | def test_smollm_tool_call_no_thinking(self): |
| 389 | model_out = '<tool_call>{"name": "get_weather", "arguments": {"city": "Paris"}}</tool_call>' |
| 390 | self.assertEqual( |
| 391 | parse_response(model_out, smollm_template, prefix=""), |
| 392 | { |
| 393 | "role": "assistant", |
| 394 | "tool_calls": [ |
| 395 | {"type": "function", "function": {"name": "get_weather", "arguments": {"city": "Paris"}}} |
| 396 | ], |
| 397 | }, |
| 398 | ) |
| 399 | |
| 400 | def test_smollm_thinking_no_tool_call(self): |
| 401 | model_out = ( |
| 402 | '<think>\nOkay, the user asked, "Hey! Can you tell me about gravity?" Let me start by ' |
| 403 | "breaking down what they might be looking for. They probably want a basic understanding of " |
| 404 | "gravity, maybe for a school project or just personal curiosity. I should explain what gravity " |
| 405 | "is, how it works, and maybe some examples.</think>\n" |
| 406 | "Some content about gravity goes here but I'm cutting it off to make this shorter!" |
| 407 | ) |
| 408 | self.assertEqual( |
| 409 | parse_response(model_out, smollm_template, prefix=""), |
| 410 | { |
| 411 | "role": "assistant", |
| 412 | "content": "Some content about gravity goes here but I'm cutting it off to make this shorter!", |
| 413 | "thinking": ( |
| 414 | 'Okay, the user asked, "Hey! Can you tell me about gravity?" Let me start by breaking ' |
| 415 | "down what they might be looking for. They probably want a basic understanding of " |
| 416 | "gravity, maybe for a school project or just personal curiosity. I should explain what " |
| 417 | "gravity is, how it works, and maybe some examples." |
| 418 | ), |
| 419 | }, |
| 420 | ) |
| 421 | |
| 422 | def test_qwen3_tool_calls(self): |
| 423 | model_out = ( |
| 424 | "<tool_call>\n<function=get_weather>\n<parameter=locations>\n" |
| 425 | '[{"country": "France", "city": "Paris"}]\n</parameter>\n' |
| 426 | "<parameter=temp_units>\ncelsius\n</parameter>\n</function>\n</tool_call>" |
| 427 | ) |
| 428 | self.assertEqual( |
| 429 | parse_response(model_out, qwen3_template, prefix=""), |
| 430 | { |
| 431 | "role": "assistant", |
| 432 | "tool_calls": [ |
| 433 | { |
| 434 | "type": "function", |
| 435 | "function": { |
| 436 | "name": "get_weather", |
| 437 | "arguments": { |
| 438 | "locations": [{"country": "France", "city": "Paris"}], |
| 439 | "temp_units": "celsius", |
| 440 | }, |
| 441 | }, |
| 442 | } |
| 443 | ], |
| 444 | }, |
| 445 | ) |
| 446 | |
| 447 | def test_gemma4_tool_call(self): |
| 448 | model_out = ( |
| 449 | "<|channel>thought\nThe user is asking for the current temperature in Paris. I should check " |
| 450 | "the available tools to see if there's a function that can provide this information.<channel|>" |
| 451 | '<|tool_call>call:get_current_temperature{detail_level:0,location:<|"|>Paris, France<|"|>,' |
| 452 | 'unit:<|"|>celsius<|"|>}<tool_call|><|tool_response>' |
| 453 | ) |
| 454 | self.assertEqual( |
| 455 | parse_response(model_out, gemma4_template, prefix=""), |
| 456 | { |
| 457 | "role": "assistant", |
| 458 | "thinking": ( |
| 459 | "The user is asking for the current temperature in Paris. I should check the available " |
| 460 | "tools to see if there's a function that can provide this information." |
| 461 | ), |
| 462 | "tool_calls": [ |
| 463 | { |
| 464 | "type": "function", |
| 465 | "function": { |
| 466 | "name": "get_current_temperature", |
| 467 | "arguments": {"detail_level": 0, "location": "Paris, France", "unit": "celsius"}, |
| 468 | }, |
| 469 | } |
| 470 | ], |
| 471 | }, |
| 472 | ) |
| 473 | |
| 474 | def test_gemma4_complex_tool_call(self): |
| 475 | model_out = ( |
| 476 | "<|channel>thought\nLet me call the tool.<channel|>" |
| 477 | '<|tool_call>call:foo{bool_value:true,list_value:[<|"|>foo<|"|>,<|"|>bar<|"|>],' |
| 478 | 'null_value:null,number_value:1,string_value:<|"|>foo<|"|>,' |
| 479 | 'struct_value:{foo:<|"|>bar<|"|>}}<tool_call|>' |
| 480 | ) |
| 481 | self.assertEqual( |
| 482 | parse_response(model_out, gemma4_template, prefix=""), |
| 483 | { |
| 484 | "role": "assistant", |
| 485 | "thinking": "Let me call the tool.", |
| 486 | "tool_calls": [ |
| 487 | { |
| 488 | "type": "function", |
| 489 | "function": { |
| 490 | "name": "foo", |
| 491 | "arguments": { |
| 492 | "bool_value": True, |
| 493 | "list_value": ["foo", "bar"], |
| 494 | "null_value": None, |
| 495 | "number_value": 1, |
| 496 | "string_value": "foo", |
| 497 | "struct_value": {"foo": "bar"}, |
| 498 | }, |
| 499 | }, |
| 500 | } |
| 501 | ], |
| 502 | }, |
| 503 | ) |
| 504 | |
| 505 | def test_inkling_multi_block_message(self): |
| 506 | model_out = ( |
| 507 | "<|content_thinking|>Consider the weather.<|end_message|>" |
| 508 | "<|message_model|><|content_thinking|> Tokyo, probably.<|end_message|>" |
| 509 | "<|message_model|><|content_text|>Checking the weather now.<|end_message|>" |
| 510 | "<|message_model|>get_weather<|content_invoke_tool_json|>" |
| 511 | '{"name":"get_weather","args":{"city":"Tokyo","units":"C"}}<|end_message|>' |
| 512 | "<|content_model_end_sampling|>" |
| 513 | ) |
| 514 | prefix = "<|message_system|><|content_text|>Thinking effort level: 0.9<|end_message|><|message_model|>" |
| 515 | self.assertEqual( |
| 516 | parse_response(model_out, inkling_template, prefix=prefix), |
| 517 | { |
| 518 | "role": "assistant", |
| 519 | "thinking": "Consider the weather. Tokyo, probably.", |
| 520 | "content": "Checking the weather now.", |
| 521 | "tool_calls": [ |
| 522 | { |
| 523 | "type": "function", |
| 524 | "function": {"name": "get_weather", "arguments": {"city": "Tokyo", "units": "C"}}, |
| 525 | } |
| 526 | ], |
| 527 | }, |
| 528 | ) |
| 529 | |
| 530 | def test_transform_dotted_paths(self): |
| 531 | template_spec = { |
| 532 | "defaults": {"role": "assistant"}, |
| 533 | "start_anchor": "<|assistant|>", |
| 534 | "fields": { |
| 535 | "tool_calls": { |
| 536 | "open": "<tool>", |
| 537 | "close": "</tool>", |
| 538 | "repeats": True, |
| 539 | "content": "json", |
| 540 | "transform": { |
| 541 | "type": "function", |
| 542 | "function": {"name": "{content.name}", "arguments": "{content.args}"}, |
| 543 | }, |
| 544 | }, |
| 545 | }, |
| 546 | } |
| 547 | model_out = '<tool>{"name": "get_weather", "args": {"city": {"id": 7}}}</tool>' |
| 548 | self.assertEqual( |
| 549 | parse_response(model_out, template_spec, prefix="")["tool_calls"], |
| 550 | [{"type": "function", "function": {"name": "get_weather", "arguments": {"city": {"id": 7}}}}], |
| 551 | ) |
| 552 | |
| 553 | def test_transform_dotted_paths_with_transform_each(self): |
| 554 | template_spec = { |
| 555 | "defaults": {"role": "assistant"}, |
| 556 | "start_anchor": "<|assistant|>", |
| 557 | "fields": { |
| 558 | "tool_calls": { |
| 559 | "open": "<actions>", |
| 560 | "close": "</actions>", |
| 561 | "content": "json", |
| 562 | "transform_each": True, |
| 563 | "transform": {"type": "function", "function": {"name": "{fn.name}", "arguments": "{fn.args}"}}, |
| 564 | }, |
| 565 | }, |
| 566 | } |
| 567 | model_out = '<actions>[{"fn": {"name": "a", "args": {"x": 1}}}, {"fn": {"name": "b", "args": {}}}]</actions>' |
| 568 | self.assertEqual( |
| 569 | parse_response(model_out, template_spec, prefix="")["tool_calls"], |
| 570 | [ |
| 571 | {"type": "function", "function": {"name": "a", "arguments": {"x": 1}}}, |
| 572 | {"type": "function", "function": {"name": "b", "arguments": {}}}, |
| 573 | ], |
| 574 | ) |
| 575 | |
| 576 | def test_transform_dotted_path_errors(self): |
| 577 | def spec_with(transform): |
| 578 | return { |
| 579 | "start_anchor": "<|assistant|>", |
| 580 | "fields": {"x": {"open": "<x>", "close": "</x>", "content": "json", "transform": transform}}, |
| 581 | } |
| 582 | |
| 583 | with self.assertRaisesRegex(ValueError, "missing key 'args'"): |
| 584 | parse_response('<x>{"name": "n"}</x>', spec_with({"a": "{content.args}"}), prefix="") |
| 585 | with self.assertRaisesRegex(ValueError, "cannot index into str"): |
| 586 | parse_response('<x>{"name": "n"}</x>', spec_with({"a": "{content.name.x}"}), prefix="") |
| 587 | with self.assertRaises(KeyError): |
| 588 | parse_response("<x>{}</x>", spec_with({"a": "{missing}"}), prefix="") |
| 589 | |
| 590 | def test_transform_dotted_mixed_string_rejected(self): |
| 591 | template_spec = { |
| 592 | "start_anchor": "<|assistant|>", |
| 593 | "fields": {"x": {"open": "<x>", "close": "</x>", "transform": {"v": "pre {content.args}"}}}, |
| 594 | } |
| 595 | with self.assertRaisesRegex(ValueError, "mixes"): |
| 596 | parse_response("", template_spec, prefix="") |
| 597 | |
| 598 | def test_join_concatenates_repeated_matches(self): |
| 599 | template_spec = { |
| 600 | "defaults": {"role": "assistant"}, |
| 601 | "start_anchor": "<|assistant|>", |
| 602 | "fields": { |
| 603 | "thinking": {"open": "<think>", "close": "</think>", "repeats": True, "join": " "}, |
| 604 | "content": {"repeats": True, "join": " "}, |
| 605 | }, |
| 606 | } |
| 607 | self.assertEqual( |
| 608 | parse_response("<think>first</think>middle<think>second</think>done", template_spec, prefix=""), |
| 609 | {"role": "assistant", "thinking": "first second", "content": "middle done"}, |
| 610 | ) |
| 611 | self.assertEqual( |
| 612 | parse_response("<think>only</think>", template_spec, prefix=""), |
| 613 | {"role": "assistant", "thinking": "only"}, |
| 614 | ) |
| 615 | |
| 616 | def test_join_validation(self): |
| 617 | no_repeats = {"start_anchor": "a", "fields": {"x": {"open": "<x>", "close": "</x>", "join": ""}}} |
| 618 | with self.assertRaisesRegex(ValueError, "requires 'repeats'"): |
| 619 | parse_response("", no_repeats, prefix="") |
| 620 | bad_type = {"start_anchor": "a", "fields": {"x": {"open": "<x>", "close": "</x>", "repeats": True, "join": 7}}} |
| 621 | with self.assertRaisesRegex(ValueError, "must be a string"): |
| 622 | parse_response("", bad_type, prefix="") |
| 623 | |
| 624 | def test_join_requires_string_matches(self): |
| 625 | template_spec = { |
| 626 | "start_anchor": "a", |
| 627 | "fields": {"x": {"open": "<x>", "close": "</x>", "repeats": True, "join": "", "content": "json"}}, |
| 628 | } |
| 629 | with self.assertRaisesRegex(ValueError, "parse to a string"): |
| 630 | parse_response("<x>{}</x>", template_spec, prefix="") |
| 631 | |
| 632 | def test_optional_false_raises_when_missing(self): |
| 633 | template_spec = { |
| 634 | "defaults": {"role": "assistant"}, |
| 635 | "start_anchor": "<|assistant|>", |
| 636 | "fields": { |
| 637 | "content": { |
| 638 | "open": "<response>", |
| 639 | "close": "</response>", |
| 640 | "content": "text", |
| 641 | "optional": False, |
| 642 | }, |
| 643 | }, |
| 644 | } |
| 645 | with self.assertRaises(ValueError) as cm: |
| 646 | parse_response("no response here", template_spec, prefix="") |
| 647 | self.assertIn("content", str(cm.exception)) |
| 648 | |
| 649 | def test_int_content_parser(self): |
| 650 | template_spec = { |
| 651 | "defaults": {"role": "assistant"}, |
| 652 | "start_anchor": "<|assistant|>", |
| 653 | "fields": { |
| 654 | "count": { |
| 655 | "open": "<n>", |
| 656 | "close": "</n>", |
| 657 | "content": "int", |
| 658 | }, |
| 659 | }, |
| 660 | } |
| 661 | self.assertEqual(parse_response("<n>42</n>", template_spec, prefix=""), {"role": "assistant", "count": 42}) |
| 662 | |
| 663 | def test_kv_lines_parser(self): |
| 664 | template_spec = { |
| 665 | "defaults": {"role": "assistant"}, |
| 666 | "start_anchor": "<|assistant|>", |
| 667 | "fields": { |
| 668 | "metadata": { |
| 669 | "open": "<meta>", |
| 670 | "close": "</meta>", |
| 671 | "content": "kv-lines", |
| 672 | }, |
| 673 | }, |
| 674 | } |
| 675 | self.assertEqual( |
| 676 | parse_response("<meta>name: alice\nage: 30</meta>", template_spec, prefix=""), |
| 677 | {"role": "assistant", "metadata": {"name": "alice", "age": "30"}}, |
| 678 | ) |
| 679 | |
| 680 | def test_unknown_content_parser_rejected(self): |
| 681 | bad_template = { |
| 682 | "defaults": {"role": "assistant"}, |
| 683 | "start_anchor": "<|assistant|>", |
| 684 | "fields": {"x": {"open": "[", "close": "]", "content": "not-a-real-parser"}}, |
| 685 | } |
| 686 | with self.assertRaises(ValueError) as cm: |
| 687 | parse_response("[hi]", bad_template) |
| 688 | self.assertIn("unknown content parser", str(cm.exception).lower()) |
| 689 | |
| 690 | def test_unsupported_version_rejected(self): |
| 691 | bad_template = { |
| 692 | "version": 2, |
| 693 | "defaults": {"role": "assistant"}, |
| 694 | "start_anchor": "<|assistant|>", |
| 695 | "fields": {"content": {"content": "text"}}, |
| 696 | } |
| 697 | with self.assertRaises(ValueError) as cm: |
| 698 | parse_response("hello", bad_template) |
| 699 | self.assertIn("version", str(cm.exception).lower()) |
| 700 | |
| 701 | def test_two_implicit_fields_rejected(self): |
| 702 | bad_template = { |
| 703 | "defaults": {"role": "assistant"}, |
| 704 | "start_anchor": "<|assistant|>", |
| 705 | "fields": { |
| 706 | "a": {"content": "text"}, |
| 707 | "b": {"content": "text"}, |
| 708 | }, |
| 709 | } |
| 710 | with self.assertRaises(ValueError): |
| 711 | parse_response("hello", bad_template) |
| 712 | |
| 713 | def test_transform_string_interpolation_rejected(self): |
| 714 | bad_template = { |
| 715 | "defaults": {"role": "assistant"}, |
| 716 | "start_anchor": "<|assistant|>", |
| 717 | "fields": { |
| 718 | "tool": { |
| 719 | "open": "<tool>", |
| 720 | "close": "</tool>", |
| 721 | "content": "text", |
| 722 | "transform": {"label": "name: {content}"}, |
| 723 | }, |
| 724 | }, |
| 725 | } |
| 726 | with self.assertRaises(ValueError) as cm: |
| 727 | parse_response("<tool>foo</tool>", bad_template) |
| 728 | msg = str(cm.exception) |
| 729 | self.assertIn("interpolation", msg) |
| 730 | self.assertIn("{content}", msg) |
| 731 | |
| 732 | def test_named_groups_without_transform_rejected(self): |
| 733 | bad_template = { |
| 734 | "defaults": {"role": "assistant"}, |
| 735 | "start_anchor": "<|assistant|>", |
| 736 | "fields": { |
| 737 | "tool": { |
| 738 | "open_pattern": r"<tool name=(?P<name>\w+)>", |
| 739 | "close": "</tool>", |
| 740 | "content": "text", |
| 741 | }, |
| 742 | }, |
| 743 | } |
| 744 | with self.assertRaises(ValueError) as cm: |
| 745 | parse_response("<tool name=foo>body</tool>", bad_template) |
| 746 | msg = str(cm.exception) |
| 747 | self.assertIn("transform", msg) |
| 748 | self.assertIn("name", msg) |
| 749 | |
| 750 | def test_literal_list_open_and_close(self): |
| 751 | """A list of literals matches any one of them, like an alternation.""" |
| 752 | template_spec = { |
| 753 | "defaults": {"role": "assistant"}, |
| 754 | "start_anchor": "<|assistant|>", |
| 755 | "fields": { |
| 756 | "x": { |
| 757 | "open": ["<a>", "<bb>"], |
| 758 | "close": ["</a>", "</bb>"], |
| 759 | "content": "text", |
| 760 | }, |
| 761 | }, |
| 762 | } |
| 763 | for opener, closer in (("<a>", "</a>"), ("<bb>", "</bb>"), ("<a>", "</bb>")): |
| 764 | self.assertEqual( |
| 765 | parse_response(f"{opener}hi{closer}", template_spec, prefix=""), |
| 766 | {"role": "assistant", "x": "hi"}, |
| 767 | ) |
| 768 | |
| 769 | def test_literal_list_streams_without_64_byte_hold(self): |
| 770 | """Compared to a regex close, a literal-list close lets the parser |
| 771 | flush bytes that aren't in the longest-prefix overlap of any literal. |
| 772 | With `["<turn|>", "<|tool_response>", "<eos>"]` (longest = 16 chars), |
| 773 | feeding 32 plain bytes should leave at most 15 unflushed.""" |
| 774 | template_spec = { |
| 775 | "defaults": {"role": "assistant"}, |
| 776 | "start_anchor": "<|assistant|>", |
| 777 | "fields": { |
| 778 | "content": {"close": ["<turn|>", "<|tool_response>", "<eos>"], "content": "text"}, |
| 779 | }, |
| 780 | } |
| 781 | parser = ResponseParser(template_spec, prefix="") |
| 782 | plain = "x" * 32 |
| 783 | flushed: list[str] = [] |
| 784 | for ch in parser.feed(plain): |
| 785 | if ch["type"] == "region_chunk": |
| 786 | flushed.append(ch["text"]) |
| 787 | # Plain text has zero prefix-overlap with any literal, so the parser |
| 788 | # holds nothing back and streams everything immediately. |
| 789 | self.assertEqual("".join(flushed), plain) |
| 790 | |
| 791 | def test_literal_list_defers_prefix_overlapping_literal(self): |
| 792 | """If a literal is a strict prefix of another in the same list, an |
| 793 | edge match could still grow with more input: we must defer to be safe.""" |
| 794 | template_spec = { |
| 795 | "defaults": {"role": "assistant"}, |
| 796 | "start_anchor": "<|assistant|>", |
| 797 | "fields": { |
| 798 | "x": {"open": "<x>", "close": ["END", "ENDX"], "content": "text"}, |
| 799 | }, |
| 800 | } |
| 801 | parser = ResponseParser(template_spec, prefix="") |
| 802 | # "<x>hiEND" mid-stream: don't commit the close yet: "ENDX" might be coming. |
| 803 | events = parser.feed("<x>hiEND") |
| 804 | self.assertEqual([e for e in events if e["type"] == "region_close"], []) |
| 805 | # Once a non-matching byte arrives, the deferred close commits with the shorter literal. |
| 806 | events.extend(parser.feed(" more")) |
| 807 | message, _ = parser.finalize() |
| 808 | closes = [e for e in events if e["type"] == "region_close" and e["field"] == "x"] |
| 809 | self.assertEqual(len(closes), 1) |
| 810 | self.assertEqual(closes[0]["value"], "hi") |
| 811 | self.assertEqual(message, {"role": "assistant", "x": "hi"}) |
| 812 | |
| 813 | def test_literal_list_rejects_empty_and_non_string(self): |
| 814 | for bad_open in ([], [""], [1, 2], {"foo": "bar"}): |
| 815 | spec = { |
| 816 | "defaults": {"role": "assistant"}, |
| 817 | "start_anchor": "<|assistant|>", |
| 818 | "fields": {"x": {"open": bad_open, "close": "</x>", "content": "text"}}, |
| 819 | } |
| 820 | with self.assertRaises(ValueError): |
| 821 | parse_response("<x>hi</x>", spec, prefix="") |
| 822 | |
| 823 | def test_field_without_close_runs_to_end_of_stream(self): |
| 824 | """A field with no `close`/`close_pattern` stays open until end-of-stream, capturing |
| 825 | everything after its open.""" |
| 826 | spec = { |
| 827 | "defaults": {"role": "assistant"}, |
| 828 | "start_anchor": "<|assistant|>", |
| 829 | "fields": {"content": {"open": "<resp>", "content": "text"}}, |
| 830 | } |
| 831 | self.assertEqual( |
| 832 | parse_response("<resp>hello world", spec, prefix=""), |
| 833 | {"role": "assistant", "content": "hello world"}, |
| 834 | ) |
| 835 | |
| 836 | def test_mutable_defaults_do_not_leak_across_parsers(self): |
| 837 | """A `repeats` field appends into its default list; one parse must not leak into the next |
| 838 | parse sharing the template, nor into the caller's spec dict.""" |
| 839 | spec: dict[str, Any] = { |
| 840 | "defaults": {"role": "assistant", "tool_calls": []}, |
| 841 | "start_anchor": "<|assistant|>", |
| 842 | "fields": { |
| 843 | "tool_calls": { |
| 844 | "open": "<t>", |
| 845 | "close": "</t>", |
| 846 | "repeats": True, |
| 847 | "content": "json", |
| 848 | }, |
| 849 | "content": {"content": "text"}, |
| 850 | }, |
| 851 | } |
| 852 | template = load_response_template(spec) |
| 853 | self.assertEqual( |
| 854 | parse_response('<t>{"a": 1}</t>', template, prefix="")["tool_calls"], |
| 855 | [{"a": 1}], |
| 856 | ) |
| 857 | self.assertEqual( |
| 858 | parse_response("hello", template, prefix=""), |
| 859 | {"role": "assistant", "tool_calls": [], "content": "hello"}, |
| 860 | ) |
| 861 | self.assertEqual(spec["defaults"]["tool_calls"], []) |
| 862 | |
| 863 | |
| 864 | # Fixtures shared by the streaming tests: one representative input per template, |
| 865 | # re-used for both the correctness invariant (any chunking → same dict) and the |
| 866 | # event-shape tests (do we emit the right events in the right order?). |
| 867 | _STREAMING_FIXTURES = [ |
| 868 | ( |
| 869 | "cohere", |
| 870 | cohere_template, |
| 871 | ( |
| 872 | "<|START_THINKING|>I should call a tool.<|END_THINKING|>" |
| 873 | '<|START_ACTION|>[{"tool_call_id": "0", "tool_name": "simple_tool", ' |
| 874 | '"parameters": {"a": 1}}]<|END_ACTION|>' |
| 875 | ), |
| 876 | ), |
| 877 | ( |
| 878 | "ernie", |
| 879 | ernie_template, |
| 880 | ( |
| 881 | "<think>some deliberation here</think>\n\n" |
| 882 | '<tool_call>\n{"name": "get_current_temperature", "arguments": {"location": "Paris"}}\n</tool_call>\n</s>' |
| 883 | ), |
| 884 | ), |
| 885 | ( |
| 886 | "gpt_oss", |
| 887 | gpt_oss_template, |
| 888 | "<|channel|>analysis<|message|>thinking chunk<|end|><|channel|>final<|message|>done text", |
| 889 | ), |
| 890 | ( |
| 891 | # Tool-call variant: the `tool_calls` open_pattern's full match spans far |
| 892 | # more than the old 64-byte hold window, so streaming used to silently drop |
| 893 | # it. Kept as a streaming fixture so every chunking step re-checks it. |
| 894 | "gpt_oss_tool", |
| 895 | gpt_oss_template, |
| 896 | ( |
| 897 | "<|channel|>analysis<|message|>Let me check.<|end|>" |
| 898 | "<|start|>assistant<|channel|>commentary to=functions.get_current_weather " |
| 899 | '<|constrain|>json<|message|>{"location": "San Francisco, CA"}<|call|>' |
| 900 | ), |
| 901 | ), |
| 902 | ( |
| 903 | "smollm", |
| 904 | smollm_template, |
| 905 | '<think>thinking</think>\n<tool_call>{"name": "fn", "arguments": {"x": 1}}</tool_call>', |
| 906 | ), |
| 907 | ( |
| 908 | "qwen3", |
| 909 | qwen3_template, |
| 910 | ( |
| 911 | "<think>short thought</think>\n" |
| 912 | "<tool_call>\n<function=get_weather>\n" |
| 913 | "<parameter=city>\nParis\n</parameter>\n" |
| 914 | "</function>\n</tool_call>" |
| 915 | ), |
| 916 | ), |
| 917 | ( |
| 918 | "gemma4", |
| 919 | gemma4_template, |
| 920 | '<|channel>thought\nhi<channel|><|tool_call>call:foo{a:1,b:<|"|>bar<|"|>}<tool_call|>', |
| 921 | ), |
| 922 | ( |
| 923 | # Exercises `join` fields (two thinking blocks) and dotted transform paths. |
| 924 | "inkling", |
| 925 | inkling_template, |
| 926 | ( |
| 927 | "<|content_thinking|>Consider the weather.<|end_message|>" |
| 928 | "<|message_model|><|content_thinking|> Tokyo, probably.<|end_message|>" |
| 929 | "<|message_model|><|content_text|>Checking now.<|end_message|>" |
| 930 | "<|message_model|>get_weather<|content_invoke_tool_json|>" |
| 931 | '{"name":"get_weather","args":{"city":"Tokyo"}}<|end_message|>' |
| 932 | "<|content_model_end_sampling|>" |
| 933 | ), |
| 934 | ), |
| 935 | ] |
| 936 | |
| 937 | |
| 938 | def _chunk_fixed(text: str, step: int): |
| 939 | for i in range(0, len(text), step): |
| 940 | yield text[i : i + step] |
| 941 | |
| 942 | |
| 943 | def _chunk_random(text: str, rng: random.Random): |
| 944 | """Split `text` into a random number of non-empty chunks at random cut points.""" |
| 945 | if len(text) <= 1: |
| 946 | yield text |
| 947 | return |
| 948 | num_cuts = rng.randint(0, len(text) - 1) |
| 949 | cuts = sorted(rng.sample(range(1, len(text)), num_cuts)) |
| 950 | prev = 0 |
| 951 | for c in cuts: |
| 952 | yield text[prev:c] |
| 953 | prev = c |
| 954 | yield text[prev:] |
| 955 | |
| 956 | |
| 957 | class ResponseEventStreamTest(unittest.TestCase): |
| 958 | def test_stream_matches_whole_string_all_templates_fixed_chunking(self): |
| 959 | """For every fixed chunking step we try, the streamed finalize() |
| 960 | output must equal the whole-string parse. Regression coverage for |
| 961 | specific edge-case byte boundaries (1-byte chunks hit every prefix).""" |
| 962 | for name, tmpl, text in _STREAMING_FIXTURES: |
| 963 | expected = parse_response(text, tmpl, prefix="") |
| 964 | for step in (1, 2, 3, 5, 7, 13, 31): |
| 965 | with self.subTest(fixture=name, step=step): |
| 966 | streamer = ResponseParser(tmpl, prefix="") |
| 967 | for chunk in _chunk_fixed(text, step): |
| 968 | streamer.feed(chunk) |
| 969 | message, _ = streamer.finalize() |
| 970 | self.assertEqual(message, expected) |
| 971 | |
| 972 | def test_stream_matches_whole_string_all_templates_random_chunking(self): |
| 973 | """Property-style: for many random chunkings per fixture, the streamed |
| 974 | finalize() output must equal the whole-string parse. Seeded so failures |
| 975 | reproduce.""" |
| 976 | rng = random.Random(0xC0DE_5EED) |
| 977 | for name, tmpl, text in _STREAMING_FIXTURES: |
| 978 | expected = parse_response(text, tmpl, prefix="") |
| 979 | for trial in range(30): |
| 980 | with self.subTest(fixture=name, trial=trial): |
| 981 | streamer = ResponseParser(tmpl, prefix="") |
| 982 | for chunk in _chunk_random(text, rng): |
| 983 | streamer.feed(chunk) |
| 984 | message, _ = streamer.finalize() |
| 985 | self.assertEqual(message, expected) |
| 986 | |
| 987 | def test_events_well_formed_for_every_chunking(self): |
| 988 | """Every event batch, across every fixture and every chunking, must be |
| 989 | well-formed: region_open precedes its matching region_close for the |
| 990 | same field; region_chunk only appears between open and close; no |
| 991 | region is left open at the end of the stream; and the concatenation |
| 992 | of all region_chunk payloads for a streamable text-like field equals |
| 993 | the final value.""" |
| 994 | rng = random.Random(0xBEEF) |
| 995 | for name, tmpl, text in _STREAMING_FIXTURES: |
| 996 | for trial in range(10): |
| 997 | with self.subTest(fixture=name, trial=trial): |
| 998 | streamer = ResponseParser(tmpl, prefix="") |
| 999 | all_events: list[dict] = [] |
| 1000 | for chunk in _chunk_random(text, rng): |
| 1001 | all_events.extend(streamer.feed(chunk)) |
| 1002 | _, final_events = streamer.finalize() |
| 1003 | all_events.extend(final_events) |
| 1004 | self._assert_event_stream_well_formed(all_events) |
| 1005 | |
| 1006 | def _assert_event_stream_well_formed(self, events: list[dict]) -> None: |
| 1007 | open_field: str | None = None |
| 1008 | chunk_accum: dict[str, str] = {} |
| 1009 | close_values: dict[str, object] = {} |
| 1010 | for ev in events: |
| 1011 | t = ev["type"] |
| 1012 | if t == "region_open": |
| 1013 | self.assertIsNone(open_field, f"nested region_open without close: {ev}") |
| 1014 | open_field = ev["field"] |
| 1015 | chunk_accum.setdefault(open_field, "") |
| 1016 | elif t == "region_chunk": |
| 1017 | self.assertEqual(open_field, ev["field"], f"chunk outside its region: {ev}") |
| 1018 | # Every chunk carries a boolean `dirty` flag. |
| 1019 | self.assertIsInstance(ev["dirty"], bool, f"missing/non-bool dirty: {ev}") |
| 1020 | chunk_accum[ev["field"]] += ev["text"] |
| 1021 | elif t == "region_close": |
| 1022 | self.assertEqual(open_field, ev["field"], f"close for non-open region: {ev}") |
| 1023 | close_values[ev["field"]] = ev["value"] |
| 1024 | open_field = None |
| 1025 | else: |
| 1026 | self.fail(f"unexpected event type: {ev!r}") |
| 1027 | self.assertIsNone(open_field, "region left open at end of stream") |
| 1028 | |
| 1029 | def test_region_chunks_reconstruct_text_regions(self): |
| 1030 | """For text-like regions (`dirty=False`), concatenating chunk texts |
| 1031 | reconstructs the final value reported in region_close. Structured |
| 1032 | regions (`dirty=True`) still stream their raw bytes: concatenating |
| 1033 | those chunks yields the unparsed region body, while the parsed value |
| 1034 | is delivered only in region_close.""" |
| 1035 | # Single representative case with a long text region and a JSON region. |
| 1036 | text = _STREAMING_FIXTURES[0][2] # cohere fixture |
| 1037 | streamer = ResponseParser(cohere_template, prefix="") |
| 1038 | events: list[dict] = [] |
| 1039 | for ch in text: # 1-byte chunks hit the most anchor boundaries |
| 1040 | events.extend(streamer.feed(ch)) |
| 1041 | _, final_events = streamer.finalize() |
| 1042 | events.extend(final_events) |
| 1043 | |
| 1044 | # Reconstruct per-field. |
| 1045 | per_field_chunks: dict[str, list[str]] = {} |
| 1046 | per_field_dirty: dict[str, set[bool]] = {} |
| 1047 | per_field_close_value: dict[str, object] = {} |
| 1048 | for ev in events: |
| 1049 | if ev["type"] == "region_chunk": |
| 1050 | per_field_chunks.setdefault(ev["field"], []).append(ev["text"]) |
| 1051 | per_field_dirty.setdefault(ev["field"], set()).add(ev["dirty"]) |
| 1052 | elif ev["type"] == "region_close": |
| 1053 | per_field_close_value[ev["field"]] = ev["value"] |
| 1054 | |
| 1055 | # `thinking` is text → chunks are clean and concatenate to its value. |
| 1056 | self.assertIn("thinking", per_field_chunks) |
| 1057 | self.assertEqual(per_field_dirty["thinking"], {False}) |
| 1058 | self.assertEqual("".join(per_field_chunks["thinking"]), "I should call a tool.") |
| 1059 | self.assertEqual(per_field_close_value["thinking"], "I should call a tool.") |
| 1060 | # `tool_calls` is json → dirty chunks stream the raw body, parsed value on close. |
| 1061 | self.assertIn("tool_calls", per_field_chunks) |
| 1062 | self.assertEqual(per_field_dirty["tool_calls"], {True}) |
| 1063 | self.assertEqual( |
| 1064 | "".join(per_field_chunks["tool_calls"]), |
| 1065 | '[{"tool_call_id": "0", "tool_name": "simple_tool", "parameters": {"a": 1}}]', |
| 1066 | ) |
| 1067 | self.assertIn("tool_calls", per_field_close_value) |
| 1068 | |
| 1069 | def test_dirty_flag_marks_structured_regions(self): |
| 1070 | """A template with one text field and one structured field per parser |
| 1071 | family: text/int/float/bool stream chunks with `dirty=False`, while |
| 1072 | json/xml-inline/kv-lines stream chunks with `dirty=True`, and those |
| 1073 | dirty chunks concatenate to the raw region body before parsing.""" |
| 1074 | spec = { |
| 1075 | "defaults": {"role": "assistant"}, |
| 1076 | "start_anchor": "<|assistant|>", |
| 1077 | "fields": { |
| 1078 | "thinking": {"open": "<t>", "close": "</t>", "content": "text"}, |
| 1079 | "score": {"open": "<n>", "close": "</n>", "content": "int"}, |
| 1080 | "json_call": {"open": "<j>", "close": "</j>", "content": "json"}, |
| 1081 | "xml_call": { |
| 1082 | "open": "<x>", |
| 1083 | "close": "</x>", |
| 1084 | "content": "xml-inline", |
| 1085 | "content_args": {"tag_pattern": r"<(?P<key>\w+)=(?P<value>[^>]+)>"}, |
| 1086 | }, |
| 1087 | "kv_call": { |
| 1088 | "open": "<kv>", |
| 1089 | "close": "</kv>", |
| 1090 | "content": "kv-lines", |
| 1091 | }, |
| 1092 | }, |
| 1093 | } |
| 1094 | 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>' |
| 1095 | # Drive byte-by-byte to maximise chunk count. |
| 1096 | streamer = ResponseParser(spec, prefix="") |
| 1097 | events: list[dict] = [] |
| 1098 | for ch in text: |
| 1099 | events.extend(streamer.feed(ch)) |
| 1100 | _, final_events = streamer.finalize() |
| 1101 | events.extend(final_events) |
| 1102 | |
| 1103 | per_field_chunks: dict[str, list[str]] = {} |
| 1104 | per_field_dirty: dict[str, set[bool]] = {} |
| 1105 | for ev in events: |
| 1106 | if ev["type"] == "region_chunk": |
| 1107 | per_field_chunks.setdefault(ev["field"], []).append(ev["text"]) |
| 1108 | per_field_dirty.setdefault(ev["field"], set()).add(ev["dirty"]) |
| 1109 | |
| 1110 | # Clean (streamable) regions. |
| 1111 | for field in ("thinking", "score"): |
| 1112 | self.assertEqual(per_field_dirty[field], {False}, f"{field} should be clean") |
| 1113 | # Dirty (structured) regions. |
| 1114 | for field in ("json_call", "xml_call", "kv_call"): |
| 1115 | self.assertEqual(per_field_dirty[field], {True}, f"{field} should be dirty") |
| 1116 | |
| 1117 | # Dirty chunks reconstruct the raw region body (un-parsed). Clean |
| 1118 | # chunks reconstruct the verbatim body too: stripping happens at close. |
| 1119 | self.assertEqual("".join(per_field_chunks["thinking"]), "hello world") |
| 1120 | self.assertEqual("".join(per_field_chunks["score"]), "42") |
| 1121 | self.assertEqual("".join(per_field_chunks["json_call"]), '{"a": 1, "b": 2}') |
| 1122 | self.assertEqual("".join(per_field_chunks["xml_call"]), "<name=foo><age=10>") |
| 1123 | self.assertEqual("".join(per_field_chunks["kv_call"]), "k1: v1\nk2: v2") |
| 1124 | |
| 1125 | def test_long_regex_open_pattern_streams_byte_by_byte(self): |
| 1126 | """Regression: a regex `open_pattern` whose full match spans well past the |
| 1127 | old fixed 64-byte hold window (gpt-oss tool calls) must not be dropped when |
| 1128 | the stream arrives in tiny chunks. Before the partial-match rewrite, the |
| 1129 | leading `<|channel|>` got flushed out of the hold window before `<|message|>` |
| 1130 | arrived, so the tool call vanished under small-chunk streaming.""" |
| 1131 | text = ( |
| 1132 | "<|channel|>analysis<|message|>Let me check.<|end|>" |
| 1133 | "<|start|>assistant<|channel|>commentary to=functions.get_current_weather " |
| 1134 | '<|constrain|>json<|message|>{"location": "San Francisco, CA"}<|call|>' |
| 1135 | ) |
| 1136 | expected = parse_response(text, gpt_oss_template, prefix="") |
| 1137 | # Sanity: the whole-string parse really does recover the tool call. |
| 1138 | self.assertEqual(len(expected["tool_calls"]), 1) |
| 1139 | self.assertEqual(expected["tool_calls"][0]["function"]["name"], "get_current_weather") |
| 1140 | streamer = ResponseParser(gpt_oss_template, prefix="") |
| 1141 | for ch in text: # one byte at a time -- the worst case for the old heuristic |
| 1142 | streamer.feed(ch) |
| 1143 | message, _ = streamer.finalize() |
| 1144 | self.assertEqual(message, expected) |
| 1145 | |
| 1146 | def test_feed_after_finalize_raises(self): |
| 1147 | streamer = ResponseParser(smollm_template, prefix="") |
| 1148 | streamer.feed("<think>x</think>") |
| 1149 | streamer.finalize() |
| 1150 | with self.assertRaises(RuntimeError): |
| 1151 | streamer.feed("more") |
| 1152 | with self.assertRaises(RuntimeError): |
| 1153 | streamer.finalize() |
| 1154 | |
| 1155 | def test_empty_input_streams_cleanly(self): |
| 1156 | streamer = ResponseParser(smollm_template, prefix="") |
| 1157 | self.assertEqual(streamer.feed(""), []) |
| 1158 | result, final_events = streamer.finalize() |
| 1159 | # Only the default fields should remain; nothing else is required. |
| 1160 | self.assertEqual(result, {"role": "assistant"}) |
| 1161 | self.assertEqual(final_events, []) |
| 1162 | |
| 1163 | def test_region_events_expose_delimiter_offsets(self): |
| 1164 | spec = { |
| 1165 | "start_anchor": "<assistant>", |
| 1166 | "fields": { |
| 1167 | "content": {"content_args": {"strip": False}}, |
| 1168 | "tag": {"open_pattern": r"<tag id=\d+>", "close": "</tag>"}, |
| 1169 | }, |
| 1170 | } |
| 1171 | parser = ResponseParser(spec, prefix="") |
| 1172 | events = parser.feed("plain<tag id=7>body</tag>tail") |
| 1173 | events += parser.finalize()[1] |
| 1174 | |
| 1175 | boundaries = [ |
| 1176 | (event["type"], event["field"], parser.input_text[event["start"] : event["end"]]) |
| 1177 | for event in events |
| 1178 | if event["type"] != "region_chunk" |
| 1179 | ] |
| 1180 | self.assertEqual( |
| 1181 | boundaries, |
| 1182 | [ |
| 1183 | ("region_open", "content", ""), |
| 1184 | ("region_close", "content", ""), |
| 1185 | ("region_open", "tag", "<tag id=7>"), |
| 1186 | ("region_close", "tag", "</tag>"), |
| 1187 | ("region_open", "content", ""), |
| 1188 | ("region_close", "content", ""), |
| 1189 | ], |
| 1190 | ) |
| 1191 | |
| 1192 | def test_open_captures_and_prefix_boundary(self): |
| 1193 | spec = { |
| 1194 | "start_anchor": "[BEGIN]", |
| 1195 | "fields": { |
| 1196 | "tool_calls": { |
| 1197 | "open_pattern": r"<call:(?P<name>\w+)>", |
| 1198 | "close": "</call>", |
| 1199 | "content": "json", |
| 1200 | "transform": {"name": "{name}", "arguments": "{content}"}, |
| 1201 | } |
| 1202 | }, |
| 1203 | } |
| 1204 | parser = ResponseParser(spec, prefix="history[BEGIN]<call:get_") |
| 1205 | prefix_end = len(parser.input_text) |
| 1206 | events = parser.feed('weather>{"city":"Paris"}</call>') |
| 1207 | opening = events[0] |
| 1208 | |
| 1209 | self.assertEqual(prefix_end, len("<call:get_")) |
| 1210 | self.assertEqual((opening["type"], opening["start"], opening["end"]), ("region_open", 0, len("<call:get_weather>"))) |
| 1211 | self.assertEqual(opening["captures"], {"name": "get_weather"}) |
| 1212 | |
| 1213 | def test_held_prefix_bytes_are_not_rechunked(self): |
| 1214 | spec = {"start_anchor": "[BEGIN]", "fields": {"content": {}, "tag": {"open": "<tag>", "close": "</tag>"}}} |
| 1215 | parser = ResponseParser(spec, prefix="[BEGIN]Existing <") |
| 1216 | chunks = [event["text"] for event in parser.feed("b> tail") if event["type"] == "region_chunk"] |
| 1217 | message, _ = parser.finalize() |
| 1218 | |
| 1219 | self.assertEqual(chunks, ["b> tail"]) |
| 1220 | self.assertEqual(message["content"], "Existing <b> tail") |
| 1221 | |
| 1222 | def test_malformed_region_events_recover_and_continue(self): |
| 1223 | spec = { |
| 1224 | "start_anchor": "[BEGIN]", |
| 1225 | "fields": { |
| 1226 | "tool_calls": { |
| 1227 | "open_pattern": r"<call:(?P<name>\w+)>", |
| 1228 | "close": "</call>", |
| 1229 | "content": "json", |
| 1230 | "repeats": True, |
| 1231 | "transform": { |
| 1232 | "type": "function", |
| 1233 | "function": {"name": "{name}", "arguments": "{content}"}, |
| 1234 | }, |
| 1235 | } |
| 1236 | }, |
| 1237 | } |
| 1238 | parser = ResponseParser(spec, prefix="") |
| 1239 | events = parser.feed('<call:bad>{"x":</call><call:good>{"x":1}</call>') |
| 1240 | message, _ = parser.finalize() |
| 1241 | opening, malformed = events[0], events[2] |
| 1242 | |
| 1243 | self.assertEqual(malformed["type"], "region_malformed") |
| 1244 | self.assertIsInstance(malformed["error"], ValueError) |
| 1245 | self.assertEqual(parser.input_text[opening["end"] : malformed["start"]], '{"x":') |
| 1246 | self.assertEqual(parser.input_text[malformed["start"] : malformed["end"]], "</call>") |
| 1247 | self.assertEqual( |
| 1248 | message["tool_calls"], |
| 1249 | [{"type": "function", "function": {"name": "good", "arguments": {"x": 1}}}], |
| 1250 | ) |
| 1251 | |
| 1252 | parser = ResponseParser(spec, prefix="") |
| 1253 | parser.feed('<call:bad>{"x":') |
| 1254 | _, events = parser.finalize() |
| 1255 | self.assertEqual(events[-1]["type"], "region_malformed") |
| 1256 | self.assertEqual(events[-1]["start"], len(parser.input_text)) |
| 1257 | |
| 1258 | def test_parse_response_raises_for_malformed_region(self): |
| 1259 | spec = { |
| 1260 | "start_anchor": "<|assistant|>", |
| 1261 | "fields": {"x": {"open": "<x>", "close": "</x>", "content": "json"}}, |
| 1262 | } |
| 1263 | with self.assertRaisesRegex(ValueError, "could not parse region as JSON"): |
| 1264 | parse_response('<x>{"name":', spec, prefix="") |
| 1265 | with self.assertRaisesRegex(ValueError, "could not parse region as JSON"): |
| 1266 | parse_response("", spec, prefix='<|assistant|><x>{"name":</x>') |
| 1267 | |
| 1268 | |
| 1269 | class PrefixAndTruncationTest(unittest.TestCase): |
| 1270 | def test_prefix_lands_inside_explicit_region(self): |
| 1271 | """A Qwen-style template emits `<|im_start|>assistant\\n<think>\\n` as the |
| 1272 | assistant prefix. The model continues from inside the thinking block.""" |
| 1273 | prompt = ( |
| 1274 | "<|im_start|>system\nYou are helpful<|im_end|>\n" |
| 1275 | "<|im_start|>user\nHi<|im_end|>\n" |
| 1276 | "<|im_start|>assistant\n<think>\n" |
| 1277 | ) |
| 1278 | generated = "Let me think...</think>" |
| 1279 | stream = ResponseParser(qwen3_template, prefix=prompt) |
| 1280 | # The region_open for `thinking` surfaces via initial_events; the |
| 1281 | # caller replays it before feeding model output. |
| 1282 | self.assertEqual( |
| 1283 | [(e["type"], e["field"]) for e in stream.initial_events], |
| 1284 | [("region_open", "thinking"), ("region_chunk", "thinking")], |
| 1285 | ) |
| 1286 | events = stream.feed(generated) |
| 1287 | result, _ = stream.finalize() |
| 1288 | # thinking ends up with the prefill + generated body; text parser strips, |
| 1289 | # so the leading "\n" from the prefix is trimmed in the final value. |
| 1290 | self.assertEqual(result, {"role": "assistant", "thinking": "Let me think..."}) |
| 1291 | # The feed stream only sees the rest of the body and the close; the |
| 1292 | # prefill already emitted region_open. |
| 1293 | self.assertEqual([e["type"] for e in events], ["region_chunk", "region_close"]) |
| 1294 | self.assertEqual(events[1]["field"], "thinking") |
| 1295 | |
| 1296 | def test_prefix_truncated_to_last_anchor(self): |
| 1297 | """Multiple `<|im_start|>assistant\\n` anchors in the prefix (multi-turn |
| 1298 | conversation): only the slice after the LAST anchor matters.""" |
| 1299 | prompt = ( |
| 1300 | "<|im_start|>system\nA<|im_end|>\n" |
| 1301 | "<|im_start|>user\nB<|im_end|>\n" |
| 1302 | "<|im_start|>assistant\nEarlier reply<|im_end|>\n" |
| 1303 | "<|im_start|>user\nFollowup<|im_end|>\n" |
| 1304 | "<|im_start|>assistant\n<think>\n" |
| 1305 | ) |
| 1306 | stream = ResponseParser(qwen3_template, prefix=prompt) |
| 1307 | # We landed inside `thinking` (from the LAST assistant turn's `<think>\n`), |
| 1308 | # not in some earlier-turn artifact. |
| 1309 | opens = [e for e in stream.initial_events if e["type"] == "region_open"] |
| 1310 | self.assertEqual([e["field"] for e in opens], ["thinking"]) |
| 1311 | # No earlier-turn content leaked into output. |
| 1312 | stream.feed("done</think>") |
| 1313 | stream.finalize() |
| 1314 | self.assertEqual(stream._output, {"role": "assistant", "thinking": "done"}) |
| 1315 | |
| 1316 | def test_template_without_anchor_rejected_at_load(self): |
| 1317 | """A template missing both `start_anchor` and `start_anchor_pattern` is |
| 1318 | rejected at load time. Without an anchor, a multi-turn prompt would be |
| 1319 | fed through the parser in full and earlier turns would pollute the |
| 1320 | current message's state.""" |
| 1321 | anchorless = {k: v for k, v in qwen3_template.items() if k != "start_anchor"} |
| 1322 | with self.assertRaises(ValueError) as cm: |
| 1323 | ResponseParser(anchorless) |
| 1324 | msg = str(cm.exception) |
| 1325 | self.assertIn("start_anchor", msg) |
| 1326 | |
| 1327 | def test_prefix_anchor_not_found_falls_back(self): |
| 1328 | """Spec has start_anchor but the prefix doesn't contain it: parser |
| 1329 | falls back to processing the entire prefix (with a logged warning).""" |
| 1330 | prompt = "<think>\n" # no <|im_start|>assistant\n |
| 1331 | stream = ResponseParser(qwen3_template, prefix=prompt) |
| 1332 | opens = [e for e in stream.initial_events if e["type"] == "region_open"] |
| 1333 | self.assertEqual([e["field"] for e in opens], ["thinking"]) |
| 1334 | stream.feed("hi</think>") |
| 1335 | stream.finalize() |
| 1336 | self.assertEqual(stream._output, {"role": "assistant", "thinking": "hi"}) |
| 1337 | |
| 1338 | def test_round_trip_equivalence_prefix_streaming_vs_oneshot(self): |
| 1339 | """The load-bearing property: `prefix=p` + chunked `feed(g)` produces the |
| 1340 | same dict as the one-shot `parse_response(g, prefix=p)`, regardless of how |
| 1341 | `g` is chunked. (Concatenating the prompt into the response is deliberately |
| 1342 | NOT equivalent -- the anchor is applied to the prefix only, never to the |
| 1343 | generation; see test_history_bleed_is_guarded_by_prefix_not_by_response_anchor.)""" |
| 1344 | prompt = "<|im_start|>system\nA<|im_end|>\n<|im_start|>user\nB<|im_end|>\n<|im_start|>assistant\n<think>\n" |
| 1345 | for name, tmpl_dict, gen_text in _STREAMING_FIXTURES: |
| 1346 | if "thinking" not in gen_text and "<think>" not in gen_text: |
| 1347 | continue |
| 1348 | # Only fixtures whose generation text is compatible with starting |
| 1349 | # inside `thinking`. Restrict to qwen3 / smollm shape for clarity. |
| 1350 | if name not in ("qwen3", "smollm"): |
| 1351 | continue |
| 1352 | tmpl_with_anchor = {**tmpl_dict, "start_anchor": "<|im_start|>assistant\n"} |
| 1353 | via_prefix = parse_response(gen_text, tmpl_with_anchor, prefix=prompt) |
| 1354 | # Streaming forms must match the one-shot prefix form for every chunking. |
| 1355 | for step in (1, 3, 7, 31): |
| 1356 | with self.subTest(fixture=name, step=step): |
| 1357 | stream = ResponseParser(tmpl_with_anchor, prefix=prompt) |
| 1358 | for chunk in _chunk_fixed(gen_text, step): |
| 1359 | stream.feed(chunk) |
| 1360 | message, _ = stream.finalize() |
| 1361 | self.assertEqual(message, via_prefix) |
| 1362 | |
| 1363 | def test_history_bleed_is_guarded_by_prefix_not_by_response_anchor(self): |
| 1364 | """The `start_anchor` guards against history bleed only via `prefix=`; it is NOT |
| 1365 | applied to the response. The response is the generation and may legitimately contain |
| 1366 | the anchor (e.g. gpt-oss harmony opens every channel with `<|start|>assistant`), so |
| 1367 | truncating it would drop real content. Passing the prompt as `prefix=` is the way to |
| 1368 | keep earlier turns out of the parse.""" |
| 1369 | spec = { |
| 1370 | "defaults": {"role": "assistant"}, |
| 1371 | "start_anchor": "<|im_start|>assistant\n", |
| 1372 | "fields": {"content": {"close_pattern": r"\Z", "content": "text"}}, |
| 1373 | } |
| 1374 | prompt = "<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n" |
| 1375 | gen = "Hello there!" |
| 1376 | clean = {"role": "assistant", "content": "Hello there!"} |
| 1377 | # The supported guard: pass the prompt as prefix= so history is truncated off the prefix. |
| 1378 | self.assertEqual(parse_response(gen, spec, prefix=prompt), clean) |
| 1379 | # Pure generation parses cleanly with the explicit no-prefix opt-out (prefix=""). |
| 1380 | self.assertEqual(parse_response(gen, spec, prefix=""), clean) |
| 1381 | # An anchor inside the response is treated as content, never as a history boundary: |
| 1382 | # gpt-oss re-emits `<|start|>assistant` between channels, and that content survives. |
| 1383 | gpt_oss_gen = ( |
| 1384 | "<|channel|>analysis<|message|>thinking<|end|><|start|>assistant<|channel|>final<|message|>answer" |
| 1385 | ) |
| 1386 | self.assertEqual( |
| 1387 | parse_response(gpt_oss_gen, gpt_oss_template, prefix=""), |
| 1388 | {"role": "assistant", "thinking": "thinking", "content": "answer"}, |
| 1389 | ) |
| 1390 | |
| 1391 | def test_prefix_with_open_close_inside_truncated_region(self): |
| 1392 | """Prefix opens AND closes a region. The full open/chunk/close event |
| 1393 | sequence is surfaced via initial_events, and the closed region lands |
| 1394 | in the output dict: so renderers can show prefill content.""" |
| 1395 | spec = { |
| 1396 | "defaults": {"role": "assistant"}, |
| 1397 | "start_anchor": "[BEGIN]", |
| 1398 | "fields": { |
| 1399 | "tag": {"open": "<tag>", "close": "</tag>", "content": "text"}, |
| 1400 | "body": {"close_pattern": r"$", "content": "text"}, # implicit |
| 1401 | }, |
| 1402 | } |
| 1403 | prefix = "noise[BEGIN]<tag>silently consumed</tag>" |
| 1404 | stream = ResponseParser(spec, prefix=prefix) |
| 1405 | types = [e["type"] for e in stream.initial_events] |
| 1406 | self.assertEqual(types, ["region_open", "region_chunk", "region_close"]) |
| 1407 | self.assertTrue(all(e["field"] == "tag" for e in stream.initial_events)) |
| 1408 | self.assertEqual(stream.initial_events[-1]["value"], "silently consumed") |
| 1409 | stream.feed("real generated body") |
| 1410 | result, _ = stream.finalize() |
| 1411 | self.assertEqual(result["tag"], "silently consumed") |
| 1412 | self.assertEqual(result["body"], "real generated body") |
| 1413 | |
| 1414 | def test_prefix_lands_inside_implicit_region(self): |
| 1415 | """Prefix wrote plaintext into the implicit region (e.g. assistant |
| 1416 | prefill before the model continues). The region_open for the implicit |
| 1417 | region must surface via initial_events so consumers don't miss it: |
| 1418 | `_opened` will already be True by the time feed runs.""" |
| 1419 | prompt = "<|im_start|>assistant\nSure, here is " |
| 1420 | stream = ResponseParser(smollm_template, prefix=prompt) |
| 1421 | opens = [e for e in stream.initial_events if e["type"] == "region_open"] |
| 1422 | self.assertEqual([e["field"] for e in opens], ["content"]) |
| 1423 | events = stream.feed("the answer<|im_end|>") |
| 1424 | # No second region_open from feed: the implicit region was already |
| 1425 | # opened during prefill and surfaced via initial_events. |
| 1426 | self.assertNotIn("region_open", [e["type"] for e in events]) |
| 1427 | |
| 1428 | def test_prefix_partial_pattern_at_boundary(self): |
| 1429 | """Post-truncation prefix ends mid-delimiter. The first feed completes |
| 1430 | the match; initial_events is empty (no region opened within the |
| 1431 | prefix yet) and the open fires from `feed()`.""" |
| 1432 | prefix = "<|im_start|>assistant\n<th" "i" # incomplete `<think>` |
| 1433 | stream = ResponseParser(qwen3_template, prefix=prefix) |
| 1434 | self.assertEqual(stream.initial_events, []) |
| 1435 | events = stream.feed("nk>real body</think>") |
| 1436 | types = [e["type"] for e in events] |
| 1437 | self.assertIn("region_open", types) # think opens during feed, not prefill |
| 1438 | stream.finalize() |
| 1439 | self.assertEqual(stream._output, {"role": "assistant", "thinking": "real body"}) |
| 1440 | |
| 1441 | def test_prefix_is_required(self): |
| 1442 | spec = { |
| 1443 | "start_anchor": "<assistant>", |
| 1444 | "fields": {"content": {"content": "text"}}, |
| 1445 | } |
| 1446 | with self.assertRaisesRegex(ValueError, "requires `prefix`"): |
| 1447 | ResponseParser(spec) |
| 1448 | |
| 1449 | |
| 1450 | # xml-inline without a value_parser: parameter bodies stay raw strings until tools= coerces them. |
| 1451 | _XML_STRING_ARGS_TEMPLATE = { |
| 1452 | "defaults": {"role": "assistant"}, |
| 1453 | "start_anchor": "<|im_start|>assistant\n", |
| 1454 | "fields": { |
| 1455 | "tool_calls": { |
| 1456 | "open_pattern": r"<tool_call>\s*<function=(?P<name>\w+)>", |
| 1457 | "close": "</tool_call>", |
| 1458 | "repeats": True, |
| 1459 | "content": "xml-inline", |
| 1460 | "content_args": { |
| 1461 | "tag_pattern": r"<parameter=(?P<key>\w+)>\s*(?P<value>.*?)\s*</parameter>", |
| 1462 | }, |
| 1463 | "transform": {"type": "function", "function": {"name": "{name}", "arguments": "{content}"}}, |
| 1464 | }, |
| 1465 | }, |
| 1466 | } |
| 1467 | |
| 1468 | # kv-lines without a value_parser: values likewise stay raw strings for tools= to cast. |
| 1469 | _KV_LINES_TOOLS_TEMPLATE = { |
| 1470 | "defaults": {"role": "assistant"}, |
| 1471 | "start_anchor": "<|im_start|>assistant\n", |
| 1472 | "fields": { |
| 1473 | "tool_calls": { |
| 1474 | "open_pattern": r"<tool_call>\s*<function=(?P<name>\w+)>\n", |
| 1475 | "close": "</tool_call>", |
| 1476 | "repeats": True, |
| 1477 | "content": "kv-lines", |
| 1478 | "transform": {"type": "function", "function": {"name": "{name}", "arguments": "{content}"}}, |
| 1479 | }, |
| 1480 | }, |
| 1481 | } |
| 1482 | |
| 1483 | _SET_ALARM_CALL = ( |
| 1484 | "<tool_call>\n<function=set_alarm>\n" |
| 1485 | "<parameter=hour>\n7\n</parameter>\n" |
| 1486 | "<parameter=enabled>\ntrue\n</parameter>\n" |
| 1487 | "<parameter=label>\nwake up\n</parameter>\n" |
| 1488 | "</function>\n</tool_call>" |
| 1489 | ) |
| 1490 | |
| 1491 | |
| 1492 | def _set_alarm_tools(**properties): |
| 1493 | return [ |
| 1494 | { |
| 1495 | "type": "function", |
| 1496 | "function": { |
| 1497 | "name": "set_alarm", |
| 1498 | "parameters": {"type": "object", "properties": properties}, |
| 1499 | }, |
| 1500 | } |
| 1501 | ] |
| 1502 | |
| 1503 | |
| 1504 | _SET_ALARM_TOOLS = _set_alarm_tools( |
| 1505 | hour={"type": "integer"}, |
| 1506 | enabled={"type": "boolean"}, |
| 1507 | label={"type": "string"}, |
| 1508 | ) |
| 1509 | |
| 1510 | |
| 1511 | def _first_tool_args(message): |
| 1512 | return message["tool_calls"][0]["function"]["arguments"] |
| 1513 | |
| 1514 | |
| 1515 | def _parser_with_tools(tools): |
| 1516 | return ResponseParser(_XML_STRING_ARGS_TEMPLATE, prefix="", tools=tools) |
| 1517 | |
| 1518 | |
| 1519 | class ToolArgCoercionTest(unittest.TestCase): |
| 1520 | def test_coerce_tool_calls_casts_declared_types(self): |
| 1521 | tools = _set_alarm_tools( |
| 1522 | count={"type": "integer"}, |
| 1523 | ratio={"type": "number"}, |
| 1524 | enabled={"type": "boolean"}, |
| 1525 | tags={"type": "array"}, |
| 1526 | note={"type": "string"}, |
| 1527 | ) |
| 1528 | arguments = { |
| 1529 | "count": "3", |
| 1530 | "ratio": "1.5", |
| 1531 | "enabled": "true", |
| 1532 | "tags": '["a", "b"]', |
| 1533 | "note": "hello", |
| 1534 | "already_typed": 7, |
| 1535 | "extra": "unscheduled", |
| 1536 | } |
| 1537 | call: dict[str, Any] = { |
| 1538 | "type": "function", |
| 1539 | "function": {"name": "set_alarm", "arguments": arguments}, |
| 1540 | } |
| 1541 | self.assertIs(_parser_with_tools(tools)._coerce_tool_calls(call), call) |
| 1542 | self.assertEqual( |
| 1543 | call["function"]["arguments"], |
| 1544 | { |
| 1545 | "count": 3, |
| 1546 | "ratio": 1.5, |
| 1547 | "enabled": True, |
| 1548 | "tags": ["a", "b"], |
| 1549 | "note": "hello", |
| 1550 | "already_typed": 7, |
| 1551 | "extra": "unscheduled", |
| 1552 | }, |
| 1553 | ) |
| 1554 | |
| 1555 | def test_coerce_falls_back_to_raw_on_failure(self): |
| 1556 | self.assertEqual(_coerce("not-a-number", ("integer",)), "not-a-number") |
| 1557 | |
| 1558 | def test_coerce_handles_any_of_and_null(self): |
| 1559 | self.assertEqual(_coerce("5", ("integer", "null")), 5) |
| 1560 | self.assertIsNone(_coerce("null", ("integer", "null"))) |
| 1561 | |
| 1562 | def test_coerce_booleans_match_bool_parser(self): |
| 1563 | # Accept the same literals as the `bool` content parser, case-insensitively. |
| 1564 | for raw, expected in [("true", True), ("True", True), ("1", True), ("false", False), ("0", False)]: |
| 1565 | self.assertEqual(_coerce(raw, ("boolean",)), expected) |
| 1566 | # Non-boolean text stays a string rather than silently becoming False. |
| 1567 | self.assertEqual(_coerce("maybe", ("boolean",)), "maybe") |
| 1568 | |
| 1569 | def test_coerce_object_array_require_matching_json(self): |
| 1570 | # A JSON object body is only accepted for an `object` param, a JSON array only for `array`. |
| 1571 | self.assertEqual(_coerce("[1, 2]", ("object",)), "[1, 2]") |
| 1572 | self.assertEqual(_coerce('{"a": 1}', ("array",)), '{"a": 1}') |
| 1573 | self.assertEqual(_coerce('{"a": 1}', ("object",)), {"a": 1}) |
| 1574 | # NaN / inf are not valid JSON numbers, so a `number` param keeps the raw text. |
| 1575 | self.assertEqual(_coerce("NaN", ("number",)), "NaN") |
| 1576 | |
| 1577 | def test_coerce_tool_calls_handles_single_and_list(self): |
| 1578 | parser = _parser_with_tools(_SET_ALARM_TOOLS) |
| 1579 | call: dict[str, Any] = { |
| 1580 | "type": "function", |
| 1581 | "function": {"name": "set_alarm", "arguments": {"hour": "7"}}, |
| 1582 | } |
| 1583 | self.assertIs(parser._coerce_tool_calls(call), call) |
| 1584 | self.assertEqual(call["function"]["arguments"], {"hour": 7}) |
| 1585 | # A list of calls (as produced by `transform_each`) is coerced element-wise. |
| 1586 | calls = [{"type": "function", "function": {"name": "set_alarm", "arguments": {"hour": "9"}}}] |
| 1587 | self.assertEqual(parser._coerce_tool_calls(calls)[0]["function"]["arguments"], {"hour": 9}) |
| 1588 | # Non-tool-call values pass through untouched. |
| 1589 | self.assertEqual(parser._coerce_tool_calls("hello"), "hello") |
| 1590 | |
| 1591 | def test_schema_types_handles_get_json_schema_dialect(self): |
| 1592 | self.assertEqual(_schema_types({"type": "integer"}), ("integer",)) |
| 1593 | self.assertEqual(_schema_types({"type": ["integer", "string"]}), ("integer", "string")) |
| 1594 | self.assertEqual(_schema_types({"anyOf": [{"type": "boolean"}, {"type": "string"}]}), ("boolean", "string")) |
| 1595 | self.assertEqual(_schema_types({"oneOf": [{"type": "number"}, {"type": "null"}]}), ("number", "null")) |
| 1596 | self.assertEqual(_schema_types({"type": "integer", "nullable": True}), ("integer", "null")) |
| 1597 | # Undescribed parameters resolve to no candidate types, making coercion a no-op. |
| 1598 | self.assertEqual(_schema_types({"description": "no type"}), ()) |
| 1599 | |
| 1600 | def test_parse_response_tools_coerces_xml_inline_string_args(self): |
| 1601 | # Without a value_parser, xml-inline argument bodies stay strings; tools= casts them. |
| 1602 | without = parse_response(_SET_ALARM_CALL, _XML_STRING_ARGS_TEMPLATE, prefix="") |
| 1603 | with_tools = parse_response(_SET_ALARM_CALL, _XML_STRING_ARGS_TEMPLATE, prefix="", tools=_SET_ALARM_TOOLS) |
| 1604 | self.assertEqual(_first_tool_args(without), {"hour": "7", "enabled": "true", "label": "wake up"}) |
| 1605 | self.assertEqual(_first_tool_args(with_tools), {"hour": 7, "enabled": True, "label": "wake up"}) |
| 1606 | |
| 1607 | def test_parse_response_tools_coerces_one_of_string_args(self): |
| 1608 | tools = _set_alarm_tools(hour={"oneOf": [{"type": "integer"}, {"type": "null"}]}) |
| 1609 | parsed = parse_response(_SET_ALARM_CALL, _XML_STRING_ARGS_TEMPLATE, prefix="", tools=tools) |
| 1610 | self.assertEqual(_first_tool_args(parsed)["hour"], 7) |
| 1611 | |
| 1612 | def test_streaming_tools_coerces_on_region_close(self): |
| 1613 | # Coercion must land on the region_close event during feed(), not only after finalize(). |
| 1614 | stream = ResponseParser(_XML_STRING_ARGS_TEMPLATE, prefix="", tools=_SET_ALARM_TOOLS) |
| 1615 | closes = [ |
| 1616 | event["value"] |
| 1617 | for chunk in _chunk_fixed(_SET_ALARM_CALL, 8) |
| 1618 | for event in stream.feed(chunk) |
| 1619 | if event["type"] == "region_close" and event["field"] == "tool_calls" |
| 1620 | ] |
| 1621 | self.assertEqual(len(closes), 1) |
| 1622 | self.assertEqual(closes[0]["function"]["arguments"], {"hour": 7, "enabled": True, "label": "wake up"}) |
| 1623 | |
| 1624 | def test_qwen3_tools_coerces_strings_left_by_value_parser(self): |
| 1625 | # qwen3's json+allow_non_json value_parser types what it can (`true`) and leaves |
| 1626 | # invalid JSON (`007`) as a string; tools= then casts those leftover strings. |
| 1627 | model_out = ( |
| 1628 | "<tool_call>\n<function=set_alarm>\n" |
| 1629 | "<parameter=hour>\n007\n</parameter>\n" |
| 1630 | "<parameter=enabled>\ntrue\n</parameter>\n" |
| 1631 | "<parameter=label>\nwake up\n</parameter>\n" |
| 1632 | "</function>\n</tool_call>" |
| 1633 | ) |
| 1634 | without = parse_response(model_out, qwen3_template, prefix="") |
| 1635 | with_tools = parse_response(model_out, qwen3_template, prefix="", tools=_SET_ALARM_TOOLS) |
| 1636 | self.assertEqual(_first_tool_args(without), {"hour": "007", "enabled": True, "label": "wake up"}) |
| 1637 | self.assertEqual(_first_tool_args(with_tools), {"hour": 7, "enabled": True, "label": "wake up"}) |
| 1638 | |
| 1639 | def test_coercion_never_reworks_already_typed_values(self): |
| 1640 | # Coercion only casts strings: values the value_parser already typed are final |
| 1641 | model_out = "<tool_call>\n<function=set_alarm>\n<parameter=label>\n1.50\n</parameter>\n</tool_call>" |
| 1642 | # qwen3's lax value_parser has already read 1.50 as the float 1.5, so the string-typed label stays a float ... |
| 1643 | self.assertEqual( |
| 1644 | _first_tool_args(parse_response(model_out, qwen3_template, prefix="", tools=_SET_ALARM_TOOLS)), |
| 1645 | {"label": 1.5}, |
| 1646 | ) |
| 1647 | # ... while without a value_parser the raw text reaches the schema cast intact |
| 1648 | self.assertEqual( |
| 1649 | _first_tool_args(parse_response(model_out, _XML_STRING_ARGS_TEMPLATE, prefix="", tools=_SET_ALARM_TOOLS)), |
| 1650 | {"label": "1.50"}, |
| 1651 | ) |
| 1652 | |
| 1653 | def test_kv_lines_string_args_are_coerced(self): |
| 1654 | model_out = "<tool_call>\n<function=set_alarm>\nhour: 7\nenabled: true\n</tool_call>" |
| 1655 | without = parse_response(model_out, _KV_LINES_TOOLS_TEMPLATE, prefix="") |
| 1656 | with_tools = parse_response(model_out, _KV_LINES_TOOLS_TEMPLATE, prefix="", tools=_SET_ALARM_TOOLS) |
| 1657 | self.assertEqual(_first_tool_args(without), {"hour": "7", "enabled": "true"}) |
| 1658 | self.assertEqual(_first_tool_args(with_tools), {"hour": 7, "enabled": True}) |
| 1659 | |
| 1660 | def test_non_tool_call_regions_are_untouched(self): |
| 1661 | # A field that captures a `name` but does not parse into a tool call must be left |
| 1662 | # alone, even when the capture happens to match a tool: its keys are not arguments. |
| 1663 | template = { |
| 1664 | "defaults": {"role": "assistant"}, |
| 1665 | "start_anchor": "<|im_start|>assistant\n", |
| 1666 | "fields": { |
| 1667 | "citation": { |
| 1668 | "open_pattern": r"<cite source=(?P<name>\w+)>", |
| 1669 | "close": "</cite>", |
| 1670 | "content": "xml-inline", |
| 1671 | "content_args": { |
| 1672 | "tag_pattern": r"<(?P<key>\w+)>\s*(?P<value>.*?)\s*</\1>", |
| 1673 | "value_parser": {"name": "json", "args": {"allow_non_json": True}}, |
| 1674 | }, |
| 1675 | "transform": {"source": "{name}", "fields": "{content}"}, |
| 1676 | }, |
| 1677 | }, |
| 1678 | } |
| 1679 | model_out = "<cite source=set_alarm><label>1.50</label><hour>7</hour></cite>" |
| 1680 | expected = {"source": "set_alarm", "fields": {"label": 1.5, "hour": 7}} |
| 1681 | self.assertEqual(parse_response(model_out, template, prefix="")["citation"], expected) |
| 1682 | self.assertEqual( |
| 1683 | parse_response(model_out, template, prefix="", tools=_SET_ALARM_TOOLS)["citation"], |
| 1684 | expected, |
| 1685 | ) |
| 1686 | |
| 1687 | def test_merge_duplicates_arguments_are_cast_element_wise(self): |
| 1688 | # merge_duplicates collects repeated tags into a list, which is cast element-wise |
| 1689 | template: dict[str, Any] = copy.deepcopy(_XML_STRING_ARGS_TEMPLATE) |
| 1690 | template["fields"]["tool_calls"]["content_args"]["merge_duplicates"] = True |
| 1691 | model_out = ( |
| 1692 | "<tool_call>\n<function=set_alarm>\n" |
| 1693 | "<parameter=hour>\n7\n</parameter>\n" |
| 1694 | "<parameter=hour>\n9\n</parameter>\n" |
| 1695 | "</function>\n</tool_call>" |
| 1696 | ) |
| 1697 | self.assertEqual(_first_tool_args(parse_response(model_out, template, prefix="")), {"hour": ["7", "9"]}) |
| 1698 | self.assertEqual( |
| 1699 | _first_tool_args(parse_response(model_out, template, prefix="", tools=_SET_ALARM_TOOLS)), |
| 1700 | {"hour": [7, 9]}, |
| 1701 | ) |
| 1702 | # Elements that don't cast, and non-string elements, are left as they are. |
| 1703 | parser = _parser_with_tools(_SET_ALARM_TOOLS) |
| 1704 | call: dict[str, Any] = { |
| 1705 | "type": "function", |
| 1706 | "function": { |
| 1707 | "name": "set_alarm", |
| 1708 | "arguments": {"hour": ["7", "x", 9]}, |
| 1709 | }, |
| 1710 | } |
| 1711 | parser._coerce_tool_calls(call) |
| 1712 | self.assertEqual(call["function"]["arguments"], {"hour": [7, "x", 9]}) |
| 1713 | |
| 1714 | def test_already_decoded_array_is_not_cast_element_wise(self): |
| 1715 | tools = _set_alarm_tools(groups={"type": "array", "items": {"type": "string"}}) |
| 1716 | call: dict[str, Any] = { |
| 1717 | "type": "function", |
| 1718 | "function": { |
| 1719 | "name": "set_alarm", |
| 1720 | "arguments": {"groups": ["[1,2]", "[]"]}, |
| 1721 | }, |
| 1722 | } |
| 1723 | _parser_with_tools(tools)._coerce_tool_calls(call) |
| 1724 | self.assertEqual(call["function"]["arguments"], {"groups": ["[1,2]", "[]"]}) |
| 1725 | |
| 1726 | def test_coerce_tool_calls_ignores_unusable_function_name(self): |
| 1727 | # A transform can hand us a name parsed from model output, so a non-string name |
| 1728 | # must be ignored rather than raising on the schema lookup. |
| 1729 | parser = _parser_with_tools(_SET_ALARM_TOOLS) |
| 1730 | call: dict[str, Any] = { |
| 1731 | "type": "function", |
| 1732 | "function": {"name": ["set_alarm"], "arguments": {"hour": "7"}}, |
| 1733 | } |
| 1734 | self.assertEqual(parser._coerce_tool_calls(call), call) |
| 1735 | self.assertEqual(call["function"]["arguments"], {"hour": "7"}) |
| 1736 | |
| 1737 | def test_union_with_container_still_keeps_scalar_text(self): |
| 1738 | model_out = "<tool_call>\n<function=set_alarm>\n<parameter=label>\n1.50\n</parameter>\n</tool_call>" |
| 1739 | union = _set_alarm_tools(label={"anyOf": [{"type": "string"}, {"type": "object"}]}) |
| 1740 | # `string` never casts, so the union keeps scalar-looking text as text ... |
| 1741 | self.assertEqual( |
| 1742 | _first_tool_args(parse_response(model_out, _XML_STRING_ARGS_TEMPLATE, prefix="", tools=union)), |
| 1743 | {"label": "1.50"}, |
| 1744 | ) |
| 1745 | # ... while the `object` branch still decodes a body that really is an object. |
| 1746 | object_body = '<tool_call>\n<function=set_alarm>\n<parameter=label>\n{"a": 1}\n</parameter>\n</tool_call>' |
| 1747 | self.assertEqual( |
| 1748 | _first_tool_args(parse_response(object_body, _XML_STRING_ARGS_TEMPLATE, prefix="", tools=union)), |
| 1749 | {"label": {"a": 1}}, |
| 1750 | ) |
| 1751 | |
| 1752 | def test_tools_is_a_noop_for_typed_json_tool_calls(self): |
| 1753 | # JSON tool-call bodies already carry their types, so passing tools= must not change them. |
| 1754 | model_out = ( |
| 1755 | "<|START_THINKING|>x<|END_THINKING|>" |
| 1756 | '<|START_ACTION|>[{"tool_call_id": "0", "tool_name": "set_alarm", ' |
| 1757 | '"parameters": {"hour": 7, "enabled": true}}]<|END_ACTION|><|END_OF_TURN_TOKEN|>' |
| 1758 | ) |
| 1759 | without = parse_response(model_out, cohere_template, prefix="") |
| 1760 | with_tools = parse_response(model_out, cohere_template, prefix="", tools=_SET_ALARM_TOOLS) |
| 1761 | self.assertEqual(without, with_tools) |
| 1762 | self.assertEqual(_first_tool_args(with_tools), {"hour": 7, "enabled": True}) |
| 1763 | |
| 1764 | |
| 1765 | if __name__ == "__main__": |
| 1766 | unittest.main() |
| 1767 | # fmt: on |