jundot--omlx
wehub-resource-sync
500 行
20 KiB
Python
500 行
20 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Gemma 4 reasoning-channel output parsing and message extraction."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
from typing import Any
|
|
|
|
from ..api.utils import _PRESERVE_BOUNDARY_KEY
|
|
from ..utils.tokenizer import create_streaming_detokenizer
|
|
from .output_parser import OutputParserFinalizeResult, OutputParserTokenResult
|
|
|
|
_OPEN_MARKER = "<|channel>thought\n"
|
|
_OPEN_MARKER_BARE = "<|channel>"
|
|
_CLOSE_MARKER = "<channel|>"
|
|
_TURN_END_MARKER = "<turn|>"
|
|
_TOOL_RESPONSE_OPEN = "<|tool_response>"
|
|
_TOOL_RESPONSE_CLOSE = "<tool_response|>"
|
|
_THINK_OPEN = "<think>\n"
|
|
_THINK_CLOSE = "</think>\n"
|
|
|
|
_LEADING_THOUGHT_RE = re.compile(
|
|
r"\A\s*(?:(?:<think>.*?</think>|<\|channel>.*?<channel\|>)\s*)+",
|
|
re.DOTALL,
|
|
)
|
|
|
|
# Matches the STRAY bare-token spellings (<|tool_call> and <tool_call|>),
|
|
# not the template's well-formed closing form (</tool_call|> with slash).
|
|
_PROTOCOL_MARKER_RE = re.compile(r"<\|tool_call>|<tool_call\|>")
|
|
|
|
|
|
def _strip_protocol_markers(text: Any) -> Any:
|
|
"""Remove stray <|tool_call> / <tool_call|> tokens from assistant content."""
|
|
if not isinstance(text, str) or not text:
|
|
return text
|
|
return _PROTOCOL_MARKER_RE.sub("", text)
|
|
|
|
|
|
def _try_parse_json(s: str) -> Any:
|
|
"""Parse string as JSON if possible, otherwise return as-is."""
|
|
if not isinstance(s, str):
|
|
return s
|
|
s = s.strip()
|
|
if not s or not (s.startswith("{") or s.startswith("[")):
|
|
return s
|
|
try:
|
|
return json.loads(s)
|
|
except (json.JSONDecodeError, ValueError):
|
|
return s
|
|
|
|
|
|
def _strip_thinking(text: Any) -> Any:
|
|
"""Remove leading ``<think>...</think>`` or raw ``<|channel>...<channel|>`` spans.
|
|
|
|
Gemma 4's multi-turn rule requires that only the final visible answer
|
|
is kept in chat history. Clients such as Open WebUI replay the full
|
|
assistant content (including the rendered ``<think>`` block, or the
|
|
raw protocol form when a client preserves it). Feeding prior thought
|
|
blocks back primes the model to emit malformed channel markers on the
|
|
next turn, which then leak into user-facing output.
|
|
|
|
The match is anchored to the start of the message: the rendered thought
|
|
block always precedes the visible answer, so this catches every
|
|
legitimate occurrence while leaving inline mentions (e.g. an assistant
|
|
explaining how ``<think>`` tags work) untouched.
|
|
"""
|
|
if not isinstance(text, str) or not text:
|
|
return text
|
|
return _LEADING_THOUGHT_RE.sub("", text, count=1)
|
|
|
|
|
|
def extract_gemma4_messages(
|
|
messages: list[Any],
|
|
max_tool_result_tokens: int | None = None,
|
|
tokenizer: Any | None = None,
|
|
consolidate_system_messages: bool = True,
|
|
) -> list[dict]:
|
|
"""Convert OpenAI-format messages to Gemma 4 chat-template format.
|
|
|
|
The Gemma 4 chat template does not handle ``role=tool`` messages.
|
|
Tool results must instead appear on a model-role turn as a
|
|
``tool_responses`` list, where each entry is::
|
|
|
|
{"name": "<function_name>", "response": <dict_or_scalar>}
|
|
|
|
This function:
|
|
- Passes non-tool messages through unchanged.
|
|
- Preserves ``tool_calls`` on assistant turns (template renders them
|
|
as ``<|tool_call>...</tool_call|>``).
|
|
- Folds consecutive ``role=tool`` messages that follow an assistant
|
|
turn into a single ``{"role": "assistant", "tool_responses": [...]}``
|
|
message, resolving function names from the preceding tool_calls by
|
|
``tool_call_id``. Falls back to the raw ``tool_call_id`` as the
|
|
name when no match is found.
|
|
- JSON-parses tool result content into a dict/list where possible so
|
|
the template renders structured responses correctly.
|
|
|
|
Args:
|
|
messages: OpenAI-format Message objects or dicts.
|
|
max_tool_result_tokens: Maximum token count for tool results
|
|
(truncation applied when tokenizer is provided).
|
|
tokenizer: Tokenizer for optional truncation.
|
|
consolidate_system_messages: When True, preserve the legacy behavior
|
|
of moving all system/developer messages to the leading system
|
|
prompt. Server routes pass False and defer that decision until the
|
|
model chat template can be probed.
|
|
|
|
Returns:
|
|
List of dicts ready for ``tokenizer.apply_chat_template``.
|
|
"""
|
|
from ..api.utils import (
|
|
_extract_text_from_content_list,
|
|
) # avoid circular at module level
|
|
|
|
processed: list[dict] = []
|
|
|
|
# Build index of message objects as plain dicts
|
|
raw: list[dict] = []
|
|
for msg in messages:
|
|
if hasattr(msg, "model_dump"):
|
|
raw.append(msg.model_dump())
|
|
elif isinstance(msg, dict):
|
|
raw.append(dict(msg))
|
|
else:
|
|
raw.append(
|
|
{
|
|
"role": getattr(msg, "role", "user"),
|
|
"content": getattr(msg, "content", ""),
|
|
}
|
|
)
|
|
|
|
i = 0
|
|
while i < len(raw):
|
|
msg = raw[i]
|
|
role = msg.get("role", "user")
|
|
|
|
if role == "developer":
|
|
role = "system"
|
|
|
|
if role == "tool":
|
|
# Orphaned tool result with no preceding assistant turn — attach
|
|
# to a synthetic assistant turn with no content.
|
|
tool_call_id = msg.get("tool_call_id", "")
|
|
content = msg.get("content", "")
|
|
if isinstance(content, list):
|
|
content = _extract_text_from_content_list(content)
|
|
if max_tool_result_tokens and tokenizer and content:
|
|
from ..api.anthropic_utils import truncate_tool_result
|
|
|
|
content = truncate_tool_result(
|
|
content, max_tool_result_tokens, tokenizer
|
|
)
|
|
response = _try_parse_json(content)
|
|
# Fallback name: use tool_call_id
|
|
processed.append(
|
|
{
|
|
"role": "assistant",
|
|
"content": "",
|
|
"tool_responses": [
|
|
{"name": tool_call_id or "unknown", "response": response}
|
|
],
|
|
_PRESERVE_BOUNDARY_KEY: True,
|
|
}
|
|
)
|
|
i += 1
|
|
continue
|
|
|
|
if role == "assistant":
|
|
# Build a tool_call_id → function_name lookup from this turn's calls.
|
|
tc_id_to_name: dict[str, str] = {}
|
|
tool_calls_raw = msg.get("tool_calls") or []
|
|
for tc in tool_calls_raw:
|
|
if isinstance(tc, dict):
|
|
tc_id = tc.get("id", "")
|
|
func_name = (tc.get("function") or {}).get("name", "")
|
|
else:
|
|
tc_id = getattr(tc, "id", "")
|
|
func = getattr(tc, "function", None)
|
|
func_name = getattr(func, "name", "") if func else ""
|
|
if tc_id:
|
|
tc_id_to_name[tc_id] = func_name
|
|
|
|
# Extract content
|
|
content = msg.get("content", "")
|
|
if isinstance(content, list):
|
|
content = _extract_text_from_content_list(content)
|
|
# Per Gemma 4's multi-turn rule, prior thought blocks must not
|
|
# be fed back into the next turn. Strip them before rendering.
|
|
content = _strip_thinking(content)
|
|
content = _strip_protocol_markers(content)
|
|
|
|
out_msg: dict = {"role": "assistant", "content": content or ""}
|
|
|
|
# Preserve tool_calls for template rendering
|
|
if tool_calls_raw:
|
|
out_calls = []
|
|
for tc in tool_calls_raw:
|
|
if isinstance(tc, dict):
|
|
func = tc.get("function") or {}
|
|
out_calls.append(
|
|
{
|
|
"id": tc.get("id", ""),
|
|
"function": {
|
|
"name": func.get("name", ""),
|
|
"arguments": _try_parse_json(
|
|
func.get("arguments", "{}")
|
|
),
|
|
},
|
|
}
|
|
)
|
|
else:
|
|
func = getattr(tc, "function", None)
|
|
args_str = getattr(func, "arguments", "{}") if func else "{}"
|
|
out_calls.append(
|
|
{
|
|
"id": getattr(tc, "id", ""),
|
|
"function": {
|
|
"name": getattr(func, "name", "") if func else "",
|
|
"arguments": _try_parse_json(args_str),
|
|
},
|
|
}
|
|
)
|
|
out_msg["tool_calls"] = out_calls
|
|
out_msg[_PRESERVE_BOUNDARY_KEY] = True
|
|
|
|
processed.append(out_msg)
|
|
i += 1
|
|
|
|
# Consume any immediately following tool results into a
|
|
# single model turn with tool_responses.
|
|
tool_responses = []
|
|
while i < len(raw) and raw[i].get("role") == "tool":
|
|
tr = raw[i]
|
|
tc_id = tr.get("tool_call_id", "")
|
|
tr_content = tr.get("content", "")
|
|
if isinstance(tr_content, list):
|
|
tr_content = _extract_text_from_content_list(tr_content)
|
|
if max_tool_result_tokens and tokenizer and tr_content:
|
|
from ..api.anthropic_utils import truncate_tool_result
|
|
|
|
tr_content = truncate_tool_result(
|
|
tr_content, max_tool_result_tokens, tokenizer
|
|
)
|
|
response = _try_parse_json(tr_content)
|
|
name = tc_id_to_name.get(tc_id) or tc_id or "unknown"
|
|
tool_responses.append({"name": name, "response": response})
|
|
i += 1
|
|
|
|
if tool_responses:
|
|
# Attach tool_responses to the SAME assistant message that
|
|
# has tool_calls. The Gemma 4 chat template checks for
|
|
# tool_responses on the current message (lines 261-267)
|
|
# BEFORE falling back to a forward-scan for role='tool'
|
|
# messages (lines 268-302). Putting them on a separate
|
|
# assistant message causes both paths to miss, producing a
|
|
# corrupt bare <|tool_response> tag and making the model
|
|
# loop on the same tool call.
|
|
out_msg["tool_responses"] = tool_responses
|
|
continue
|
|
|
|
# All other roles (user, system)
|
|
# Preserve image_url and input_audio parts for VLM processing
|
|
content = msg.get("content", "")
|
|
if isinstance(content, list):
|
|
from ..api.utils import _extract_multimodal_content_list
|
|
|
|
multimodal_parts = _extract_multimodal_content_list(content)
|
|
multimodal_types = {"image_url", "input_audio"}
|
|
has_multimodal = any(
|
|
p.get("type") in multimodal_types for p in multimodal_parts
|
|
)
|
|
if has_multimodal:
|
|
content = multimodal_parts
|
|
else:
|
|
content = _extract_text_from_content_list(content)
|
|
out: dict = {"role": role, "content": content if content is not None else ""}
|
|
processed.append(out)
|
|
i += 1
|
|
|
|
# Standard cleanup passes shared with other extractors
|
|
from ..api.utils import (
|
|
_consolidate_system_messages,
|
|
_drop_void_assistant_messages,
|
|
_merge_consecutive_roles,
|
|
)
|
|
|
|
cleaned = processed
|
|
if consolidate_system_messages:
|
|
cleaned = _consolidate_system_messages(cleaned)
|
|
cleaned = _drop_void_assistant_messages(cleaned)
|
|
return _merge_consecutive_roles(cleaned)
|
|
|
|
|
|
def _matching_prefix_len(text: str, marker: str) -> int:
|
|
"""Return longest suffix of ``text`` that is a prefix of ``marker``."""
|
|
max_len = min(len(text), len(marker) - 1)
|
|
for size in range(max_len, 0, -1):
|
|
if text.endswith(marker[:size]):
|
|
return size
|
|
return 0
|
|
|
|
|
|
class Gemma4OutputParserSession:
|
|
"""Suppress Gemma 4 protocol markers and re-emit thought blocks as ``<think>`` tags."""
|
|
|
|
def __init__(self, tokenizer: Any, model_path: str | None = None):
|
|
self._tokenizer = tokenizer
|
|
self._buffer = ""
|
|
self._in_thought = False
|
|
self._text_mode = False
|
|
|
|
self._detokenizer = create_streaming_detokenizer(tokenizer, model_path)
|
|
if self._detokenizer is not None:
|
|
self._detokenizer.reset()
|
|
|
|
def _append_text(
|
|
self,
|
|
stream_parts: list[str],
|
|
visible_parts: list[str],
|
|
text: str,
|
|
) -> None:
|
|
if not text:
|
|
return
|
|
stream_parts.append(text)
|
|
visible_parts.append(text)
|
|
|
|
def _active_markers(self) -> list[str]:
|
|
# Channel open/close are tracked unconditionally so a stray
|
|
# ``<channel|>`` outside a thought block (occasionally emitted in long
|
|
# multi-turn contexts) is absorbed instead of leaking into visible
|
|
# text. ``_OPEN_MARKER_BARE`` is a defensive fallback for malformed
|
|
# opens (e.g. ``<|channel>thought<channel|>`` with no newline, or a
|
|
# bare ``<|channel>`` emitted when the model is confused by polluted
|
|
# history). Tool-call markup is intentionally not tracked here — the
|
|
# downstream ``ToolCallStreamFilter`` removes it from stream deltas
|
|
# while ``parse_tool_calls`` still sees the raw markers in
|
|
# ``output_text`` for extraction.
|
|
return [
|
|
_OPEN_MARKER,
|
|
_OPEN_MARKER_BARE,
|
|
_CLOSE_MARKER,
|
|
_TURN_END_MARKER,
|
|
_TOOL_RESPONSE_OPEN,
|
|
_TOOL_RESPONSE_CLOSE,
|
|
]
|
|
|
|
@staticmethod
|
|
def _find_next_marker(
|
|
source: str, pos: int, markers: list[str]
|
|
) -> tuple[int, str] | tuple[None, None]:
|
|
next_idx: int | None = None
|
|
next_marker: str | None = None
|
|
for marker in markers:
|
|
idx = source.find(marker, pos)
|
|
if idx == -1:
|
|
continue
|
|
if next_idx is None or idx < next_idx:
|
|
next_idx = idx
|
|
next_marker = marker
|
|
return next_idx, next_marker
|
|
|
|
def _consume_text(
|
|
self, text: str, *, final: bool = False
|
|
) -> OutputParserTokenResult:
|
|
source = self._buffer + text
|
|
self._buffer = ""
|
|
|
|
stream_parts: list[str] = []
|
|
visible_parts: list[str] = []
|
|
pos = 0
|
|
|
|
while pos < len(source):
|
|
markers = self._active_markers()
|
|
idx, marker = self._find_next_marker(source, pos, markers)
|
|
|
|
if idx is None or marker is None:
|
|
remainder = source[pos:]
|
|
if not final:
|
|
keep = max(
|
|
_matching_prefix_len(remainder, marker_text)
|
|
for marker_text in markers
|
|
)
|
|
if keep:
|
|
emit = remainder[:-keep]
|
|
self._buffer = remainder[-keep:]
|
|
else:
|
|
emit = remainder
|
|
else:
|
|
emit = remainder
|
|
|
|
self._append_text(stream_parts, visible_parts, emit)
|
|
break
|
|
|
|
# Streaming defer: a bare ``<|channel>`` (or ``<|channel>thought``
|
|
# without trailing newline) at the end of the source could still
|
|
# extend to the canonical ``<|channel>thought\n`` once more
|
|
# tokens arrive. Buffer and wait so the canonical match wins.
|
|
if not final and marker == _OPEN_MARKER_BARE:
|
|
suffix = source[idx:]
|
|
if len(suffix) < len(_OPEN_MARKER) and _OPEN_MARKER.startswith(suffix):
|
|
self._append_text(stream_parts, visible_parts, source[pos:idx])
|
|
self._buffer = suffix
|
|
return OutputParserTokenResult(
|
|
stream_text="".join(stream_parts),
|
|
visible_text="".join(visible_parts),
|
|
)
|
|
|
|
self._append_text(stream_parts, visible_parts, source[pos:idx])
|
|
|
|
advance = len(marker)
|
|
|
|
if marker == _OPEN_MARKER:
|
|
# Nested open while already in a thought block: drop the stray
|
|
# marker without re-emitting ``<think>`` to keep the structure
|
|
# well-formed.
|
|
if not self._in_thought:
|
|
stream_parts.append(_THINK_OPEN)
|
|
visible_parts.append(_THINK_OPEN)
|
|
self._in_thought = True
|
|
elif marker == _OPEN_MARKER_BARE:
|
|
# Defensive fallback for malformed opens: ``<|channel>thought``
|
|
# without the trailing newline, or a bare ``<|channel>`` with
|
|
# an unrecognised channel name. Treat as a thought open and
|
|
# absorb the optional ``thought`` keyword and newline so they
|
|
# don't leak as visible text.
|
|
if not self._in_thought:
|
|
stream_parts.append(_THINK_OPEN)
|
|
visible_parts.append(_THINK_OPEN)
|
|
self._in_thought = True
|
|
after = idx + advance
|
|
if source.startswith("thought\n", after):
|
|
advance += len("thought\n")
|
|
elif source.startswith("thought", after):
|
|
advance += len("thought")
|
|
elif marker == _CLOSE_MARKER:
|
|
# Stray close outside a thought block: drop silently to keep
|
|
# the marker out of visible content.
|
|
if self._in_thought:
|
|
stream_parts.append(_THINK_CLOSE)
|
|
visible_parts.append(_THINK_CLOSE)
|
|
self._in_thought = False
|
|
# _TURN_END_MARKER, _TOOL_RESPONSE_OPEN / _CLOSE: silent drop.
|
|
|
|
pos = idx + advance
|
|
|
|
return OutputParserTokenResult(
|
|
stream_text="".join(stream_parts),
|
|
visible_text="".join(visible_parts),
|
|
)
|
|
|
|
def process_token(self, token_id: int) -> OutputParserTokenResult:
|
|
if self._detokenizer is not None:
|
|
self._detokenizer.add_token(token_id)
|
|
text = self._detokenizer.last_segment
|
|
else:
|
|
text = self._tokenizer.decode([token_id])
|
|
return self._consume_text(text)
|
|
|
|
def process_text(self, text: str) -> OutputParserTokenResult:
|
|
"""Process an already-detokenized text segment.
|
|
|
|
Engines that emit text segments instead of token ids (the serial
|
|
diffusion lane detokenizes inside ``stream_diffusion_generate``)
|
|
feed their output through this entry point so protocol markers
|
|
are handled identically to the token-id path. Switches the
|
|
session to text mode so ``finalize`` does not flush the unused
|
|
token detokenizer.
|
|
"""
|
|
self._text_mode = True
|
|
if not text:
|
|
return OutputParserTokenResult(stream_text="", visible_text="")
|
|
return self._consume_text(text)
|
|
|
|
def finalize(self) -> OutputParserFinalizeResult:
|
|
text = ""
|
|
if self._detokenizer is not None and not self._text_mode:
|
|
self._detokenizer.finalize()
|
|
text = self._detokenizer.last_segment
|
|
|
|
token_result = self._consume_text(text, final=True)
|
|
|
|
stream_text = token_result.stream_text
|
|
visible_text = token_result.visible_text
|
|
|
|
if self._buffer:
|
|
stream_text += self._buffer
|
|
visible_text += self._buffer
|
|
self._buffer = ""
|
|
|
|
if self._in_thought:
|
|
stream_text += _THINK_CLOSE
|
|
visible_text += _THINK_CLOSE
|
|
self._in_thought = False
|
|
|
|
return OutputParserFinalizeResult(
|
|
stream_text=stream_text,
|
|
visible_text=visible_text,
|
|
)
|