[Bugfix][GPT-OSS] Fix strict tool-call grammar to accept Harmony renders (#52222)

Signed-off-by: Yifan Zong <yzong@redhat.com>
Co-authored-by: Ben Browning <56071+bbrowning@users.noreply.github.com>
This commit is contained in:
yzong-rh
2026-08-26 18:05:09 -07:00
committed by GitHub
co-authored by Ben Browning
parent d1e5e66ee3
commit 51d1b45182
2 changed files with 111 additions and 81 deletions
+87 -61
View File
@@ -89,8 +89,11 @@ def encode_output(harmony_str: str) -> list[int]:
return get_encoding().encode(harmony_str, allowed_special="all")
def assistant(content: str, channel: str) -> Message:
return Message.from_role_and_content(Role.ASSISTANT, content).with_channel(channel)
def assistant(content: str, channel: str, content_type: str | None = None) -> Message:
message = Message.from_role_and_content(Role.ASSISTANT, content).with_channel(
channel
)
return message if content_type is None else message.with_content_type(content_type)
def tool_call(
@@ -103,16 +106,14 @@ def tool_call(
return message if content_type is None else message.with_content_type(content_type)
def get_model_output_tokens(
prompt_messages: Sequence[Message],
response_messages: Sequence[Message],
) -> list[int]:
def get_model_output_tokens(response_messages: Sequence[Message]) -> list[int]:
enc = get_encoding()
prompt_messages = [Message.from_role_and_content(Role.USER, "x")]
# Keep analysis messages when synthesizing model-output-only token sequences
# for parser tests; the default render path drops them after a later final turn.
config = RenderConversationConfig(auto_drop_analysis=False)
prompt_ids = enc.render_conversation_for_completion(
Conversation.from_messages(list(prompt_messages)),
Conversation.from_messages(prompt_messages),
Role.ASSISTANT,
config=config,
)
@@ -124,6 +125,10 @@ def get_model_output_tokens(
return full_ids[len(prompt_ids) :]
def get_model_output_str(response_messages: Sequence[Message]) -> str:
return get_encoding().decode_utf8(get_model_output_tokens(response_messages))
def get_text(msg: Message) -> str:
return msg.content[0].text if msg.content else ""
@@ -215,13 +220,12 @@ class TestParse:
# Rendered conversation outputs.
def test_reasoning_only(self, harmony_parser, chat_request):
prompt = [Message.from_role_and_content(Role.USER, "Why?")]
response = [assistant("This is reasoning", "analysis")]
reasoning, content, tool_calls = harmony_parser.parse(
"",
chat_request,
model_output_token_ids=get_model_output_tokens(prompt, response),
model_output_token_ids=get_model_output_tokens(response),
)
assert reasoning == "This is reasoning"
@@ -229,13 +233,12 @@ class TestParse:
assert tool_calls is None
def test_content_only(self, harmony_parser, chat_request):
prompt = [Message.from_role_and_content(Role.USER, "Hello")]
response = [assistant("This is a test", "final")]
reasoning, content, tool_calls = harmony_parser.parse(
"",
chat_request,
model_output_token_ids=get_model_output_tokens(prompt, response),
model_output_token_ids=get_model_output_tokens(response),
)
assert reasoning is None
@@ -243,7 +246,6 @@ class TestParse:
assert tool_calls is None
def test_reasoning_and_content(self, harmony_parser, chat_request):
prompt = [Message.from_role_and_content(Role.USER, "What is 2+2?")]
response = [
assistant("I should think first.", "analysis"),
assistant("The answer is 4.", "final"),
@@ -252,7 +254,7 @@ class TestParse:
reasoning, content, tool_calls = harmony_parser.parse(
"",
chat_request,
model_output_token_ids=get_model_output_tokens(prompt, response),
model_output_token_ids=get_model_output_tokens(response),
)
assert reasoning == "I should think first."
@@ -270,15 +272,12 @@ class TestParse:
def test_single_tool_call(
self, harmony_parser, chat_request, tool_args, tool_channel
):
prompt = [
Message.from_role_and_content(Role.USER, "What is the weather in Tokyo?")
]
response = [tool_call("functions.get_current_weather", tool_args, tool_channel)]
reasoning, content, tool_calls = harmony_parser.parse(
"",
chat_request,
model_output_token_ids=get_model_output_tokens(prompt, response),
model_output_token_ids=get_model_output_tokens(response),
)
assert reasoning is None
@@ -288,11 +287,6 @@ class TestParse:
]
def test_multiple_tool_calls_varied_formats(self, harmony_parser, chat_request):
prompt = [
Message.from_role_and_content(
Role.USER, "What is the weather in Tokyo based on where I'm at?"
)
]
response = [
tool_call("functions.get_current_weather", '{"location": "Tokyo"}'),
tool_call("functions.get_user_location", '{"location": "Tokyo"}'),
@@ -309,7 +303,7 @@ class TestParse:
_, content, tool_calls = harmony_parser.parse(
"",
chat_request,
model_output_token_ids=get_model_output_tokens(prompt, response),
model_output_token_ids=get_model_output_tokens(response),
)
assert content is None
@@ -323,13 +317,12 @@ class TestParse:
]
def test_tool_call_bare_recipient(self, harmony_parser, chat_request):
prompt = [Message.from_role_and_content(Role.USER, "Weather?")]
response = [tool_call("get_current_weather", '{"location": "Tokyo"}')]
_, _, tool_calls = harmony_parser.parse(
"",
chat_request,
model_output_token_ids=get_model_output_tokens(prompt, response),
model_output_token_ids=get_model_output_tokens(response),
)
assert tool_call_tuples(tool_calls) == [
@@ -337,7 +330,6 @@ class TestParse:
]
def test_multiple_tool_calls_bare_recipients(self, harmony_parser, chat_request):
prompt = [Message.from_role_and_content(Role.USER, "Use both tools.")]
response = [
tool_call("get_current_weather", '{"location": "Tokyo"}'),
tool_call("get_user_location", "{}"),
@@ -346,7 +338,7 @@ class TestParse:
_, _, tool_calls = harmony_parser.parse(
"",
chat_request,
model_output_token_ids=get_model_output_tokens(prompt, response),
model_output_token_ids=get_model_output_tokens(response),
)
assert tool_call_tuples(tool_calls) == [
@@ -355,7 +347,6 @@ class TestParse:
]
def test_assistant_recipient_not_tool(self, harmony_parser, chat_request):
prompt = [Message.from_role_and_content(Role.USER, "Hello")]
response = [
tool_call("assistant", "Some tool response", content_type=None),
assistant("Here is the answer", "final"),
@@ -364,7 +355,7 @@ class TestParse:
reasoning, content, tool_calls = harmony_parser.parse(
"",
chat_request,
model_output_token_ids=get_model_output_tokens(prompt, response),
model_output_token_ids=get_model_output_tokens(response),
)
assert reasoning is None
@@ -372,13 +363,12 @@ class TestParse:
assert tool_calls is None
def test_tool_call_dotted_name(self, harmony_parser, chat_request):
prompt = [Message.from_role_and_content(Role.USER, "Compute 2+3")]
response = [tool_call("math.sum", '{"a": 2, "b": 3}')]
_, _, tool_calls = harmony_parser.parse(
"",
chat_request,
model_output_token_ids=get_model_output_tokens(prompt, response),
model_output_token_ids=get_model_output_tokens(response),
)
assert tool_call_tuples(tool_calls) == [
@@ -386,7 +376,6 @@ class TestParse:
]
def test_tool_calls_with_final_content(self, harmony_parser, chat_request):
prompt = [Message.from_role_and_content(Role.USER, "What is the weather?")]
response = [
assistant("User asked about the weather.", "analysis"),
tool_call("functions.get_current_weather", '{"location": "Tokyo"}'),
@@ -396,7 +385,7 @@ class TestParse:
reasoning, content, tool_calls = harmony_parser.parse(
"",
chat_request,
model_output_token_ids=get_model_output_tokens(prompt, response),
model_output_token_ids=get_model_output_tokens(response),
)
assert reasoning == "User asked about the weather."
@@ -700,12 +689,11 @@ class TestParseDelta:
recipient,
):
parser = HarmonyParser(gpt_oss_tokenizer)
prompt = [Message.from_role_and_content(Role.USER, "Hello")]
response = [tool_call(recipient, "Ignore this", content_type=None)]
delta = parser.parse_delta(
delta_text="",
delta_token_ids=get_model_output_tokens(prompt, response),
delta_token_ids=get_model_output_tokens(response),
request=chat_request,
finished=False,
)
@@ -914,35 +902,61 @@ class TestAdjustRequest:
"required": ["answer"],
}
ANALYSIS = "<|channel|>analysis<|message|>analysis message<|end|><|start|>assistant"
COMMENTARY = (
"<|channel|>commentary<|message|>commentary message<|end|><|start|>assistant"
ANALYSIS = get_model_output_str([assistant("reasoning", "analysis")])
COMMENTARY = get_model_output_str([assistant("commentary", "commentary")])
TOOL_CALL_1_CHANNEL_FIRST = (
ANALYSIS + f"<|start|>assistant<|channel|>commentary to=functions.{TOOL_1_NAME}"
" <|constrain|>json<|message|>{}<|call|>"
)
TOOL_CALL_1 = (
ANALYSIS + f"<|channel|>commentary to=functions.{TOOL_1_NAME} json<|message|>"
"{}<|call|>"
TOOL_CALL_1_FUNCTION_FIRST = get_model_output_str(
[
assistant("reasoning", "analysis"),
tool_call(f"functions.{TOOL_1_NAME}", "{}"),
]
)
TOOL_CALL_2 = (
ANALYSIS + f"<|channel|>commentary to=functions.{TOOL_2_NAME} json<|message|>"
'{"city": "Tokyo"}<|call|>'
TOOL_CALL_2_CHANNEL_FIRST = (
ANALYSIS + f"<|start|>assistant<|channel|>commentary to=functions.{TOOL_2_NAME}"
' <|constrain|>json<|message|>{"city": "Tokyo"}<|call|>'
)
FINAL_JSON_SCHEMA = (
ANALYSIS
+ '<|channel|>final <|constrain|>json<|message|>{"answer": "Tokyo"}<|end|>'
TOOL_CALL_2_FUNCTION_FIRST = get_model_output_str(
[
assistant("reasoning", "analysis"),
tool_call(f"functions.{TOOL_2_NAME}", '{"city": "Tokyo"}'),
]
)
FINAL_JSON_OBJECT = (
ANALYSIS
+ '<|channel|>final <|constrain|>json<|message|>{"city": "Tokyo"}<|end|>'
FINAL_JSON_SCHEMA = get_model_output_str(
[
assistant("reasoning", "analysis"),
assistant('{"answer": "Tokyo"}', "final", "json"),
]
)
FINAL_JSON_OBJECT = get_model_output_str(
[
assistant("reasoning", "analysis"),
assistant('{"city": "Tokyo"}', "final", "<|constrain|>json"),
]
)
FINAL_TEXT_ONLY = get_model_output_str(
[assistant("reasoning", "analysis"), assistant("any", "final")]
)
FINAL_REGEX = get_model_output_str(
[assistant("reasoning", "analysis"), assistant("regex", "final")]
)
FINAL_CHOICE = get_model_output_str(
[assistant("reasoning", "analysis"), assistant("choice1", "final")]
)
FINAL_GRAMMAR = get_model_output_str(
[assistant("reasoning", "analysis"), assistant("grammar", "final")]
)
FINAL_STRUCTURAL_TAG = get_model_output_str(
[assistant("reasoning", "analysis"), assistant("tag content", "final")]
)
FINAL_TEXT_ONLY = ANALYSIS + "<|channel|>final<|message|>any<|end|>"
FINAL_REGEX = ANALYSIS + "<|channel|>final<|message|>regex<|end|>"
FINAL_CHOICE = ANALYSIS + "<|channel|>final<|message|>choice1<|end|>"
FINAL_GRAMMAR = ANALYSIS + "<|channel|>final<|message|>grammar<|end|>"
FINAL_STRUCTURAL_TAG = ANALYSIS + "<|channel|>final<|message|>tag content<|end|>"
ADMISSION_SAMPLES = (
"COMMENTARY",
"TOOL_CALL_1",
"TOOL_CALL_2",
"TOOL_CALL_1_CHANNEL_FIRST",
"TOOL_CALL_1_FUNCTION_FIRST",
"TOOL_CALL_2_CHANNEL_FIRST",
"TOOL_CALL_2_FUNCTION_FIRST",
"FINAL_JSON_SCHEMA",
"FINAL_JSON_OBJECT",
"FINAL_TEXT_ONLY",
@@ -1090,8 +1104,10 @@ class TestAdjustRequest:
{"tool_choice": "auto", "strict_tools": True},
[
"COMMENTARY",
"TOOL_CALL_1",
"TOOL_CALL_2",
"TOOL_CALL_1_CHANNEL_FIRST",
"TOOL_CALL_1_FUNCTION_FIRST",
"TOOL_CALL_2_CHANNEL_FIRST",
"TOOL_CALL_2_FUNCTION_FIRST",
"FINAL_JSON_SCHEMA",
"FINAL_JSON_OBJECT",
"FINAL_TEXT_ONLY",
@@ -1103,11 +1119,21 @@ class TestAdjustRequest:
),
(
{"tool_choice": "required"},
["COMMENTARY", "TOOL_CALL_1", "TOOL_CALL_2"],
[
"COMMENTARY",
"TOOL_CALL_1_CHANNEL_FIRST",
"TOOL_CALL_1_FUNCTION_FIRST",
"TOOL_CALL_2_CHANNEL_FIRST",
"TOOL_CALL_2_FUNCTION_FIRST",
],
),
(
{"tool_choice": "named"},
["COMMENTARY", "TOOL_CALL_2"],
[
"COMMENTARY",
"TOOL_CALL_2_CHANNEL_FIRST",
"TOOL_CALL_2_FUNCTION_FIRST",
],
),
(
{"response_format_type": "json_schema"},
+24 -20
View File
@@ -387,6 +387,7 @@ class HarmonyParser(DelegatingParser):
return recipient[:constrain_index].rstrip() or None
_JSON_CONSTRAINS = [" json", " <|constrain|>json"]
# Harmomy's stop tokens are <|return|>, <|call|>, <|endoftext|>
# They are represented as "" since xgrammar disallows stop tokens while
# under constraints, leading to bad or infinite generations.
@@ -398,10 +399,8 @@ _TOOL_CALL_CHANNELS = [
"<|channel|>final",
]
_FUNCTION_CALL_BEGINS = [
"to=functions.{name} {channel} json<|message|>",
"to=functions.{name} {channel} <|constrain|>json<|message|>",
"{channel} to=functions.{name} json<|message|>",
"{channel} to=functions.{name} <|constrain|>json<|message|>",
" to=functions.{name}{channel}{constrain}<|message|>",
"{channel} to=functions.{name}{constrain}<|message|>",
]
_JSON_CONTENT = JSONSchemaFormat(json_schema={"type": "object"})
_ANY_CONTENT = AnyTextFormat()
@@ -476,7 +475,9 @@ def get_harmony_structural_tag(
else:
tags = [
TagFormat(
begin=pattern.format(name=tool.function.name, channel=channel),
begin=pattern.format(
name=tool.function.name, channel=channel, constrain=constrain
),
content=JSONSchemaFormat(
json_schema=get_function_parameters(tool.function)
),
@@ -484,23 +485,18 @@ def get_harmony_structural_tag(
)
for tool in tools
for pattern in _FUNCTION_CALL_BEGINS
for constrain in _JSON_CONSTRAINS
for channel in _TOOL_CALL_CHANNELS
]
if tool_choice == "auto":
tags.append(
tags.extend(
TagFormat(
begin=_FINAL_BEGIN.format(constrain=" <|constrain|>json"),
content=_ANY_CONTENT,
end=_END_TAG,
)
)
tags.append(
TagFormat(
begin=_FINAL_BEGIN.format(constrain=""),
begin=_FINAL_BEGIN.format(constrain=constrain),
content=_ANY_CONTENT,
end=_END_TAG,
)
for constrain in _JSON_CONSTRAINS + [""]
)
return _assemble_tag(
@@ -558,14 +554,22 @@ def _adjust_output_format(
return request
if isinstance(final_content, JSONSchemaFormat):
begin = _FINAL_BEGIN.format(constrain=" <|constrain|>json")
content = OrFormat(
elements=[
TagFormat(
begin=_FINAL_BEGIN.format(constrain=constrain),
content=final_content,
end=_END_TAG,
)
for constrain in _JSON_CONSTRAINS
]
)
else:
begin = _FINAL_BEGIN.format(constrain="")
content = TagFormat(
begin=_FINAL_BEGIN.format(constrain=""), content=final_content, end=_END_TAG
)
structural_tag = _assemble_tag(
allow_analysis=True,
allow_commentary=False,
content=TagFormat(begin=begin, content=final_content, end=_END_TAG),
allow_analysis=True, allow_commentary=False, content=content
)
request.structured_outputs = replace(