wehub-resource-sync
526 lines
19 KiB
Python
526 lines
19 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Harmony format streaming parser for gpt-oss models.
|
|
|
|
Uses the official openai-harmony package for robust parsing.
|
|
|
|
Harmony protocol uses special tokens to structure messages:
|
|
- <|start|>: Begin message header
|
|
- <|channel|>: Mark channel type
|
|
- <|message|>: Transition to content
|
|
- <|end|>: End message
|
|
- <|return|>: Model completion signal
|
|
- <|call|>: Tool invocation signal
|
|
|
|
Message structure: <|start|>{header}<|channel|>{channel_name}<|message|>{content}<|end|>
|
|
|
|
Channels:
|
|
- final: User-visible response (plain text)
|
|
- analysis: Chain-of-thought reasoning (wrapped in <think>...</think> for streaming)
|
|
- commentary: Tool/function calls (non-streaming only)
|
|
|
|
gpt-oss also emits tool calls on the analysis channel with an explicit
|
|
``to=functions.*`` recipient; those are honored as tool calls when the
|
|
arguments form a JSON object (#2216).
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import re
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from functools import lru_cache
|
|
from typing import Any
|
|
|
|
from openai_harmony import (
|
|
HarmonyEncoding,
|
|
Role,
|
|
StreamableParser,
|
|
load_harmony_encoding,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Pattern to match <think>...</think> blocks
|
|
_THINK_TAG_PATTERN = re.compile(r"<think>.*?</think>\s*", re.DOTALL)
|
|
|
|
# Pre-allocated constants
|
|
_THINK_OPEN = "<think>\n"
|
|
_THINK_CLOSE = "</think>\n"
|
|
|
|
# Harmony special tokens that should not be streamed
|
|
_HARMONY_SPECIAL_TOKENS = [
|
|
"<|start|>",
|
|
"<|end|>",
|
|
"<|message|>",
|
|
"<|channel|>",
|
|
"<|return|>",
|
|
"<|call|>",
|
|
"<|constrain|>",
|
|
]
|
|
|
|
|
|
def _has_no_real_recipient(recipient: str | None) -> bool:
|
|
"""Return True when the parser only preserved the primed assistant header."""
|
|
return recipient is None or recipient == "<|start|>assistant"
|
|
|
|
|
|
def _message_content_text(msg: Any) -> str:
|
|
"""Concatenate the text parts of a parsed Harmony message."""
|
|
text = ""
|
|
content = getattr(msg, "content", None)
|
|
if content is not None:
|
|
for part in content:
|
|
part_text = getattr(part, "text", None)
|
|
if isinstance(part_text, str):
|
|
text += part_text
|
|
return text
|
|
|
|
|
|
def _is_tool_call_message(msg: Any) -> bool:
|
|
"""Return True when a parsed message is a genuine tool invocation.
|
|
|
|
Commentary is the canonical tool-call channel. gpt-oss also emits tool
|
|
calls on the analysis channel with an explicit functions.* recipient
|
|
(#2216); those count only when the arguments form a JSON object, so
|
|
recipient-less tool syntax reasoned about in thinking and prose
|
|
addressed to a tool both stay in the reasoning output (#2032).
|
|
"""
|
|
recipient = getattr(msg, "recipient", None)
|
|
if not (isinstance(recipient, str) and recipient.startswith("functions.")):
|
|
return False
|
|
channel = getattr(msg, "channel", None)
|
|
if channel == "commentary":
|
|
return True
|
|
if channel != "analysis":
|
|
return False
|
|
try:
|
|
return isinstance(json.loads(_message_content_text(msg)), dict)
|
|
except ValueError:
|
|
return False
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def load_harmony_gpt_oss_encoding() -> HarmonyEncoding:
|
|
"""Load the Harmony gpt-oss encoding with a small retry window."""
|
|
last_error: Exception | None = None
|
|
for attempt in range(3):
|
|
try:
|
|
return load_harmony_encoding("HarmonyGptOss")
|
|
except Exception as exc:
|
|
last_error = exc
|
|
if attempt == 2:
|
|
break
|
|
delay = 0.5 * (2**attempt)
|
|
logger.warning(
|
|
"Failed to load HarmonyGptOss encoding "
|
|
"(attempt %d/3): %s; retrying in %.1fs",
|
|
attempt + 1,
|
|
exc,
|
|
delay,
|
|
)
|
|
time.sleep(delay)
|
|
|
|
assert last_error is not None
|
|
raise last_error
|
|
|
|
|
|
def preprocess_harmony_messages(
|
|
messages: list[dict[str, Any]]
|
|
) -> list[dict[str, Any]]:
|
|
"""
|
|
Preprocess messages for Harmony (gpt-oss) models.
|
|
|
|
- Strips <think> tags from assistant messages
|
|
- Keeps tool role messages unchanged (chat_template handles conversion)
|
|
|
|
The chat_template expects standard OpenAI format:
|
|
- role: "tool" with tool_call_id and content
|
|
- It uses last_tool_call.name from the previous assistant message
|
|
- Generates: <|start|>functions.{name} to=assistant<|channel|>commentary<|message|>{content|tojson}<|end|>
|
|
|
|
Args:
|
|
messages: List of message dicts with 'role' and 'content' keys
|
|
|
|
Returns:
|
|
Messages preprocessed for Harmony format
|
|
"""
|
|
if not messages:
|
|
return []
|
|
|
|
result = []
|
|
|
|
for msg in messages:
|
|
# Validate message is a dict
|
|
if not isinstance(msg, dict):
|
|
logger.warning(f"Skipping non-dict message: {type(msg)}")
|
|
continue
|
|
|
|
role = msg.get("role")
|
|
|
|
if role == "assistant":
|
|
content = msg.get("content", "")
|
|
# Ensure content is a string (could be list in some formats)
|
|
if isinstance(content, str):
|
|
# Strip <think> tags
|
|
if content and "<think>" in content:
|
|
content = _THINK_TAG_PATTERN.sub("", content).strip()
|
|
msg = {**msg, "content": content}
|
|
elif content is not None:
|
|
# Non-string content (e.g., list) - log but don't modify
|
|
logger.debug(f"Assistant message has non-string content: {type(content)}")
|
|
|
|
result.append(msg)
|
|
|
|
else:
|
|
# Pass through all other messages (user, tool, system, etc.) unchanged
|
|
# Chat template handles tool messages directly using last_tool_call.name
|
|
result.append(msg)
|
|
|
|
return result
|
|
|
|
|
|
def _get_special_token_ids(tokenizer: Any) -> set[int]:
|
|
"""
|
|
Get special token IDs from model tokenizer.
|
|
|
|
Args:
|
|
tokenizer: The model's tokenizer
|
|
|
|
Returns:
|
|
Set of special token IDs
|
|
"""
|
|
special_ids = set()
|
|
for token in _HARMONY_SPECIAL_TOKENS:
|
|
try:
|
|
token_id = tokenizer.convert_tokens_to_ids(token)
|
|
if isinstance(token_id, int) and token_id >= 0:
|
|
special_ids.add(token_id)
|
|
else:
|
|
logger.debug(f"Harmony special token '{token}' not found in tokenizer")
|
|
except Exception as e:
|
|
logger.debug(f"Failed to get ID for Harmony token '{token}': {e}")
|
|
return special_ids
|
|
|
|
|
|
@dataclass
|
|
class HarmonyStreamingParser:
|
|
"""
|
|
Streaming parser for Harmony format using official openai-harmony package.
|
|
|
|
Parses tokens incrementally and routes them to appropriate channels.
|
|
Returns token IDs instead of decoded text to allow proper UTF-8 handling
|
|
via streaming detokenizer in the caller.
|
|
|
|
Output routing:
|
|
- analysis channel -> stream only (wrapped in <think>...</think>)
|
|
- final channel -> stream and visible (stored in output_text)
|
|
- commentary channel -> buffered for tool calls (non-streaming)
|
|
|
|
The parser returns:
|
|
- control_text: Control strings like <think>, </think>
|
|
- stream_token: Token ID to stream (None if not streaming)
|
|
- visible_token: Token ID to store (None if not storing)
|
|
- is_stop: Whether this is a stop signal
|
|
"""
|
|
|
|
tokenizer: Any
|
|
|
|
# Internal state (initialized in __post_init__)
|
|
_encoding: HarmonyEncoding = field(init=False, repr=False)
|
|
_parser: StreamableParser = field(init=False, repr=False)
|
|
_stop_tokens: set[int] = field(init=False, default_factory=set)
|
|
_special_tokens: set[int] = field(init=False, default_factory=set)
|
|
|
|
# <think> tag state
|
|
_in_think_tag: bool = field(init=False, default=False)
|
|
_prev_channel: str | None = field(init=False, default=None)
|
|
|
|
# Passthrough mode: activated when streaming parser encounters an
|
|
# unrecoverable error. Tokens are still accumulated by the scheduler
|
|
# (request.append_output_token) so parse_tool_calls_from_tokens can
|
|
# extract tool calls at finalization.
|
|
_passthrough_mode: bool = field(init=False, default=False)
|
|
|
|
def __post_init__(self):
|
|
"""Initialize the official Harmony parser."""
|
|
self._encoding = load_harmony_gpt_oss_encoding()
|
|
# role=None allows the parser to handle tool-call headers
|
|
# (e.g. "assistant to=functions.Write") which Role.ASSISTANT rejects.
|
|
self._parser = StreamableParser(self._encoding, None, strict=False)
|
|
self._stop_tokens = set(self._encoding.stop_tokens_for_assistant_actions())
|
|
self._special_tokens = _get_special_token_ids(self.tokenizer)
|
|
|
|
# Prime the parser with "<|start|>assistant" tokens. The chat
|
|
# template already includes these in the prompt, so the model's
|
|
# first output token is <|channel|>, not <|start|>. Without
|
|
# priming, the parser rejects <|channel|> as unexpected.
|
|
self._prime_parser(self._parser)
|
|
|
|
logger.info(
|
|
f"Harmony parser initialized: {len(self._special_tokens)} special tokens, "
|
|
f"{len(self._stop_tokens)} stop tokens"
|
|
)
|
|
|
|
def _prime_parser(self, parser: StreamableParser) -> None:
|
|
"""Feed '<|start|>assistant' header tokens so parser expects <|channel|> next."""
|
|
start_tokens = self._encoding.encode(
|
|
"<|start|>assistant", allowed_special="all"
|
|
)
|
|
for t in start_tokens:
|
|
parser.process(t)
|
|
|
|
def process_token(
|
|
self, token_id: int
|
|
) -> tuple[str, int | None, int | None, bool]:
|
|
"""
|
|
Process a single token and return routing information.
|
|
|
|
This method routes tokens to appropriate channels without decoding.
|
|
The caller should use streaming detokenizer to decode the returned
|
|
token IDs for proper UTF-8 handling.
|
|
|
|
Args:
|
|
token_id: The token ID to process.
|
|
|
|
Returns:
|
|
Tuple of:
|
|
- control_text: Control strings (<think>, </think>, etc.)
|
|
- stream_token: Token ID to stream (None to skip)
|
|
- visible_token: Token ID to store in output_text (None to skip)
|
|
- is_stop: True if this is a stop token
|
|
"""
|
|
# Check if this is a special token (should not be streamed)
|
|
is_special_token = token_id in self._special_tokens
|
|
is_stop = token_id in self._stop_tokens
|
|
was_analysis = self._prev_channel == "analysis"
|
|
|
|
# Passthrough: parser crashed earlier, buffer all tokens silently.
|
|
# Tokens are still tracked by the scheduler for non-streaming tool
|
|
# call extraction at finalization.
|
|
if self._passthrough_mode:
|
|
return "", None, None, is_stop
|
|
|
|
try:
|
|
self._parser.process(token_id)
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"Harmony streaming parser error, switching to passthrough: {e}"
|
|
)
|
|
self._passthrough_mode = True
|
|
control_text = ""
|
|
if self._in_think_tag:
|
|
control_text = _THINK_CLOSE
|
|
self._in_think_tag = False
|
|
return control_text, None, None, is_stop
|
|
|
|
channel = self._parser.current_channel
|
|
control_text = ""
|
|
|
|
# Harmony uses the same end token for analysis, final, and tool/action
|
|
# messages. Ending analysis should let generation continue into the
|
|
# final channel; ending other channels should stop the request.
|
|
if was_analysis and is_stop:
|
|
is_stop = False
|
|
|
|
# Handle channel transitions for <think> tags
|
|
if channel != self._prev_channel:
|
|
# Close previous analysis channel
|
|
if self._in_think_tag and self._prev_channel == "analysis":
|
|
control_text += _THINK_CLOSE
|
|
self._in_think_tag = False
|
|
# Open new analysis channel
|
|
if channel == "analysis" and not self._in_think_tag:
|
|
control_text += _THINK_OPEN
|
|
self._in_think_tag = True
|
|
self._prev_channel = channel
|
|
|
|
# Special tokens should never be streamed or stored
|
|
if is_special_token:
|
|
return control_text, None, None, is_stop
|
|
|
|
# Route based on channel
|
|
if channel == "final":
|
|
# final: stream AND store (same token for both)
|
|
return control_text, token_id, token_id, is_stop
|
|
elif channel == "analysis":
|
|
# analysis: stream only (wrapped in <think>)
|
|
return control_text, token_id, None, is_stop
|
|
elif channel is None:
|
|
# Channel not yet determined (still in header parsing)
|
|
# Buffer token but don't stream
|
|
return control_text, None, None, is_stop
|
|
elif channel != "commentary" and _has_no_real_recipient(
|
|
self._parser.current_recipient
|
|
):
|
|
# Some fine-tunes emit misspelled channels (for example "mardown").
|
|
# If there is no recipient, preserve the text as user-visible output.
|
|
return control_text, token_id, token_id, is_stop
|
|
else:
|
|
# commentary etc: buffer only (for tool calls)
|
|
return control_text, None, None, is_stop
|
|
|
|
def get_stop_token_ids(self) -> set[int]:
|
|
"""Get Harmony stop token IDs."""
|
|
return self._stop_tokens
|
|
|
|
def get_tool_calls(self) -> list[dict[str, str]]:
|
|
"""Get accumulated tool calls from parsed messages."""
|
|
tool_calls = []
|
|
try:
|
|
messages = self._parser.messages
|
|
if not messages:
|
|
return tool_calls
|
|
|
|
for msg in messages:
|
|
if not _is_tool_call_message(msg):
|
|
continue
|
|
|
|
name = msg.recipient[10:] # Remove "functions." prefix
|
|
content = _message_content_text(msg)
|
|
|
|
tool_calls.append({"name": name, "arguments": content})
|
|
logger.info(f"Extracted tool call: {name}, arguments={content}")
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Error extracting tool calls: {e}")
|
|
|
|
return tool_calls
|
|
|
|
def finalize(self) -> str:
|
|
"""
|
|
Finalize parsing and close any open tags.
|
|
|
|
Returns:
|
|
Any remaining control text (e.g., closing </think> tag).
|
|
"""
|
|
try:
|
|
self._parser.process_eos()
|
|
except Exception as e:
|
|
# Can fail if message is incomplete (e.g., missing <|end|>)
|
|
# This is expected in some cases, so just log and continue
|
|
logger.debug(f"Harmony parser process_eos failed (expected for incomplete messages): {e}")
|
|
|
|
if self._in_think_tag:
|
|
self._in_think_tag = False
|
|
return _THINK_CLOSE
|
|
return ""
|
|
|
|
def reset(self) -> None:
|
|
"""Reset parser state for a new request."""
|
|
self._parser = StreamableParser(self._encoding, None, strict=False)
|
|
self._prime_parser(self._parser)
|
|
self._in_think_tag = False
|
|
self._prev_channel = None
|
|
self._passthrough_mode = False
|
|
|
|
@property
|
|
def current_channel(self) -> str | None:
|
|
"""Get current channel name."""
|
|
return self._parser.current_channel
|
|
|
|
@property
|
|
def current_recipient(self) -> str | None:
|
|
"""Get current recipient (for tool calls)."""
|
|
return self._parser.current_recipient
|
|
|
|
|
|
def parse_tool_calls_from_tokens(
|
|
token_ids: list[int],
|
|
prepend_start: bool = True,
|
|
) -> tuple[str, str, list[dict[str, str]]]:
|
|
"""
|
|
Parse a complete Harmony token sequence (non-streaming).
|
|
|
|
Args:
|
|
token_ids: Model output token ID list
|
|
prepend_start: Whether to prepend "<|start|>assistant" tokens.
|
|
Set to False if token_ids already includes start tokens.
|
|
|
|
Returns:
|
|
(output_text, analysis_text, tool_calls)
|
|
- output_text: Text from the final channel
|
|
- analysis_text: Chain-of-thought text from the analysis channel
|
|
- tool_calls: [{"name": "...", "arguments": "..."}]
|
|
"""
|
|
if not token_ids:
|
|
return "", "", []
|
|
|
|
try:
|
|
encoding = load_harmony_gpt_oss_encoding()
|
|
|
|
start_tokens = encoding.encode("<|start|>assistant", allowed_special="all")
|
|
has_start = list(token_ids[: len(start_tokens)]) == start_tokens
|
|
|
|
# The normal chat template includes "<|start|>assistant" in the prompt,
|
|
# so completions start from "<|channel|>" and need the prefix restored.
|
|
# Budget-forced or recovered Harmony completions may already include it.
|
|
if prepend_start and not has_start:
|
|
full_token_ids = start_tokens + list(token_ids)
|
|
else:
|
|
full_token_ids = list(token_ids)
|
|
|
|
# Decode tokens for debugging
|
|
decoded_text = encoding.decode(full_token_ids)
|
|
logger.info(f"parse_tool_calls input ({len(full_token_ids)} tokens): {decoded_text[:300]}...")
|
|
|
|
messages = encoding.parse_messages_from_completion_tokens(
|
|
full_token_ids,
|
|
role=Role.ASSISTANT,
|
|
strict=False,
|
|
)
|
|
|
|
logger.info(f"Parsed {len(messages)} messages")
|
|
for i, msg in enumerate(messages):
|
|
content_count = len(msg.content) if msg.content else 0
|
|
logger.info(
|
|
f"Message {i}: channel={msg.channel}, recipient={msg.recipient}, "
|
|
f"content_count={content_count}"
|
|
)
|
|
|
|
output_text = ""
|
|
analysis_text = ""
|
|
tool_calls = []
|
|
|
|
for msg in messages:
|
|
# Safely get content
|
|
msg_content = getattr(msg, "content", None)
|
|
if msg_content is None:
|
|
continue
|
|
|
|
# Checked before the analysis branch so analysis-channel tool
|
|
# calls do not leak their arguments into reasoning text.
|
|
if _is_tool_call_message(msg):
|
|
name = msg.recipient[10:] # Remove "functions." prefix
|
|
tool_calls.append(
|
|
{"name": name, "arguments": _message_content_text(msg)}
|
|
)
|
|
|
|
elif msg.channel == "final":
|
|
# Extract text from final channel
|
|
for content in msg_content:
|
|
text = getattr(content, "text", None)
|
|
if isinstance(text, str):
|
|
output_text += text
|
|
|
|
elif msg.channel == "analysis":
|
|
# Extract chain-of-thought text from analysis channel
|
|
for content in msg_content:
|
|
text = getattr(content, "text", None)
|
|
if isinstance(text, str):
|
|
analysis_text += text
|
|
|
|
elif msg.channel != "commentary" and _has_no_real_recipient(msg.recipient):
|
|
# Preserve malformed/unknown assistant channels as visible text
|
|
# instead of returning an empty assistant message.
|
|
for content in msg_content:
|
|
text = getattr(content, "text", None)
|
|
if isinstance(text, str):
|
|
output_text += text
|
|
|
|
return output_text, analysis_text, tool_calls
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Error parsing tool calls from tokens: {e}")
|
|
return "", "", []
|