414 lines
15 KiB
Python
414 lines
15 KiB
Python
import re
|
|
import shutil
|
|
import sys
|
|
import time
|
|
|
|
# ── Global singleton for live terminal streaming ──────────────────────────
|
|
THINKING_STREAM_DISPLAY = None
|
|
|
|
|
|
def get_thinking_stream_display():
|
|
"""Return or create a compact-mode TerminalStreamDisplay for live
|
|
token output during LLM inference. Use the same label across all
|
|
nodes so the user sees a single coherent stream in the terminal."""
|
|
global THINKING_STREAM_DISPLAY
|
|
if THINKING_STREAM_DISPLAY is None:
|
|
THINKING_STREAM_DISPLAY = TerminalStreamDisplay(
|
|
label="ThinkingLLM",
|
|
flush_chars=160,
|
|
min_chunk_chars=48,
|
|
flush_interval=0.8,
|
|
compact=True,
|
|
)
|
|
return THINKING_STREAM_DISPLAY
|
|
|
|
|
|
def reset_thinking_stream_display():
|
|
"""Kill the singleton so the next call to *get_thinking_stream_display*
|
|
creates a fresh display. Safe to call between queues."""
|
|
global THINKING_STREAM_DISPLAY
|
|
if THINKING_STREAM_DISPLAY is not None:
|
|
THINKING_STREAM_DISPLAY.end_stage()
|
|
THINKING_STREAM_DISPLAY = None
|
|
|
|
|
|
# ── Shared chunk normalizer for OpenAI-style streamed responses ────────────
|
|
|
|
def _get_stream_field(container, key: str, default=None):
|
|
if isinstance(container, dict):
|
|
return container.get(key, default)
|
|
return getattr(container, key, default)
|
|
|
|
|
|
def _stream_value_to_text(value) -> str:
|
|
if value is None:
|
|
return ""
|
|
if isinstance(value, str):
|
|
return value
|
|
if isinstance(value, (int, float)):
|
|
return str(value)
|
|
if isinstance(value, list):
|
|
return "".join(_stream_value_to_text(part) for part in value)
|
|
if isinstance(value, dict):
|
|
for key in ("text", "content", "reasoning_content", "reasoning", "thinking"):
|
|
text = _stream_value_to_text(value.get(key))
|
|
if text:
|
|
return text
|
|
return ""
|
|
return str(value)
|
|
|
|
|
|
def _first_stream_text(containers, keys: tuple[str, ...]) -> str:
|
|
for container in containers:
|
|
if container is None:
|
|
continue
|
|
for key in keys:
|
|
text = _stream_value_to_text(_get_stream_field(container, key, None))
|
|
if text:
|
|
return text
|
|
return ""
|
|
|
|
|
|
def extract_stream_token(chunk) -> dict[str, str]:
|
|
"""Return a dict with 'reasoning' and 'content' keys from an OpenAI-style
|
|
streaming delta chunk.
|
|
|
|
Populates both fields from any supported field name so callers get reasoning
|
|
and answer text regardless of backend (llama.cpp, vLLM, HF, etc.).
|
|
Empty values are returned as empty strings.
|
|
"""
|
|
choices = _get_stream_field(chunk, "choices", []) or []
|
|
choice = choices[0] if isinstance(choices, (list, tuple)) and choices else {}
|
|
delta = _get_stream_field(choice, "delta", {}) or {}
|
|
message = _get_stream_field(choice, "message", {}) or {}
|
|
|
|
containers = (delta, message, choice, chunk)
|
|
content = _first_stream_text(containers, ("content", "text", "output_text"))
|
|
reasoning = _first_stream_text(
|
|
containers,
|
|
("reasoning_content", "thinking", "reasoning", "reasoning_text"),
|
|
)
|
|
|
|
return {"reasoning": reasoning, "content": content}
|
|
|
|
|
|
class StreamDegenerationError(RuntimeError):
|
|
def __init__(self, reason: str):
|
|
super().__init__(reason)
|
|
self.reason = reason
|
|
|
|
|
|
class StreamDegenerationGuard:
|
|
def __init__(self, repeated_word_limit: int = 8, repeated_ngram_limit: int = 8):
|
|
self.repeated_word_limit = max(4, int(repeated_word_limit))
|
|
self.repeated_ngram_limit = max(3, int(repeated_ngram_limit))
|
|
self._text = ""
|
|
self._words: list[str] = []
|
|
|
|
def push(self, text: str):
|
|
if not text:
|
|
return
|
|
self._text = (self._text + text)[-4096:]
|
|
words = [w.casefold() for w in re.findall(r"[A-Za-z][A-Za-z'\-]{2,}", self._text)]
|
|
if not words:
|
|
return
|
|
self._words = words[-256:]
|
|
reason = self._detect()
|
|
if reason:
|
|
raise StreamDegenerationError(reason)
|
|
|
|
def _detect(self) -> str:
|
|
if not self._words:
|
|
return ""
|
|
last_word = self._words[-1]
|
|
if len(last_word) >= 4:
|
|
run = 0
|
|
for word in reversed(self._words):
|
|
if word != last_word:
|
|
break
|
|
run += 1
|
|
if run >= self.repeated_word_limit:
|
|
return f"repeated word '{last_word}' {run} times"
|
|
|
|
for size in range(2, 6):
|
|
needed = size * self.repeated_ngram_limit
|
|
if len(self._words) < needed:
|
|
continue
|
|
tail = self._words[-size:]
|
|
if all(self._words[-offset - size : -offset] == tail for offset in range(size, needed, size)):
|
|
return f"repeated {size}-word phrase {' '.join(tail)!r}"
|
|
return ""
|
|
|
|
|
|
def strip_degenerate_repetition(text: str, repeated_word_limit: int = 8) -> str:
|
|
if not text:
|
|
return ""
|
|
matches = list(re.finditer(r"[A-Za-z][A-Za-z'\-]{2,}", text))
|
|
if not matches:
|
|
return text
|
|
last = matches[-1].group(0).casefold()
|
|
if len(last) < 4:
|
|
return text
|
|
run_start = len(matches) - 1
|
|
while run_start >= 0 and matches[run_start].group(0).casefold() == last:
|
|
run_start -= 1
|
|
run_count = len(matches) - run_start - 1
|
|
if run_count < repeated_word_limit:
|
|
return text
|
|
cut_at = matches[run_start + 1].start()
|
|
return text[:cut_at].rstrip(" \t\r\n,.;:!?")
|
|
|
|
|
|
class TerminalStreamDisplay:
|
|
"""Buffered terminal display for streamed model output.
|
|
|
|
Keeps the raw output accumulation in the caller, but renders only readable
|
|
chunks to the terminal with stage headers and completion summaries.
|
|
|
|
When *compact* is True, token content is printed as clean wrapped text
|
|
without node labels, carriage returns, or tail truncation.
|
|
"""
|
|
|
|
_SENTENCE_BOUNDARY_RE = re.compile(r"[.!?](?:['\")\]]+)?(?:\s+|$)")
|
|
_PARAGRAPH_BOUNDARY_RE = re.compile(r"\n\s*\n")
|
|
_PLANNING_RE = re.compile(
|
|
r"(?im)^\s*(?:[-*]\s*)?(?:\*\*)?\s*(okay[,.:]?|first[,.:]?|next[,.:]?|then[,.:]?|wait[,.:]?|final\s+plan|final\s+check)\b"
|
|
)
|
|
_FIRST_PERSON_RE = re.compile(r"(?i)\b(i\s+(should|need|must|will|am\s+going\s+to|have\s+to))\b")
|
|
TAIL_LENGTH = 200
|
|
|
|
def __init__(
|
|
self,
|
|
label: str,
|
|
flush_chars: int = 160,
|
|
min_chunk_chars: int = 48,
|
|
flush_interval: float = 1.0,
|
|
suppress_planning: bool = False,
|
|
compact: bool = True,
|
|
line_width: int | None = None,
|
|
):
|
|
self.label = label
|
|
self.flush_chars = max(32, int(flush_chars))
|
|
self.min_chunk_chars = max(1, int(min_chunk_chars))
|
|
self.flush_interval = max(0.1, float(flush_interval))
|
|
self.suppress_planning = suppress_planning
|
|
self.compact = compact
|
|
self._buffer = ""
|
|
self._last_flush_at = time.monotonic()
|
|
self._stage_name = None
|
|
self._stage_started_at = None
|
|
self._stage_chars = 0
|
|
self._stage_updates = 0
|
|
self._last_compact_emit_at = 0.0
|
|
self._last_compact_tail = ""
|
|
self._compact_active = False
|
|
self._stream_pending_text = ""
|
|
self._stream_col = 0
|
|
self._stream_need_space = False
|
|
self._line_width_override = int(line_width) if line_width else None
|
|
|
|
def start_stage(self, stage_name: str):
|
|
self.end_stage()
|
|
self._stage_name = stage_name
|
|
self._stage_started_at = time.monotonic()
|
|
self._stage_chars = 0
|
|
self._stage_updates = 0
|
|
if self.compact:
|
|
return
|
|
self._write_line(f"[{self.label}] {stage_name}")
|
|
|
|
def push(self, text: str):
|
|
if not text:
|
|
return
|
|
self._stage_chars += len(text)
|
|
self._buffer += text
|
|
if self._should_flush():
|
|
self.flush()
|
|
|
|
def flush(self, force: bool = False):
|
|
while True:
|
|
flush_at = self._find_flush_index(force=force)
|
|
if flush_at <= 0:
|
|
break
|
|
chunk = self._buffer[:flush_at]
|
|
self._buffer = self._buffer[flush_at:]
|
|
display_text = self._normalize_chunk(chunk)
|
|
self._last_flush_at = time.monotonic()
|
|
if display_text:
|
|
self._stage_updates += 1
|
|
self._write_line(display_text)
|
|
if not force:
|
|
break
|
|
|
|
def _compact_tail(self) -> str:
|
|
"""Return a single line containing the rolling last ~200 chars."""
|
|
raw = (self._buffer or "")
|
|
if len(raw) > self.TAIL_LENGTH:
|
|
raw = "…" + raw[-self.TAIL_LENGTH:]
|
|
return re.sub(r"[ \t]+", " ", raw.replace("\r", "").replace("\n", " ")).strip()
|
|
|
|
def end_stage(self):
|
|
if self._compact_active:
|
|
self.end_compact()
|
|
else:
|
|
self.flush(force=True)
|
|
if not self._stage_name:
|
|
return
|
|
if self.compact:
|
|
self._stage_name = None
|
|
self._stage_started_at = None
|
|
self._stage_chars = 0
|
|
self._stage_updates = 0
|
|
return
|
|
elapsed = 0.0
|
|
if self._stage_started_at is not None:
|
|
elapsed = max(0.0, time.monotonic() - self._stage_started_at)
|
|
self._write_line(
|
|
f"[{self.label}] {self._stage_name} complete ({elapsed:.1f}s, {self._stage_chars} chars, {self._stage_updates} updates)"
|
|
)
|
|
self._stage_name = None
|
|
self._stage_started_at = None
|
|
self._stage_chars = 0
|
|
self._stage_updates = 0
|
|
|
|
def _should_flush(self) -> bool:
|
|
if not self._buffer:
|
|
return False
|
|
if len(self._buffer) >= self.flush_chars:
|
|
return True
|
|
if self._last_boundary_index() >= self.min_chunk_chars:
|
|
return True
|
|
if (time.monotonic() - self._last_flush_at) >= self.flush_interval and len(self._buffer) >= self.min_chunk_chars:
|
|
return True
|
|
return False
|
|
|
|
def _find_flush_index(self, force: bool) -> int:
|
|
if not self._buffer:
|
|
return 0
|
|
if force:
|
|
return len(self._buffer)
|
|
boundary_index = self._last_boundary_index()
|
|
if boundary_index >= self.min_chunk_chars:
|
|
return boundary_index
|
|
if len(self._buffer) >= self.flush_chars:
|
|
whitespace_index = self._buffer.rfind(" ", 0, self.flush_chars)
|
|
return whitespace_index + 1 if whitespace_index > 0 else self.flush_chars
|
|
if (time.monotonic() - self._last_flush_at) >= self.flush_interval and len(self._buffer) >= self.min_chunk_chars:
|
|
whitespace_index = self._buffer.rfind(" ")
|
|
return whitespace_index + 1 if whitespace_index > 0 else len(self._buffer)
|
|
return 0
|
|
|
|
def _last_boundary_index(self) -> int:
|
|
last_end = 0
|
|
for match in self._PARAGRAPH_BOUNDARY_RE.finditer(self._buffer):
|
|
last_end = max(last_end, match.end())
|
|
for match in self._SENTENCE_BOUNDARY_RE.finditer(self._buffer):
|
|
last_end = max(last_end, match.end())
|
|
return last_end
|
|
|
|
def _normalize_chunk(self, chunk: str) -> str:
|
|
text = chunk.replace("\r", "")
|
|
parts = []
|
|
for block in re.split(r"\n\s*\n", text):
|
|
normalized = re.sub(r"[ \t]+", " ", block).strip()
|
|
if not normalized:
|
|
continue
|
|
if self.suppress_planning and self._looks_like_planning(normalized):
|
|
continue
|
|
parts.append(normalized)
|
|
return "\n".join(parts)
|
|
|
|
def _looks_like_planning(self, text: str) -> bool:
|
|
return bool(self._PLANNING_RE.search(text) or self._FIRST_PERSON_RE.search(text))
|
|
|
|
def _write_line(self, text: str):
|
|
if not text:
|
|
return
|
|
msg = f"{text}\n"
|
|
sys.stdout.write(msg)
|
|
sys.stdout.flush()
|
|
|
|
def _stream_line_width(self) -> int:
|
|
if self._line_width_override:
|
|
return max(24, self._line_width_override)
|
|
columns = shutil.get_terminal_size(fallback=(120, 24)).columns
|
|
return max(40, min(96, columns // 2))
|
|
|
|
def _write_wrapped_stream_text(self, text: str):
|
|
if not text:
|
|
return
|
|
width = self._stream_line_width()
|
|
for part in re.findall(r"\S+|\s+", text.replace("\r", "")):
|
|
if not part:
|
|
continue
|
|
if part.isspace():
|
|
newline_count = part.count("\n")
|
|
if newline_count:
|
|
sys.stdout.write("\n" * newline_count)
|
|
self._stream_col = 0
|
|
self._stream_need_space = False
|
|
elif self._stream_col > 0:
|
|
self._stream_need_space = True
|
|
continue
|
|
|
|
word = part
|
|
prefix = 1 if self._stream_need_space and self._stream_col > 0 else 0
|
|
if self._stream_col > 0 and self._stream_col + prefix + len(word) > width:
|
|
sys.stdout.write("\n")
|
|
self._stream_col = 0
|
|
prefix = 0
|
|
if prefix:
|
|
sys.stdout.write(" ")
|
|
self._stream_col += 1
|
|
sys.stdout.write(word)
|
|
self._stream_col += len(word)
|
|
self._stream_need_space = False
|
|
sys.stdout.flush()
|
|
|
|
def _drain_stream_pending(self, force: bool = False):
|
|
if not self._stream_pending_text:
|
|
return
|
|
text = self._stream_pending_text
|
|
if force:
|
|
emit = text
|
|
self._stream_pending_text = ""
|
|
else:
|
|
match = None
|
|
for match in re.finditer(r"\s+", text):
|
|
pass
|
|
if match is None:
|
|
if len(text) < self._stream_line_width():
|
|
return
|
|
emit = text
|
|
self._stream_pending_text = ""
|
|
else:
|
|
end = match.end()
|
|
emit = text[:end]
|
|
self._stream_pending_text = text[end:]
|
|
self._write_wrapped_stream_text(emit)
|
|
|
|
def push_compact(self, text: str):
|
|
"""Append *text* as clean wrapped terminal output without node prefixes."""
|
|
if not text:
|
|
return
|
|
self._compact_active = True
|
|
self._stage_chars += len(text)
|
|
self._stage_updates += 1
|
|
self._stream_pending_text += text
|
|
self._drain_stream_pending(force=False)
|
|
|
|
def end_compact(self):
|
|
"""Finish compact mode with one final newline-terminated snapshot."""
|
|
had_pending = bool(self._stream_pending_text)
|
|
self._drain_stream_pending(force=True)
|
|
if (had_pending or self._stream_col > 0) and self._stage_chars:
|
|
sys.stdout.write("\n")
|
|
sys.stdout.flush()
|
|
self._last_compact_tail = ""
|
|
self._last_compact_emit_at = 0.0
|
|
self._compact_active = False
|
|
self._buffer = ""
|
|
self._stream_pending_text = ""
|
|
self._stream_col = 0
|
|
self._stream_need_space = False
|