wehub-resource-sync
481 lines
18 KiB
Python
481 lines
18 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for Harmony streaming parser (omlx.adapter.harmony)."""
|
|
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from omlx.adapter.harmony import (
|
|
HarmonyStreamingParser,
|
|
load_harmony_gpt_oss_encoding,
|
|
parse_tool_calls_from_tokens,
|
|
preprocess_harmony_messages,
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def encoding():
|
|
"""Load HarmonyGptOss encoding."""
|
|
return load_harmony_gpt_oss_encoding()
|
|
|
|
|
|
@pytest.fixture
|
|
def tokenizer(encoding):
|
|
"""Create a mock tokenizer that resolves Harmony special tokens."""
|
|
tok = MagicMock()
|
|
special_map = {}
|
|
for name in [
|
|
"<|start|>", "<|end|>", "<|message|>", "<|channel|>",
|
|
"<|return|>", "<|call|>", "<|constrain|>",
|
|
]:
|
|
tokens = encoding.encode(name, allowed_special="all")
|
|
if tokens:
|
|
special_map[name] = tokens[0]
|
|
|
|
def _convert(token_str):
|
|
return special_map.get(token_str, -1)
|
|
|
|
tok.convert_tokens_to_ids = _convert
|
|
return tok
|
|
|
|
|
|
@pytest.fixture
|
|
def parser(tokenizer):
|
|
"""Create a HarmonyStreamingParser instance (pre-primed with <|start|>assistant)."""
|
|
return HarmonyStreamingParser(tokenizer=tokenizer)
|
|
|
|
|
|
# ── Tool call header parsing ──────────────────────────────────────────
|
|
# The parser is pre-primed with "<|start|>assistant", so the first
|
|
# model output token is <|channel|> (regular message) or " to=..."
|
|
# (tool call continuation). Subsequent messages start with <|start|>.
|
|
|
|
|
|
class TestToolCallParsing:
|
|
"""StreamableParser with role=None correctly parses tool call headers."""
|
|
|
|
def test_tool_call_header_parsed(self, parser, encoding):
|
|
"""Tool call header 'to=functions.Write' sets commentary channel.
|
|
|
|
The first message's <|start|>assistant is already primed.
|
|
Model output for a tool call starts with ' to=functions.Write<|channel|>...'
|
|
"""
|
|
# Model continues from primed "assistant" with " to=functions.Write..."
|
|
tokens = encoding.encode(
|
|
" to=functions.Write<|channel|>commentary<|message|>args",
|
|
allowed_special="all",
|
|
)
|
|
for t in tokens:
|
|
parser.process_token(t)
|
|
|
|
assert parser.current_channel == "commentary"
|
|
assert parser.current_recipient == "functions.Write"
|
|
|
|
def test_regular_message_still_works(self, parser, encoding):
|
|
"""Regular assistant message parses correctly (primed parser)."""
|
|
# Model output starts from <|channel|> (first message)
|
|
tokens = encoding.encode(
|
|
"<|channel|>final<|message|>Hello",
|
|
allowed_special="all",
|
|
)
|
|
results = [parser.process_token(t) for t in tokens]
|
|
|
|
assert parser.current_channel == "final"
|
|
|
|
# "Hello" token should be streamed and visible
|
|
last = results[-1]
|
|
control_text, stream_token, visible_token, is_stop = last
|
|
assert stream_token is not None
|
|
assert visible_token is not None
|
|
|
|
def test_tool_call_tokens_not_streamed(self, parser, encoding):
|
|
"""Commentary channel tokens return None for stream/visible."""
|
|
tokens = encoding.encode(
|
|
" to=functions.Write<|channel|>commentary<|message|>args",
|
|
allowed_special="all",
|
|
)
|
|
results = [parser.process_token(t) for t in tokens]
|
|
|
|
# "args" is in commentary channel — should not be streamed
|
|
last = results[-1]
|
|
_, stream_token, visible_token, _ = last
|
|
assert stream_token is None
|
|
assert visible_token is None
|
|
|
|
def test_get_tool_calls_channel_rules(self):
|
|
"""Commentary always counts; analysis counts only with JSON arguments."""
|
|
parser = object.__new__(HarmonyStreamingParser)
|
|
parser._parser = SimpleNamespace(
|
|
messages=[
|
|
SimpleNamespace(
|
|
channel="analysis",
|
|
recipient="functions.Read",
|
|
content=[SimpleNamespace(text="analysis-noise")],
|
|
),
|
|
SimpleNamespace(
|
|
channel="analysis",
|
|
recipient="functions.Read",
|
|
content=[SimpleNamespace(text='{"path":"from-analysis.py"}')],
|
|
),
|
|
SimpleNamespace(
|
|
channel="commentary",
|
|
recipient="functions.Read",
|
|
content=[SimpleNamespace(text='{"path":"ok.py"}')],
|
|
),
|
|
SimpleNamespace(
|
|
channel="final",
|
|
recipient="functions.Read",
|
|
content=[SimpleNamespace(text="final-noise")],
|
|
),
|
|
]
|
|
)
|
|
|
|
assert parser.get_tool_calls() == [
|
|
{"name": "Read", "arguments": '{"path":"from-analysis.py"}'},
|
|
{"name": "Read", "arguments": '{"path":"ok.py"}'},
|
|
]
|
|
|
|
|
|
# ── Multi-message sequences ──────────────────────────────────────────
|
|
|
|
|
|
class TestMultiMessageSequence:
|
|
"""Correct channel transitions across analysis → final → tool call."""
|
|
|
|
def test_analysis_final_tool_sequence(self, parser, encoding):
|
|
"""Channels transition correctly: analysis → final → commentary.
|
|
|
|
First message uses primed header (starts from <|channel|>).
|
|
Subsequent messages include full <|start|>assistant header.
|
|
"""
|
|
# First message: analysis (primed, starts from <|channel|>)
|
|
analysis = encoding.encode(
|
|
"<|channel|>analysis<|message|>thinking<|end|>",
|
|
allowed_special="all",
|
|
)
|
|
# Second message: final (full header)
|
|
final = encoding.encode(
|
|
"<|start|>assistant<|channel|>final<|message|>result<|end|>",
|
|
allowed_special="all",
|
|
)
|
|
# Third message: tool call (full header with routing)
|
|
tool = encoding.encode(
|
|
"<|start|>assistant to=functions.Read"
|
|
'<|channel|>commentary<|message|>{"path":"f.py"}<|end|>',
|
|
allowed_special="all",
|
|
)
|
|
|
|
channels_seen = []
|
|
for t in analysis + final + tool:
|
|
parser.process_token(t)
|
|
ch = parser.current_channel
|
|
if not channels_seen or channels_seen[-1] != ch:
|
|
channels_seen.append(ch)
|
|
|
|
assert "analysis" in channels_seen
|
|
assert "final" in channels_seen
|
|
assert "commentary" in channels_seen
|
|
|
|
def test_think_tags_emitted(self, parser, encoding):
|
|
"""<think>/<think> control text emitted on analysis channel transitions."""
|
|
# First message: analysis (primed)
|
|
# Second message: final (full header)
|
|
tokens = encoding.encode(
|
|
"<|channel|>analysis<|message|>thought<|end|>"
|
|
"<|start|>assistant<|channel|>final<|message|>answer",
|
|
allowed_special="all",
|
|
)
|
|
control_texts = []
|
|
for t in tokens:
|
|
control_text, _, _, _ = parser.process_token(t)
|
|
if control_text:
|
|
control_texts.append(control_text)
|
|
|
|
assert "<think>\n" in control_texts
|
|
assert "</think>\n" in control_texts
|
|
|
|
|
|
# ── Passthrough mode ─────────────────────────────────────────────────
|
|
|
|
|
|
class TestPassthroughMode:
|
|
"""Passthrough mode activates on parser error and prevents cascading."""
|
|
|
|
def test_passthrough_activates_on_error(self, tokenizer, encoding):
|
|
"""Parser switches to passthrough when process() raises."""
|
|
parser = HarmonyStreamingParser(tokenizer=tokenizer)
|
|
|
|
# Force an error by monkey-patching the inner parser
|
|
def _failing_process(token_id):
|
|
raise RuntimeError("simulated parser failure")
|
|
|
|
parser._parser.process = _failing_process
|
|
|
|
result = parser.process_token(42)
|
|
control_text, stream_token, visible_token, is_stop = result
|
|
|
|
assert parser._passthrough_mode is True
|
|
assert stream_token is None
|
|
assert visible_token is None
|
|
|
|
def test_passthrough_buffers_silently(self, tokenizer, encoding):
|
|
"""After passthrough activation, all tokens return None."""
|
|
parser = HarmonyStreamingParser(tokenizer=tokenizer)
|
|
parser._passthrough_mode = True
|
|
|
|
for token_id in [100, 200, 300]:
|
|
control_text, stream_token, visible_token, is_stop = (
|
|
parser.process_token(token_id)
|
|
)
|
|
assert stream_token is None
|
|
assert visible_token is None
|
|
assert control_text == ""
|
|
|
|
def test_stop_token_detected_in_passthrough(self, tokenizer, encoding):
|
|
"""Stop tokens are still detected in passthrough mode."""
|
|
parser = HarmonyStreamingParser(tokenizer=tokenizer)
|
|
parser._passthrough_mode = True
|
|
|
|
stop_ids = parser.get_stop_token_ids()
|
|
assert len(stop_ids) > 0
|
|
|
|
stop_token = next(iter(stop_ids))
|
|
_, _, _, is_stop = parser.process_token(stop_token)
|
|
assert is_stop is True
|
|
|
|
def test_analysis_end_does_not_stop_generation(self, tokenizer, encoding):
|
|
"""Ending analysis closes thinking but lets the final channel continue."""
|
|
parser = HarmonyStreamingParser(tokenizer=tokenizer)
|
|
|
|
tokens = encoding.encode(
|
|
"<|channel|>analysis<|message|>thinking<|end|>",
|
|
allowed_special="all",
|
|
)
|
|
|
|
result = None
|
|
for token_id in tokens:
|
|
result = parser.process_token(token_id)
|
|
|
|
assert result is not None
|
|
control_text, stream_token, visible_token, is_stop = result
|
|
assert "</think>\n" in control_text
|
|
assert stream_token is None
|
|
assert visible_token is None
|
|
assert is_stop is False
|
|
|
|
def test_passthrough_closes_think_tag(self, tokenizer, encoding):
|
|
"""Passthrough activation while in analysis channel closes think tag."""
|
|
parser = HarmonyStreamingParser(tokenizer=tokenizer)
|
|
|
|
# Process analysis channel header (primed, starts from <|channel|>)
|
|
analysis_header = encoding.encode(
|
|
"<|channel|>analysis<|message|>",
|
|
allowed_special="all",
|
|
)
|
|
for t in analysis_header:
|
|
parser.process_token(t)
|
|
|
|
assert parser._in_think_tag is True
|
|
|
|
# Now trigger passthrough
|
|
def _failing_process(token_id):
|
|
raise RuntimeError("simulated failure")
|
|
|
|
parser._parser.process = _failing_process
|
|
|
|
control_text, _, _, _ = parser.process_token(999)
|
|
assert parser._passthrough_mode is True
|
|
assert "</think>\n" in control_text
|
|
assert parser._in_think_tag is False
|
|
|
|
def test_reset_clears_passthrough(self, tokenizer):
|
|
"""reset() clears passthrough mode and re-primes parser."""
|
|
parser = HarmonyStreamingParser(tokenizer=tokenizer)
|
|
parser._passthrough_mode = True
|
|
|
|
parser.reset()
|
|
assert parser._passthrough_mode is False
|
|
|
|
def test_reset_parser_accepts_channel_token(self, parser, encoding):
|
|
"""After reset, parser accepts <|channel|> (re-primed)."""
|
|
parser.reset()
|
|
|
|
# Should not raise — parser is re-primed with <|start|>assistant
|
|
tokens = encoding.encode(
|
|
"<|channel|>final<|message|>test",
|
|
allowed_special="all",
|
|
)
|
|
for t in tokens:
|
|
parser.process_token(t)
|
|
|
|
assert parser.current_channel == "final"
|
|
|
|
|
|
# ── Preprocessing ─────────────────────────────────────────────────────
|
|
|
|
|
|
class TestPreprocessHarmonyMessages:
|
|
"""preprocess_harmony_messages strips think tags from assistant content."""
|
|
|
|
def test_strips_think_tags(self):
|
|
msgs = [
|
|
{"role": "assistant", "content": "<think>reasoning</think>answer"}
|
|
]
|
|
result = preprocess_harmony_messages(msgs)
|
|
assert result[0]["content"] == "answer"
|
|
|
|
def test_passes_tool_messages(self):
|
|
msgs = [
|
|
{"role": "tool", "tool_call_id": "123", "content": "result"}
|
|
]
|
|
result = preprocess_harmony_messages(msgs)
|
|
assert result == msgs
|
|
|
|
|
|
# ── parse_tool_calls_from_tokens ──────────────────────────────────────
|
|
|
|
|
|
class TestParseToolCallsFromTokens:
|
|
"""Non-streaming tool call extraction from complete token sequences."""
|
|
|
|
def test_extracts_tool_call(self, encoding):
|
|
"""Extracts function name and arguments from tool call tokens."""
|
|
# Model output starts from <|channel|> (prompt includes <|start|>assistant)
|
|
tokens = encoding.encode(
|
|
'<|channel|>commentary<|message|>{"path":"t.py"}<|end|>'
|
|
"<|return|>",
|
|
allowed_special="all",
|
|
)
|
|
# prepend_start=True adds <|start|>assistant
|
|
output_text, analysis_text, tool_calls = parse_tool_calls_from_tokens(
|
|
tokens, prepend_start=True
|
|
)
|
|
assert analysis_text == ""
|
|
assert isinstance(tool_calls, list)
|
|
|
|
def test_extracts_final_text(self, encoding):
|
|
"""Extracts text from final channel."""
|
|
tokens = encoding.encode(
|
|
"<|channel|>final<|message|>Hello world<|end|>",
|
|
allowed_special="all",
|
|
)
|
|
output_text, analysis_text, tool_calls = parse_tool_calls_from_tokens(
|
|
tokens, prepend_start=True
|
|
)
|
|
assert analysis_text == ""
|
|
assert "Hello world" in output_text
|
|
|
|
def test_unknown_channel_without_recipient_falls_back_to_final(self, encoding):
|
|
"""Malformed channel names without recipients remain visible."""
|
|
tokens = encoding.encode(
|
|
"<|channel|>mardown<|message|>visible text<|return|>",
|
|
allowed_special="all",
|
|
)
|
|
|
|
output_text, analysis_text, tool_calls = parse_tool_calls_from_tokens(
|
|
tokens, prepend_start=True
|
|
)
|
|
|
|
assert "visible text" in output_text
|
|
assert analysis_text == ""
|
|
assert tool_calls == []
|
|
|
|
def test_does_not_prepend_duplicate_start_header(self, encoding):
|
|
"""Budget-forced Harmony completions may already include the start header."""
|
|
tokens = encoding.encode(
|
|
"<|start|>assistant<|channel|>analysis<|message|>thinking<|end|>",
|
|
allowed_special="all",
|
|
)
|
|
output_text, analysis_text, tool_calls = parse_tool_calls_from_tokens(
|
|
tokens, prepend_start=True
|
|
)
|
|
assert output_text == ""
|
|
assert "thinking" in analysis_text
|
|
assert tool_calls == []
|
|
|
|
def test_channel_rules_for_function_recipient(self, monkeypatch):
|
|
"""Commentary and JSON-argument analysis recipients become tool calls."""
|
|
messages = [
|
|
SimpleNamespace(
|
|
channel="analysis",
|
|
recipient="functions.Read",
|
|
content=[SimpleNamespace(text="analysis-noise")],
|
|
),
|
|
SimpleNamespace(
|
|
channel="analysis",
|
|
recipient="functions.Read",
|
|
content=[SimpleNamespace(text='{"path":"from-analysis.py"}')],
|
|
),
|
|
SimpleNamespace(
|
|
channel="commentary",
|
|
recipient="functions.Read",
|
|
content=[SimpleNamespace(text='{"path":"ok.py"}')],
|
|
),
|
|
SimpleNamespace(
|
|
channel="final",
|
|
recipient="functions.Read",
|
|
content=[SimpleNamespace(text="final-noise")],
|
|
),
|
|
SimpleNamespace(
|
|
channel="other",
|
|
recipient="functions.Read",
|
|
content=[SimpleNamespace(text="other-noise")],
|
|
),
|
|
]
|
|
|
|
class FakeEncoding:
|
|
def encode(self, text, allowed_special="all"):
|
|
return [1, 2] if text == "<|start|>assistant" else [3]
|
|
|
|
def decode(self, token_ids):
|
|
return "decoded"
|
|
|
|
def parse_messages_from_completion_tokens(self, token_ids, role, strict):
|
|
return messages
|
|
|
|
monkeypatch.setattr(
|
|
"omlx.adapter.harmony.load_harmony_gpt_oss_encoding",
|
|
lambda: FakeEncoding(),
|
|
)
|
|
|
|
output_text, analysis_text, tool_calls = parse_tool_calls_from_tokens([99])
|
|
|
|
assert analysis_text == "analysis-noise"
|
|
assert output_text == "final-noise"
|
|
assert tool_calls == [
|
|
{"name": "Read", "arguments": '{"path":"from-analysis.py"}'},
|
|
{"name": "Read", "arguments": '{"path":"ok.py"}'},
|
|
]
|
|
|
|
def test_analysis_channel_tool_call_with_recipient(self, encoding):
|
|
"""Analysis-channel tool calls with explicit recipients are honored (#2216)."""
|
|
tokens = encoding.encode(
|
|
"<|channel|>analysis<|message|>Now need files.<|end|>"
|
|
"<|start|>assistant<|channel|>analysis to=functions.read code"
|
|
'<|message|>{"filePath": "/tmp/x.py", "offset": 1}<|call|>',
|
|
allowed_special="all",
|
|
)
|
|
output_text, analysis_text, tool_calls = parse_tool_calls_from_tokens(
|
|
tokens, prepend_start=True
|
|
)
|
|
assert tool_calls == [
|
|
{"name": "read", "arguments": '{"filePath": "/tmp/x.py", "offset": 1}'}
|
|
]
|
|
assert "Now need files." in analysis_text
|
|
assert "filePath" not in analysis_text
|
|
|
|
def test_analysis_recipient_with_prose_stays_reasoning(self, encoding):
|
|
"""Analysis messages addressed to a tool without JSON args stay reasoning."""
|
|
tokens = encoding.encode(
|
|
"<|channel|>analysis<|message|>ok<|end|>"
|
|
"<|start|>assistant<|channel|>analysis to=functions.read code"
|
|
"<|message|>maybe I should read the file<|call|>",
|
|
allowed_special="all",
|
|
)
|
|
output_text, analysis_text, tool_calls = parse_tool_calls_from_tokens(
|
|
tokens, prepend_start=True
|
|
)
|
|
assert tool_calls == []
|
|
assert "maybe I should read the file" in analysis_text
|