Yoni GozlanandBugen Zhao
69d00bc264
Signed-off-by: Yoni Gozlan <74535834+yonigozlan@users.noreply.github.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Bugen Zhao <i@bugenzhao.com>
1776 lines
80 KiB
Python
1776 lines
80 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
|
# ruff: noqa
|
|
# fmt: off
|
|
#
|
|
# 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 unittest
|
|
from typing import Any
|
|
|
|
from vllm.parser.chat_parsing import ResponseParser, parse_response
|
|
from vllm.parser.chat_parsing.response_parser import _coerce, _schema_types
|
|
from vllm.parser.chat_parsing.response_templates import load_response_template
|
|
|
|
|
|
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_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"},
|
|
)
|
|
|
|
def test_mutable_defaults_do_not_leak_across_parsers(self):
|
|
"""A `repeats` field appends into its default list; one parse must not leak into the next
|
|
parse sharing the template, nor into the caller's spec dict."""
|
|
spec: dict[str, Any] = {
|
|
"defaults": {"role": "assistant", "tool_calls": []},
|
|
"start_anchor": "<|assistant|>",
|
|
"fields": {
|
|
"tool_calls": {
|
|
"open": "<t>",
|
|
"close": "</t>",
|
|
"repeats": True,
|
|
"content": "json",
|
|
},
|
|
"content": {"content": "text"},
|
|
},
|
|
}
|
|
template = load_response_template(spec)
|
|
self.assertEqual(
|
|
parse_response('<t>{"a": 1}</t>', template, prefix="")["tool_calls"],
|
|
[{"a": 1}],
|
|
)
|
|
self.assertEqual(
|
|
parse_response("hello", template, prefix=""),
|
|
{"role": "assistant", "tool_calls": [], "content": "hello"},
|
|
)
|
|
self.assertEqual(spec["defaults"]["tool_calls"], [])
|
|
|
|
|
|
# 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[ev["field"]] += ev["text"]
|
|
elif t == "region_close":
|
|
self.assertEqual(open_field, ev["field"], f"close for non-open region: {ev}")
|
|
close_values[ev["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, [])
|
|
|
|
def test_region_events_expose_delimiter_offsets(self):
|
|
spec = {
|
|
"start_anchor": "<assistant>",
|
|
"fields": {
|
|
"content": {"content_args": {"strip": False}},
|
|
"tag": {"open_pattern": r"<tag id=\d+>", "close": "</tag>"},
|
|
},
|
|
}
|
|
parser = ResponseParser(spec, prefix="")
|
|
events = parser.feed("plain<tag id=7>body</tag>tail")
|
|
events += parser.finalize()[1]
|
|
|
|
boundaries = [
|
|
(event["type"], event["field"], parser.input_text[event["start"] : event["end"]])
|
|
for event in events
|
|
if event["type"] != "region_chunk"
|
|
]
|
|
self.assertEqual(
|
|
boundaries,
|
|
[
|
|
("region_open", "content", ""),
|
|
("region_close", "content", ""),
|
|
("region_open", "tag", "<tag id=7>"),
|
|
("region_close", "tag", "</tag>"),
|
|
("region_open", "content", ""),
|
|
("region_close", "content", ""),
|
|
],
|
|
)
|
|
|
|
def test_open_captures_and_prefix_boundary(self):
|
|
spec = {
|
|
"start_anchor": "[BEGIN]",
|
|
"fields": {
|
|
"tool_calls": {
|
|
"open_pattern": r"<call:(?P<name>\w+)>",
|
|
"close": "</call>",
|
|
"content": "json",
|
|
"transform": {"name": "{name}", "arguments": "{content}"},
|
|
}
|
|
},
|
|
}
|
|
parser = ResponseParser(spec, prefix="history[BEGIN]<call:get_")
|
|
prefix_end = len(parser.input_text)
|
|
events = parser.feed('weather>{"city":"Paris"}</call>')
|
|
opening = events[0]
|
|
|
|
self.assertEqual(prefix_end, len("<call:get_"))
|
|
self.assertEqual((opening["type"], opening["start"], opening["end"]), ("region_open", 0, len("<call:get_weather>")))
|
|
self.assertEqual(opening["captures"], {"name": "get_weather"})
|
|
|
|
def test_held_prefix_bytes_are_not_rechunked(self):
|
|
spec = {"start_anchor": "[BEGIN]", "fields": {"content": {}, "tag": {"open": "<tag>", "close": "</tag>"}}}
|
|
parser = ResponseParser(spec, prefix="[BEGIN]Existing <")
|
|
chunks = [event["text"] for event in parser.feed("b> tail") if event["type"] == "region_chunk"]
|
|
message, _ = parser.finalize()
|
|
|
|
self.assertEqual(chunks, ["b> tail"])
|
|
self.assertEqual(message["content"], "Existing <b> tail")
|
|
|
|
def test_malformed_region_events_recover_and_continue(self):
|
|
spec = {
|
|
"start_anchor": "[BEGIN]",
|
|
"fields": {
|
|
"tool_calls": {
|
|
"open_pattern": r"<call:(?P<name>\w+)>",
|
|
"close": "</call>",
|
|
"content": "json",
|
|
"repeats": True,
|
|
"transform": {
|
|
"type": "function",
|
|
"function": {"name": "{name}", "arguments": "{content}"},
|
|
},
|
|
}
|
|
},
|
|
}
|
|
parser = ResponseParser(spec, prefix="")
|
|
events = parser.feed('<call:bad>{"x":</call><call:good>{"x":1}</call>')
|
|
message, _ = parser.finalize()
|
|
opening, malformed = events[0], events[2]
|
|
|
|
self.assertEqual(malformed["type"], "region_malformed")
|
|
self.assertIsInstance(malformed["error"], ValueError)
|
|
self.assertEqual(parser.input_text[opening["end"] : malformed["start"]], '{"x":')
|
|
self.assertEqual(parser.input_text[malformed["start"] : malformed["end"]], "</call>")
|
|
self.assertEqual(
|
|
message["tool_calls"],
|
|
[{"type": "function", "function": {"name": "good", "arguments": {"x": 1}}}],
|
|
)
|
|
|
|
parser = ResponseParser(spec, prefix="")
|
|
parser.feed('<call:bad>{"x":')
|
|
_, events = parser.finalize()
|
|
self.assertEqual(events[-1]["type"], "region_malformed")
|
|
self.assertEqual(events[-1]["start"], len(parser.input_text))
|
|
|
|
def test_parse_response_raises_for_malformed_region(self):
|
|
spec = {
|
|
"start_anchor": "<|assistant|>",
|
|
"fields": {"x": {"open": "<x>", "close": "</x>", "content": "json"}},
|
|
}
|
|
with self.assertRaisesRegex(ValueError, "could not parse region as JSON"):
|
|
parse_response('<x>{"name":', spec, prefix="")
|
|
with self.assertRaisesRegex(ValueError, "could not parse region as JSON"):
|
|
parse_response("", spec, prefix='<|assistant|><x>{"name":</x>')
|
|
|
|
|
|
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<th" "i" # 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_is_required(self):
|
|
spec = {
|
|
"start_anchor": "<assistant>",
|
|
"fields": {"content": {"content": "text"}},
|
|
}
|
|
with self.assertRaisesRegex(ValueError, "requires `prefix`"):
|
|
ResponseParser(spec)
|
|
|
|
|
|
# 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: dict[str, Any] = {
|
|
"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_integer_accepts_whole_number_floats(self):
|
|
self.assertEqual(_coerce("3.0", ("integer",)), 3)
|
|
self.assertIsInstance(_coerce("3.0", ("integer",)), int)
|
|
self.assertEqual(_coerce("3.5", ("integer",)), "3.5")
|
|
self.assertEqual(_coerce("3.5", ("integer", "number")), 3.5)
|
|
self.assertEqual(_coerce("inf", ("integer",)), "inf")
|
|
self.assertEqual(_coerce("NaN", ("integer",)), "NaN")
|
|
|
|
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: dict[str, Any] = {
|
|
"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({"oneOf": [{"type": "number"}, {"type": "null"}]}), ("number", "null"))
|
|
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_parse_response_tools_coerces_one_of_string_args(self):
|
|
tools = _set_alarm_tools(hour={"oneOf": [{"type": "integer"}, {"type": "null"}]})
|
|
parsed = parse_response(_SET_ALARM_CALL, _XML_STRING_ARGS_TEMPLATE, prefix="", tools=tools)
|
|
self.assertEqual(_first_tool_args(parsed)["hour"], 7)
|
|
|
|
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: dict[str, Any] = 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: dict[str, Any] = {
|
|
"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_already_decoded_array_is_not_cast_element_wise(self):
|
|
tools = _set_alarm_tools(groups={"type": "array", "items": {"type": "string"}})
|
|
call: dict[str, Any] = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "set_alarm",
|
|
"arguments": {"groups": ["[1,2]", "[]"]},
|
|
},
|
|
}
|
|
_parser_with_tools(tools)._coerce_tool_calls(call)
|
|
self.assertEqual(call["function"]["arguments"], {"groups": ["[1,2]", "[]"]})
|
|
|
|
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: dict[str, Any] = {
|
|
"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_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()
|
|
# fmt: on
|