243 lines
9.9 KiB
Python
243 lines
9.9 KiB
Python
"""Safe helpers for provider text streams.
|
|
|
|
Some OpenAI-compatible providers place reasoning in the normal ``content``
|
|
stream as ``<think>...</think>`` instead of using LangChain's dedicated
|
|
reasoning fields. Those fragments must be separated from report/file content
|
|
while callers can still surface them in an explicitly marked live UI trace.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
|
|
class InlineThinkTagFilter:
|
|
"""Separate inline ``<think>`` blocks while preserving token streaming.
|
|
|
|
Tags can be split across arbitrary provider chunks, so regex-per-chunk is
|
|
not safe: it can leak a partial tag or the model's reasoning into the
|
|
visible draft. This small state machine holds only the suffix that might
|
|
still form a tag and emits all confirmed answer text immediately.
|
|
|
|
``leading_only=True`` reserves only one optional response-leading envelope
|
|
for reasoning. It is suitable for exact payloads such as file bodies,
|
|
where a later literal ``<think>`` must remain ordinary content.
|
|
"""
|
|
|
|
_OPEN_TAG = "<think>"
|
|
_CLOSE_TAG = "</think>"
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
allow_implicit_close: bool = False,
|
|
leading_only: bool = False,
|
|
) -> None:
|
|
self._inside_think = False
|
|
self._pending = ""
|
|
self._channel_decided = not allow_implicit_close
|
|
# File-body transports need a stricter interpretation than ordinary
|
|
# chat rendering: only one optional reasoning envelope at the start of
|
|
# the response is protocol metadata. Once file bytes begin, a literal
|
|
# ``<think>`` is user content and must never be rewritten.
|
|
self._leading_only = leading_only
|
|
self._leading_phase = (
|
|
"implicit_probe"
|
|
if leading_only and allow_implicit_close
|
|
else ("probe" if leading_only else "disabled")
|
|
)
|
|
|
|
@staticmethod
|
|
def _possible_tag_prefix_size(content: str, tag: str) -> int:
|
|
"""Return the longest suffix which may become the next tag."""
|
|
lowered = content.lower()
|
|
for size in range(min(len(content), len(tag) - 1), 0, -1):
|
|
if lowered[-size:] == tag[:size]:
|
|
return size
|
|
return 0
|
|
|
|
def push_parts(self, chunk: str) -> tuple[str, str]:
|
|
"""Return ``(visible_text, thinking_text)`` for one provider chunk."""
|
|
if not chunk:
|
|
return "", ""
|
|
|
|
if self._leading_only:
|
|
return self._push_leading_parts(chunk)
|
|
|
|
remaining = self._pending + chunk
|
|
self._pending = ""
|
|
visible: list[str] = []
|
|
thinking: list[str] = []
|
|
|
|
if not self._channel_decided:
|
|
lowered = remaining.lower()
|
|
open_at = lowered.find(self._OPEN_TAG)
|
|
close_at = lowered.find(self._CLOSE_TAG)
|
|
if close_at >= 0 and (open_at < 0 or close_at < open_at):
|
|
# Some templates inject the opening tag outside returned
|
|
# content. Keep the ambiguous prefix private until its close
|
|
# marker arrives, even when that marker spans chunks.
|
|
thinking.append(remaining[:close_at])
|
|
remaining = remaining[close_at + len(self._CLOSE_TAG) :]
|
|
self._channel_decided = True
|
|
elif open_at >= 0:
|
|
self._channel_decided = True
|
|
else:
|
|
self._pending = remaining
|
|
return "", ""
|
|
|
|
while remaining:
|
|
tag = self._CLOSE_TAG if self._inside_think else self._OPEN_TAG
|
|
marker_at = remaining.lower().find(tag)
|
|
if marker_at >= 0:
|
|
(thinking if self._inside_think else visible).append(remaining[:marker_at])
|
|
remaining = remaining[marker_at + len(tag) :]
|
|
self._inside_think = not self._inside_think
|
|
continue
|
|
|
|
suffix_size = self._possible_tag_prefix_size(remaining, tag)
|
|
confirmed = remaining[:-suffix_size] if suffix_size else remaining
|
|
(thinking if self._inside_think else visible).append(confirmed)
|
|
self._pending = remaining[-suffix_size:] if suffix_size else ""
|
|
break
|
|
|
|
return "".join(visible), "".join(thinking)
|
|
|
|
def _push_leading_parts(self, chunk: str) -> tuple[str, str]:
|
|
"""Route only a leading ``<think>`` envelope as provider reasoning.
|
|
|
|
This mode is intended for protocol payloads such as streamed file
|
|
bodies. It waits only while the response prefix could still become an
|
|
opening tag, supports tags split across arbitrary chunks, and locks to
|
|
the visible channel immediately after the closing tag or the first
|
|
confirmed file byte. Later literal tags are therefore preserved.
|
|
"""
|
|
|
|
remaining = self._pending + chunk
|
|
self._pending = ""
|
|
|
|
if self._leading_phase == "visible":
|
|
return remaining, ""
|
|
|
|
if self._leading_phase in {"probe", "implicit_probe"}:
|
|
candidate = remaining.lstrip()
|
|
lowered = candidate.lower()
|
|
if lowered.startswith(self._OPEN_TAG):
|
|
# Whitespace before a leading reasoning envelope belongs to
|
|
# the provider's presentation, not to the prepared file.
|
|
remaining = candidate[len(self._OPEN_TAG) :]
|
|
self._leading_phase = "thinking"
|
|
self._inside_think = True
|
|
elif self._OPEN_TAG.startswith(lowered):
|
|
# Empty/whitespace-only prefixes and partial tags remain
|
|
# ambiguous until the next provider chunk arrives.
|
|
self._pending = remaining
|
|
return "", ""
|
|
elif self._leading_phase == "implicit_probe":
|
|
close_at = remaining.lower().find(self._CLOSE_TAG)
|
|
if close_at >= 0:
|
|
# Some chat templates inject ``<think>`` outside the
|
|
# returned token stream. This mode is enabled only when
|
|
# the immediately preceding bridge-selection response
|
|
# demonstrated that exact close-only convention. Until
|
|
# the closing marker arrives the prefix is intentionally
|
|
# buffered: ordinary file text and implicit reasoning are
|
|
# otherwise indistinguishable.
|
|
thinking = remaining[:close_at]
|
|
visible = remaining[close_at + len(self._CLOSE_TAG) :]
|
|
self._leading_phase = "visible"
|
|
return visible, thinking
|
|
self._pending = remaining
|
|
return "", ""
|
|
else:
|
|
self._leading_phase = "visible"
|
|
return remaining, ""
|
|
|
|
marker_at = remaining.lower().find(self._CLOSE_TAG)
|
|
if marker_at >= 0:
|
|
thinking = remaining[:marker_at]
|
|
visible = remaining[marker_at + len(self._CLOSE_TAG) :]
|
|
self._leading_phase = "visible"
|
|
self._inside_think = False
|
|
return visible, thinking
|
|
|
|
suffix_size = self._possible_tag_prefix_size(remaining, self._CLOSE_TAG)
|
|
thinking = remaining[:-suffix_size] if suffix_size else remaining
|
|
self._pending = remaining[-suffix_size:] if suffix_size else ""
|
|
return "", thinking
|
|
|
|
def push(self, chunk: str) -> str:
|
|
"""Compatibility helper returning only confirmed visible text."""
|
|
return self.push_parts(chunk)[0]
|
|
|
|
def confirm_visible_channel(self) -> str:
|
|
"""Stop close-only probing and return any buffered visible prelude.
|
|
|
|
A provider that emits a dedicated reasoning field has already told us
|
|
how to classify its trace. Its ordinary ``content`` channel can then
|
|
stream immediately even when the model supports thinking.
|
|
"""
|
|
if self._leading_only:
|
|
if self._leading_phase != "probe":
|
|
if self._leading_phase != "implicit_probe":
|
|
return ""
|
|
self._leading_phase = "visible"
|
|
pending = self._pending
|
|
self._pending = ""
|
|
return pending
|
|
if self._channel_decided:
|
|
return ""
|
|
self._channel_decided = True
|
|
pending = self._pending
|
|
self._pending = ""
|
|
return pending
|
|
|
|
def finish_parts(self) -> tuple[str, str]:
|
|
"""Flush the remaining visible or thinking text at end of stream."""
|
|
if self._leading_only:
|
|
tail = self._pending
|
|
self._pending = ""
|
|
if self._leading_phase == "implicit_probe":
|
|
# No close-only marker ever arrived. The buffered bytes are a
|
|
# normal file body, not reasoning; preserve them exactly.
|
|
self._leading_phase = "visible"
|
|
return tail, ""
|
|
if self._leading_phase == "thinking":
|
|
self._inside_think = False
|
|
self._leading_phase = "visible"
|
|
return "", tail
|
|
self._leading_phase = "visible"
|
|
return tail, ""
|
|
if not self._channel_decided:
|
|
tail = self._pending
|
|
self._pending = ""
|
|
self._channel_decided = True
|
|
return tail, ""
|
|
if self._inside_think:
|
|
thinking = self._pending
|
|
self._pending = ""
|
|
self._inside_think = False
|
|
return "", thinking
|
|
tail = self._pending
|
|
self._pending = ""
|
|
return tail, ""
|
|
|
|
def finish(self) -> str:
|
|
"""Compatibility helper returning only the visible stream tail."""
|
|
return self.finish_parts()[0]
|
|
|
|
@property
|
|
def has_unclosed_thinking(self) -> bool:
|
|
"""Whether the stream is currently inside a reasoning envelope."""
|
|
|
|
if self._leading_only:
|
|
return self._leading_phase == "thinking"
|
|
return self._inside_think
|
|
|
|
|
|
def strip_inline_think_tags(content: str) -> str:
|
|
"""Return complete text with inline model-thinking blocks removed."""
|
|
text_filter = InlineThinkTagFilter()
|
|
return text_filter.push(content) + text_filter.finish()
|
|
|
|
|
|
__all__ = ["InlineThinkTagFilter", "strip_inline_think_tags"]
|