goodguy1963-ComfyUI-ThinkingLLM/AILab_StreamDisplay.py

import re
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 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 shown as a rolling 200-character
    tail on a single line so older tokens scroll off naturally.
    """

    _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,
    ):
        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

    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
        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
        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 _write_rolling_line(self, label: str, tail: str, final: bool = False):
        """Emit a carriage-return line showing the last ~200 chars, clearing to EOL.
        'final' adds a trailing newline so the last frame stays visible."""
        if not tail:
            return
        msg = f"\r[{label}] {tail}\033[K"
        sys.stdout.write(msg)
        if final:
            sys.stdout.write("\n")
        sys.stdout.flush()

    def push_compact(self, text: str):
        """Append *text* and emit a carriage-return rolling tail every throttle interval.

        Uses \\r to redraw a single line so the terminal stays clean without
        flooding the ComfyUI log buffer.
        """
        if not text:
            return
        self._compact_active = True
        self._stage_chars += len(text)
        self._buffer += text
        if len(self._buffer) > self.TAIL_LENGTH * 2:
            self._buffer = self._buffer[-self.TAIL_LENGTH:]
        tail = self._compact_tail()
        now = time.monotonic()
        should_emit = False
        if tail and tail != self._last_compact_tail:
            if self._last_compact_emit_at <= 0.0:
                should_emit = len(self._buffer) >= self.min_chunk_chars
            elif (now - self._last_compact_emit_at) >= self.flush_interval:
                should_emit = True
        if should_emit:
            self._stage_updates += 1
            self._last_compact_tail = tail
            self._last_compact_emit_at = now
            self._write_rolling_line(self.label, tail)

    def end_compact(self):
        """Finish compact mode with one final newline-terminated snapshot."""
        tail = self._compact_tail()
        if tail and tail != self._last_compact_tail:
            self._stage_updates += 1
        self._last_compact_tail = ""
        self._last_compact_emit_at = 0.0
        self._compact_active = False
        self._buffer = ""
        self._write_rolling_line(self.label, tail or "", final=True)