[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:
@@ -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
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user