wehub-resource-sync
821 lines
28 KiB
Python
821 lines
28 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
"""Generic streamed output parser sessions.
|
||
|
||
This module provides a tiny scheduler-facing abstraction for protocol-specific
|
||
output parsing. A parser session owns any protocol state needed while a single
|
||
request is generating (e.g. Harmony channel parsing or Gemma 4 reasoning marker
|
||
suppression) and exposes a uniform token-by-token interface.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
from collections.abc import Callable
|
||
from dataclasses import dataclass, field
|
||
from typing import Any, Protocol
|
||
|
||
from ..utils.tokenizer import (
|
||
create_streaming_detokenizer,
|
||
is_gemma4_model,
|
||
is_harmony_model,
|
||
)
|
||
from .harmony import HarmonyStreamingParser, parse_tool_calls_from_tokens
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@dataclass
|
||
class OutputParserTokenResult:
|
||
"""Per-token parser result returned during streaming."""
|
||
|
||
stream_text: str = ""
|
||
visible_text: str = ""
|
||
is_stop: bool = False
|
||
record_token: bool | None = None
|
||
|
||
|
||
@dataclass
|
||
class OutputParserFinalizeResult:
|
||
"""Final parser result returned once a request finishes."""
|
||
|
||
stream_text: str = ""
|
||
visible_text: str = ""
|
||
output_text_prefix: str = ""
|
||
tool_calls: list[dict[str, str]] = field(default_factory=list)
|
||
finish_reason: str | None = None
|
||
|
||
|
||
class OutputParserSession(Protocol):
|
||
"""Protocol implemented by per-request output parser sessions."""
|
||
|
||
def process_token(self, token_id: int) -> OutputParserTokenResult:
|
||
"""Process one generated token."""
|
||
|
||
def finalize(self) -> OutputParserFinalizeResult:
|
||
"""Flush any buffered output when generation ends."""
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class OutputParserFactory:
|
||
"""Factory for creating per-request parser sessions."""
|
||
|
||
kind: str
|
||
create_session: Callable[[Any], OutputParserSession]
|
||
stop_token_ids: set[int] = field(default_factory=set)
|
||
thinking_start_text: str | None = None
|
||
thinking_start_output_text: str | None = None
|
||
thinking_end_text: str | None = None
|
||
thinking_end_trailing_text: str | None = None
|
||
# Marker strings that must survive special-token stripping so the
|
||
# parser session can see them in the text stream. Engines that strip
|
||
# special tokens during detokenization (e.g. the serial diffusion
|
||
# lane) preserve the token ids of these markers and let the parser
|
||
# session remove them instead.
|
||
protocol_marker_texts: tuple[str, ...] = ()
|
||
|
||
|
||
class HarmonyOutputParserSession:
|
||
"""Scheduler-facing wrapper around ``HarmonyStreamingParser``."""
|
||
|
||
def __init__(self, tokenizer: Any, model_path: str | None = None):
|
||
self._tokenizer = tokenizer
|
||
self._parser = HarmonyStreamingParser(tokenizer)
|
||
self._raw_token_ids: list[int] = []
|
||
|
||
self._detokenizer = create_streaming_detokenizer(tokenizer, model_path)
|
||
if self._detokenizer is not None:
|
||
self._detokenizer.reset()
|
||
|
||
def process_token(self, token_id: int) -> OutputParserTokenResult:
|
||
control_text, stream_token, visible_token, is_stop = self._parser.process_token(
|
||
token_id
|
||
)
|
||
self._raw_token_ids.append(token_id)
|
||
|
||
stream_text = control_text
|
||
visible_text = ""
|
||
|
||
if stream_token is not None:
|
||
if self._detokenizer is not None:
|
||
self._detokenizer.add_token(stream_token)
|
||
decoded_text = self._detokenizer.last_segment
|
||
else:
|
||
decoded_text = self._tokenizer.decode([stream_token])
|
||
|
||
stream_text += decoded_text
|
||
if visible_token is not None:
|
||
visible_text += decoded_text
|
||
elif visible_token is not None:
|
||
if self._detokenizer is not None:
|
||
self._detokenizer.add_token(visible_token)
|
||
visible_text += self._detokenizer.last_segment
|
||
else:
|
||
visible_text += self._tokenizer.decode([visible_token])
|
||
|
||
return OutputParserTokenResult(
|
||
stream_text=stream_text,
|
||
visible_text=visible_text,
|
||
is_stop=is_stop,
|
||
record_token=True,
|
||
)
|
||
|
||
def finalize(self) -> OutputParserFinalizeResult:
|
||
stream_text = self._parser.finalize()
|
||
visible_text = ""
|
||
|
||
if self._detokenizer is not None:
|
||
self._detokenizer.finalize()
|
||
final_text = self._detokenizer.last_segment
|
||
if final_text:
|
||
stream_text += final_text
|
||
if self._parser.current_channel == "final":
|
||
visible_text += final_text
|
||
|
||
_, analysis_text, tool_calls = parse_tool_calls_from_tokens(self._raw_token_ids)
|
||
finish_reason = "tool_calls" if tool_calls else None
|
||
|
||
output_text_prefix = (
|
||
f"<think>\n{analysis_text}\n</think>\n" if analysis_text else ""
|
||
)
|
||
|
||
return OutputParserFinalizeResult(
|
||
stream_text=stream_text,
|
||
visible_text=visible_text,
|
||
output_text_prefix=output_text_prefix,
|
||
tool_calls=tool_calls,
|
||
finish_reason=finish_reason,
|
||
)
|
||
|
||
|
||
def _is_cohere2_moe_model(
|
||
model_name: str,
|
||
model_config: dict[str, Any] | None = None,
|
||
) -> bool:
|
||
return model_config is not None and model_config.get("model_type") == "cohere2_moe"
|
||
|
||
|
||
_MINIMAX_M3_MODEL_TYPES = {"minimax_m3", "minimax_m3_vl"}
|
||
_MINIMAX_THINK_START = "<mm:think>"
|
||
_MINIMAX_THINK_END = "</mm:think>"
|
||
_MINIMAX_EOS_TOKEN = "[e~["
|
||
_MINIMAX_SPECIAL_TOKENS = (_MINIMAX_EOS_TOKEN, "]~b]", "]~!b[", "]!p~[", "]!d~[")
|
||
_MINIMAX_TOOL_CALL_START = "]<]minimax[>[<tool_call>"
|
||
_MINIMAX_TOOL_CALL_END = "]<]minimax[>[</tool_call>"
|
||
_DEEPSEEK_V4_TOOL_CALL_START = "<｜DSML｜tool_calls>"
|
||
_DEEPSEEK_V4_TOOL_CALL_END = "</｜DSML｜tool_calls>"
|
||
|
||
|
||
def _is_deepseek_v4_model(
|
||
model_name: str,
|
||
tokenizer: Any,
|
||
model_config: dict[str, Any] | None = None,
|
||
) -> bool:
|
||
model_type = str(model_config.get("model_type", "")) if model_config else ""
|
||
if model_type.startswith("deepseek_v4"):
|
||
return True
|
||
|
||
if (
|
||
getattr(tokenizer, "tool_call_start", None) == _DEEPSEEK_V4_TOOL_CALL_START
|
||
and getattr(tokenizer, "tool_call_end", None) == _DEEPSEEK_V4_TOOL_CALL_END
|
||
):
|
||
return True
|
||
|
||
return "deepseek-v4" in model_name.lower() or "deepseek_v4" in model_name.lower()
|
||
|
||
|
||
def _serialize_minimax_tool_arguments(arguments: Any) -> str:
|
||
if isinstance(arguments, str):
|
||
return arguments or "{}"
|
||
if arguments is None:
|
||
return "{}"
|
||
try:
|
||
return json.dumps(arguments, ensure_ascii=False, separators=(",", ":"))
|
||
except TypeError:
|
||
return str(arguments)
|
||
|
||
|
||
def _is_minimax_m3_model(
|
||
model_name: str,
|
||
model_config: dict[str, Any] | None = None,
|
||
) -> bool:
|
||
model_type = model_config.get("model_type") if model_config else None
|
||
if model_type in _MINIMAX_M3_MODEL_TYPES:
|
||
return True
|
||
lowered = model_name.lower()
|
||
return "minimax" in lowered and "m3" in lowered
|
||
|
||
|
||
class _MiniMaxM3ProtocolNormalizer:
|
||
"""Normalize MiniMax M3 protocol markers to oMLX-visible markers."""
|
||
|
||
_REPLACEMENTS = (
|
||
(_MINIMAX_THINK_START, "<think>"),
|
||
(_MINIMAX_THINK_END, "</think>"),
|
||
*tuple((token, "") for token in _MINIMAX_SPECIAL_TOKENS),
|
||
)
|
||
_MARKERS = tuple(marker for marker, _ in _REPLACEMENTS)
|
||
|
||
def __init__(self) -> None:
|
||
self._buffer = ""
|
||
|
||
@classmethod
|
||
def _replace_markers(cls, text: str) -> str:
|
||
for marker, replacement in cls._REPLACEMENTS:
|
||
text = text.replace(marker, replacement)
|
||
return text
|
||
|
||
@classmethod
|
||
def _partial_suffix_len(cls, text: str) -> int:
|
||
max_len = min(len(text), max(len(marker) for marker in cls._MARKERS) - 1)
|
||
for size in range(max_len, 0, -1):
|
||
suffix = text[-size:]
|
||
if any(marker.startswith(suffix) for marker in cls._MARKERS):
|
||
return size
|
||
return 0
|
||
|
||
def feed(self, text: str) -> str:
|
||
if not text:
|
||
return ""
|
||
|
||
self._buffer += text
|
||
keep = self._partial_suffix_len(self._buffer)
|
||
if keep:
|
||
ready = self._buffer[:-keep]
|
||
self._buffer = self._buffer[-keep:]
|
||
else:
|
||
ready = self._buffer
|
||
self._buffer = ""
|
||
return self._replace_markers(ready)
|
||
|
||
def finish(self) -> str:
|
||
text = self._replace_markers(self._buffer)
|
||
self._buffer = ""
|
||
return text
|
||
|
||
|
||
def _token_id_for_text(tokenizer: Any, text: str) -> int | None:
|
||
try:
|
||
token_id = tokenizer.convert_tokens_to_ids(text)
|
||
except (AttributeError, KeyError, TypeError, ValueError):
|
||
token_id = None
|
||
if token_id is not None and token_id != getattr(tokenizer, "unk_token_id", None):
|
||
try:
|
||
return int(token_id)
|
||
except (TypeError, ValueError):
|
||
pass
|
||
|
||
try:
|
||
token_ids = tokenizer.encode(text, add_special_tokens=False)
|
||
except TypeError:
|
||
try:
|
||
token_ids = tokenizer.encode(text)
|
||
except Exception:
|
||
return None
|
||
except Exception:
|
||
return None
|
||
|
||
if len(token_ids) == 1:
|
||
try:
|
||
return int(token_ids[0])
|
||
except (TypeError, ValueError):
|
||
return None
|
||
return None
|
||
|
||
|
||
class DeepSeekV4OutputParserSession:
|
||
"""Parser session for DeepSeek V4 DSML tool-call output.
|
||
|
||
A completed DSML tool-call block ends the assistant turn. Without a
|
||
parser-owned stop, batched decode keeps the row alive after
|
||
``</｜DSML｜tool_calls>`` and the model may emit additional or malformed
|
||
DSML fragments as visible assistant text.
|
||
"""
|
||
|
||
def __init__(self, tokenizer: Any, model_path: str | None = None):
|
||
self._tokenizer = tokenizer
|
||
self._raw_text = ""
|
||
self._stopped = False
|
||
self._detokenizer = create_streaming_detokenizer(tokenizer, model_path)
|
||
if self._detokenizer is not None:
|
||
self._detokenizer.reset()
|
||
|
||
try:
|
||
from ..api.tool_calling import ToolCallStreamFilter
|
||
|
||
self._stream_filter = ToolCallStreamFilter(tokenizer)
|
||
self._visible_filter = ToolCallStreamFilter(tokenizer)
|
||
except Exception as e: # noqa: BLE001
|
||
logger.debug("DeepSeek V4 stream filter unavailable: %s", e)
|
||
self._stream_filter = None
|
||
self._visible_filter = None
|
||
|
||
def _decode_token(self, token_id: int) -> str:
|
||
if self._detokenizer is not None:
|
||
self._detokenizer.add_token(token_id)
|
||
return self._detokenizer.last_segment
|
||
try:
|
||
return self._tokenizer.decode([token_id], skip_special_tokens=False)
|
||
except TypeError:
|
||
return self._tokenizer.decode([token_id])
|
||
|
||
def _filtered_text(self, text: str, tool_filter: Any) -> str:
|
||
if not text:
|
||
return ""
|
||
if tool_filter is not None:
|
||
return tool_filter.feed(text)
|
||
return text
|
||
|
||
def _finish_filtered_text(self, tool_filter: Any) -> str:
|
||
if tool_filter is None:
|
||
return ""
|
||
return tool_filter.finish()
|
||
|
||
def _trim_at_first_tool_block_end(self, text: str) -> tuple[str, bool]:
|
||
start_idx = text.find(_DEEPSEEK_V4_TOOL_CALL_START)
|
||
if start_idx < 0:
|
||
return text, False
|
||
end_idx = text.find(_DEEPSEEK_V4_TOOL_CALL_END, start_idx)
|
||
if end_idx < 0:
|
||
return text, False
|
||
cutoff = end_idx + len(_DEEPSEEK_V4_TOOL_CALL_END)
|
||
return text[:cutoff], True
|
||
|
||
def process_token(self, token_id: int) -> OutputParserTokenResult:
|
||
if self._stopped:
|
||
return OutputParserTokenResult(is_stop=True, record_token=False)
|
||
|
||
decoded_text = self._decode_token(token_id)
|
||
combined = self._raw_text + decoded_text
|
||
trimmed, is_stop = self._trim_at_first_tool_block_end(combined)
|
||
|
||
feed_text = trimmed[len(self._raw_text) :]
|
||
self._raw_text = trimmed
|
||
self._stopped = is_stop
|
||
|
||
return OutputParserTokenResult(
|
||
stream_text=self._filtered_text(feed_text, self._stream_filter),
|
||
visible_text=self._filtered_text(feed_text, self._visible_filter),
|
||
is_stop=is_stop,
|
||
record_token=True,
|
||
)
|
||
|
||
def finalize(self) -> OutputParserFinalizeResult:
|
||
stream_text = ""
|
||
visible_text = ""
|
||
if self._detokenizer is not None and not self._stopped:
|
||
self._detokenizer.finalize()
|
||
final_text = self._detokenizer.last_segment
|
||
if final_text:
|
||
prev_len = len(self._raw_text)
|
||
combined = self._raw_text + final_text
|
||
self._raw_text, self._stopped = self._trim_at_first_tool_block_end(
|
||
combined
|
||
)
|
||
final_text = self._raw_text[prev_len:]
|
||
stream_text += self._filtered_text(final_text, self._stream_filter)
|
||
visible_text += self._filtered_text(final_text, self._visible_filter)
|
||
|
||
stream_text += self._finish_filtered_text(self._stream_filter)
|
||
visible_text += self._finish_filtered_text(self._visible_filter)
|
||
|
||
tool_calls: list[dict[str, str]] = []
|
||
try:
|
||
from ..api.tool_calling import parse_tool_calls
|
||
|
||
_, parsed_calls = parse_tool_calls(self._raw_text, self._tokenizer)
|
||
for call in parsed_calls or []:
|
||
tool_calls.append(
|
||
{
|
||
"id": getattr(call, "id", ""),
|
||
"name": call.function.name,
|
||
"arguments": call.function.arguments,
|
||
}
|
||
)
|
||
except Exception as e: # noqa: BLE001
|
||
logger.debug("DeepSeek V4 tool-call parse failed: %s", e)
|
||
|
||
return OutputParserFinalizeResult(
|
||
stream_text=stream_text,
|
||
visible_text=visible_text,
|
||
tool_calls=tool_calls,
|
||
finish_reason="tool_calls" if tool_calls else None,
|
||
)
|
||
|
||
|
||
class MiniMaxM3OutputParserSession:
|
||
"""Parser session for MiniMax M3 XML-style tool calls."""
|
||
|
||
def __init__(self, tokenizer: Any, model_path: str | None = None):
|
||
self._tokenizer = tokenizer
|
||
self._raw_text = ""
|
||
self._detokenizer = create_streaming_detokenizer(tokenizer, model_path)
|
||
if self._detokenizer is not None:
|
||
self._detokenizer.reset()
|
||
|
||
try:
|
||
from ..api.tool_calling import ToolCallStreamFilter
|
||
|
||
self._stream_filter = ToolCallStreamFilter(tokenizer)
|
||
self._visible_filter = ToolCallStreamFilter(tokenizer)
|
||
except Exception as e: # noqa: BLE001
|
||
logger.debug("MiniMax M3 stream filter unavailable: %s", e)
|
||
self._stream_filter = None
|
||
self._visible_filter = None
|
||
self._stream_normalizer = _MiniMaxM3ProtocolNormalizer()
|
||
self._visible_normalizer = _MiniMaxM3ProtocolNormalizer()
|
||
|
||
def _decode_token(self, token_id: int) -> str:
|
||
if self._detokenizer is not None:
|
||
self._detokenizer.add_token(token_id)
|
||
return self._detokenizer.last_segment
|
||
try:
|
||
return self._tokenizer.decode([token_id], skip_special_tokens=False)
|
||
except TypeError:
|
||
return self._tokenizer.decode([token_id])
|
||
|
||
def _filtered_text(
|
||
self,
|
||
text: str,
|
||
tool_filter: Any,
|
||
normalizer: _MiniMaxM3ProtocolNormalizer,
|
||
) -> str:
|
||
if not text:
|
||
return ""
|
||
if tool_filter is not None:
|
||
text = tool_filter.feed(text)
|
||
return normalizer.feed(text)
|
||
|
||
def _finish_filtered_text(
|
||
self,
|
||
tool_filter: Any,
|
||
normalizer: _MiniMaxM3ProtocolNormalizer,
|
||
) -> str:
|
||
text = ""
|
||
if tool_filter is not None:
|
||
text += normalizer.feed(tool_filter.finish())
|
||
text += normalizer.finish()
|
||
return text
|
||
|
||
def process_token(self, token_id: int) -> OutputParserTokenResult:
|
||
decoded_text = self._decode_token(token_id)
|
||
self._raw_text += decoded_text
|
||
is_stop = decoded_text == _MINIMAX_EOS_TOKEN
|
||
return OutputParserTokenResult(
|
||
stream_text=self._filtered_text(
|
||
decoded_text,
|
||
self._stream_filter,
|
||
self._stream_normalizer,
|
||
),
|
||
visible_text=self._filtered_text(
|
||
decoded_text,
|
||
self._visible_filter,
|
||
self._visible_normalizer,
|
||
),
|
||
is_stop=is_stop,
|
||
record_token=not is_stop,
|
||
)
|
||
|
||
def finalize(self) -> OutputParserFinalizeResult:
|
||
stream_text = ""
|
||
visible_text = ""
|
||
if self._detokenizer is not None:
|
||
self._detokenizer.finalize()
|
||
final_text = self._detokenizer.last_segment
|
||
if final_text:
|
||
self._raw_text += final_text
|
||
stream_text += self._filtered_text(
|
||
final_text,
|
||
self._stream_filter,
|
||
self._stream_normalizer,
|
||
)
|
||
visible_text += self._filtered_text(
|
||
final_text,
|
||
self._visible_filter,
|
||
self._visible_normalizer,
|
||
)
|
||
|
||
stream_text += self._finish_filtered_text(
|
||
self._stream_filter,
|
||
self._stream_normalizer,
|
||
)
|
||
visible_text += self._finish_filtered_text(
|
||
self._visible_filter,
|
||
self._visible_normalizer,
|
||
)
|
||
|
||
tool_calls: list[dict[str, str]] = []
|
||
if _MINIMAX_TOOL_CALL_START in self._raw_text:
|
||
try:
|
||
from ..patches.mlx_vlm_minimax_m3_compat import (
|
||
apply_mlx_vlm_minimax_m3_compat_patch,
|
||
)
|
||
|
||
apply_mlx_vlm_minimax_m3_compat_patch()
|
||
|
||
from mlx_vlm.tool_parsers.minimax_m3 import parse_tool_call
|
||
|
||
parsed = parse_tool_call(self._raw_text)
|
||
parsed_calls = parsed if isinstance(parsed, list) else [parsed]
|
||
tool_calls = [
|
||
{
|
||
"name": str(call.get("name", "")),
|
||
"arguments": _serialize_minimax_tool_arguments(
|
||
call.get("arguments")
|
||
),
|
||
}
|
||
for call in parsed_calls
|
||
if isinstance(call, dict) and call.get("name")
|
||
]
|
||
except Exception as e: # noqa: BLE001
|
||
logger.debug("MiniMax M3 tool-call parse failed: %s", e)
|
||
|
||
return OutputParserFinalizeResult(
|
||
stream_text=stream_text,
|
||
visible_text=visible_text,
|
||
tool_calls=tool_calls,
|
||
finish_reason="tool_calls" if tool_calls else None,
|
||
)
|
||
|
||
|
||
def _create_cohere2_moe_filter():
|
||
try:
|
||
from cohere_melody import PyFilter, PyFilterOptions
|
||
except ImportError:
|
||
return None
|
||
|
||
return PyFilter(PyFilterOptions().cmd4().stream_tool_actions())
|
||
|
||
|
||
def _reserialize_cohere_tool_arguments(args: str) -> str:
|
||
if not args:
|
||
return "{}"
|
||
try:
|
||
return json.dumps(
|
||
json.loads(args, strict=False),
|
||
ensure_ascii=False,
|
||
separators=(",", ":"),
|
||
)
|
||
except (json.JSONDecodeError, ValueError):
|
||
return args or "{}"
|
||
|
||
|
||
class Cohere2MoeOutputParserSession:
|
||
"""Parser session for Cohere2 MoE / Command-style Melody output."""
|
||
|
||
def __init__(self, tokenizer: Any, model_path: str | None = None):
|
||
self._tokenizer = tokenizer
|
||
self._melody = _create_cohere2_moe_filter()
|
||
if self._melody is None:
|
||
raise RuntimeError("cohere_melody is not installed")
|
||
|
||
self._detokenizer = create_streaming_detokenizer(tokenizer, model_path)
|
||
if self._detokenizer is not None:
|
||
self._detokenizer.reset()
|
||
|
||
self._thinking_started = False
|
||
self._thinking_closed = False
|
||
self._tool_calls: dict[int, dict[str, str]] = {}
|
||
|
||
def _decode_token(self, token_id: int) -> str:
|
||
if self._detokenizer is not None:
|
||
self._detokenizer.add_token(token_id)
|
||
return self._detokenizer.last_segment
|
||
try:
|
||
return self._tokenizer.decode([token_id], skip_special_tokens=False)
|
||
except TypeError:
|
||
return self._tokenizer.decode([token_id])
|
||
|
||
def _accumulate_tool_calls(self, tool_calls: list[Any]) -> None:
|
||
for tool_call in tool_calls:
|
||
index = int(getattr(tool_call, "index", 0) or 0)
|
||
current = self._tool_calls.setdefault(
|
||
index,
|
||
{"id": "", "name": "", "arguments": ""},
|
||
)
|
||
current["id"] += getattr(tool_call, "id", "") or ""
|
||
current["name"] += getattr(tool_call, "name", "") or ""
|
||
current["arguments"] += getattr(tool_call, "arguments", "") or ""
|
||
|
||
def _apply_melody_result(self, result: Any) -> tuple[str, str]:
|
||
stream_text = ""
|
||
visible_text = ""
|
||
|
||
reasoning = getattr(result, "reasoning", None)
|
||
if reasoning:
|
||
if not self._thinking_started:
|
||
self._thinking_started = True
|
||
stream_text += "<think>\n"
|
||
visible_text += "<think>\n"
|
||
stream_text += reasoning
|
||
visible_text += reasoning
|
||
|
||
content = getattr(result, "content", None)
|
||
if content:
|
||
if self._thinking_started and not self._thinking_closed:
|
||
self._thinking_closed = True
|
||
stream_text += "</think>\n"
|
||
visible_text += "</think>\n"
|
||
stream_text += content
|
||
visible_text += content
|
||
|
||
self._accumulate_tool_calls(getattr(result, "tool_calls", []) or [])
|
||
return stream_text, visible_text
|
||
|
||
def process_token(self, token_id: int) -> OutputParserTokenResult:
|
||
decoded_text = self._decode_token(token_id)
|
||
if not decoded_text:
|
||
return OutputParserTokenResult(record_token=True)
|
||
|
||
result = self._melody.write_decoded(decoded_text)
|
||
stream_text, visible_text = self._apply_melody_result(result)
|
||
return OutputParserTokenResult(
|
||
stream_text=stream_text,
|
||
visible_text=visible_text,
|
||
record_token=True,
|
||
)
|
||
|
||
def finalize(self) -> OutputParserFinalizeResult:
|
||
stream_text = ""
|
||
visible_text = ""
|
||
|
||
if self._detokenizer is not None:
|
||
self._detokenizer.finalize()
|
||
final_text = self._detokenizer.last_segment
|
||
if final_text:
|
||
result = self._melody.write_decoded(final_text)
|
||
s_text, v_text = self._apply_melody_result(result)
|
||
stream_text += s_text
|
||
visible_text += v_text
|
||
|
||
result = self._melody.flush_partials()
|
||
s_text, v_text = self._apply_melody_result(result)
|
||
stream_text += s_text
|
||
visible_text += v_text
|
||
|
||
if self._thinking_started and not self._thinking_closed:
|
||
self._thinking_closed = True
|
||
stream_text += "</think>\n"
|
||
visible_text += "</think>\n"
|
||
|
||
tool_calls = [
|
||
{
|
||
"id": value["id"],
|
||
"name": value["name"],
|
||
"arguments": _reserialize_cohere_tool_arguments(value["arguments"]),
|
||
}
|
||
for _, value in sorted(self._tool_calls.items())
|
||
if value["name"]
|
||
]
|
||
|
||
return OutputParserFinalizeResult(
|
||
stream_text=stream_text,
|
||
visible_text=visible_text,
|
||
tool_calls=tool_calls,
|
||
finish_reason="tool_calls" if tool_calls else None,
|
||
)
|
||
|
||
|
||
def detect_output_parser(
|
||
model_name: str,
|
||
tokenizer: Any,
|
||
model_config: dict[str, Any] | None = None,
|
||
model_path: str | None = None,
|
||
) -> OutputParserFactory | None:
|
||
"""Detect a protocol-specific output parser for the model, if needed.
|
||
|
||
``model_name`` drives detection (string matching) and may be a display
|
||
id rather than a directory since #2178. Pass ``model_path`` when the
|
||
filesystem path is available so parser sessions can locate
|
||
tokenizer.json for their streaming detokenizers.
|
||
"""
|
||
session_model_path = model_path or model_name
|
||
|
||
if is_harmony_model(model_name, model_config):
|
||
temp_parser = HarmonyStreamingParser(tokenizer)
|
||
return OutputParserFactory(
|
||
kind="harmony",
|
||
create_session=lambda session_tokenizer: HarmonyOutputParserSession(
|
||
session_tokenizer,
|
||
model_path=session_model_path,
|
||
),
|
||
stop_token_ids=temp_parser.get_stop_token_ids(),
|
||
thinking_end_text="<|end|>",
|
||
thinking_end_trailing_text="<|start|>assistant<|channel|>final<|message|>",
|
||
)
|
||
|
||
if is_gemma4_model(model_name, model_config):
|
||
from .gemma4 import (
|
||
_CLOSE_MARKER,
|
||
_OPEN_MARKER_BARE,
|
||
_TOOL_RESPONSE_CLOSE,
|
||
_TOOL_RESPONSE_OPEN,
|
||
_TURN_END_MARKER,
|
||
Gemma4OutputParserSession,
|
||
)
|
||
|
||
return OutputParserFactory(
|
||
kind="gemma4",
|
||
create_session=lambda session_tokenizer: Gemma4OutputParserSession(
|
||
session_tokenizer,
|
||
model_path=session_model_path,
|
||
),
|
||
stop_token_ids=set(),
|
||
thinking_end_text="<channel|>",
|
||
protocol_marker_texts=(
|
||
_OPEN_MARKER_BARE,
|
||
_CLOSE_MARKER,
|
||
_TURN_END_MARKER,
|
||
_TOOL_RESPONSE_OPEN,
|
||
_TOOL_RESPONSE_CLOSE,
|
||
),
|
||
)
|
||
|
||
if _is_deepseek_v4_model(model_name, tokenizer, model_config):
|
||
return OutputParserFactory(
|
||
kind="deepseek_v4",
|
||
create_session=lambda session_tokenizer: DeepSeekV4OutputParserSession(
|
||
session_tokenizer,
|
||
model_path=session_model_path,
|
||
),
|
||
stop_token_ids=set(),
|
||
protocol_marker_texts=(
|
||
_DEEPSEEK_V4_TOOL_CALL_START,
|
||
_DEEPSEEK_V4_TOOL_CALL_END,
|
||
),
|
||
)
|
||
|
||
if _is_cohere2_moe_model(model_name, model_config):
|
||
if _create_cohere2_moe_filter() is None:
|
||
logger.warning(
|
||
"cohere_melody is not installed; Cohere2 MoE output parser "
|
||
"is disabled for %s",
|
||
model_name,
|
||
)
|
||
return None
|
||
|
||
return OutputParserFactory(
|
||
kind="cohere2_moe",
|
||
create_session=lambda session_tokenizer: Cohere2MoeOutputParserSession(
|
||
session_tokenizer,
|
||
model_path=session_model_path,
|
||
),
|
||
stop_token_ids=set(),
|
||
thinking_end_text="</think>",
|
||
)
|
||
|
||
if _is_minimax_m3_model(model_name, model_config):
|
||
minimax_stop_ids = set()
|
||
eos_id = _token_id_for_text(tokenizer, _MINIMAX_EOS_TOKEN)
|
||
if eos_id is not None:
|
||
minimax_stop_ids.add(eos_id)
|
||
|
||
return OutputParserFactory(
|
||
kind="minimax_m3",
|
||
create_session=lambda session_tokenizer: MiniMaxM3OutputParserSession(
|
||
session_tokenizer,
|
||
model_path=session_model_path,
|
||
),
|
||
stop_token_ids=minimax_stop_ids,
|
||
thinking_start_text=_MINIMAX_THINK_START,
|
||
thinking_start_output_text="<think>\n",
|
||
thinking_end_text=_MINIMAX_THINK_END,
|
||
protocol_marker_texts=(
|
||
_MINIMAX_THINK_START,
|
||
_MINIMAX_THINK_END,
|
||
_MINIMAX_TOOL_CALL_START,
|
||
_MINIMAX_TOOL_CALL_END,
|
||
),
|
||
)
|
||
|
||
return None
|
||
|
||
|
||
def detect_message_extractor(
|
||
model_name: str,
|
||
model_config: dict[str, Any] | None = None,
|
||
) -> Callable:
|
||
"""Return the appropriate message extractor function for the model.
|
||
|
||
The returned callable has the signature::
|
||
|
||
extractor(messages, max_tool_result_tokens=None, tokenizer=None) -> list[dict]
|
||
|
||
This mirrors how ``detect_output_parser`` decouples model-specific
|
||
knowledge from the server layer — the engine stores the extractor at
|
||
load time and the server just calls ``engine.message_extractor(...)``.
|
||
"""
|
||
if is_harmony_model(model_name, model_config):
|
||
from ..api.utils import extract_harmony_messages
|
||
|
||
return extract_harmony_messages
|
||
|
||
if is_gemma4_model(model_name, model_config):
|
||
from .gemma4 import extract_gemma4_messages
|
||
|
||
return extract_gemma4_messages
|
||
|
||
# Default: caller decides between extract_text_content and
|
||
# extract_multimodal_content based on engine type (VLM vs text).
|
||
return None
|