Stream qwen3 thinking tokens with collapsible UI panel

Route <think>…</think> content from the llama-server stream into
thinking_delta SSE events instead of discarding it. A small look-ahead
buffer in _route_content() handles tags split across chunk boundaries.
Strip thinking tags from the persisted final_message so the DB stores
only the actual response.

Frontend: ThinkingBubble component accumulates the live reasoning text
in a scrollable panel that auto-collapses to "Thought for Xs" once the
first content token arrives. Falls back to the existing heartbeat
behaviour for models that don't emit <think> tags.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-09-12 16:26:25 -04:00
parent 2a1f125fa5
commit bb2103e81e
5 changed files with 174 additions and 18 deletions
+70 -16
View File
@@ -2,6 +2,7 @@ import asyncio
import json
import logging
import os
import re
import time
import uuid
@@ -35,6 +36,54 @@ from .prompts import SIMBA_SYSTEM_PROMPT
TTFT_TIMEOUT = float(os.getenv("TTFT_TIMEOUT_SECONDS", "30"))
_THINK_OPEN = "<think>"
_THINK_CLOSE = "</think>"
_THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL)
def _route_content(
buf: str, text: str, in_thinking: bool
) -> tuple[str, str, bool, str]:
"""Split a streaming content chunk into (thinking_text, content_text).
Handles tags split across chunk boundaries by keeping a small trailing
buffer. Returns the updated in_thinking state and any leftover buffer.
"""
combined = buf + text
thinking_parts: list[str] = []
content_parts: list[str] = []
while combined:
tag, hold = (
(_THINK_CLOSE, len(_THINK_CLOSE) - 1)
if in_thinking
else (_THINK_OPEN, len(_THINK_OPEN) - 1)
)
idx = combined.find(tag)
if idx == -1:
# Tag not yet complete — hold back enough bytes to catch a split.
if len(combined) <= hold:
return (
"".join(thinking_parts),
"".join(content_parts),
in_thinking,
combined,
)
safe, leftover = combined[:-hold], combined[-hold:]
(thinking_parts if in_thinking else content_parts).append(safe)
return (
"".join(thinking_parts),
"".join(content_parts),
in_thinking,
leftover,
)
(thinking_parts if in_thinking else content_parts).append(combined[:idx])
combined = combined[idx + len(tag) :]
in_thinking = not in_thinking
return "".join(thinking_parts), "".join(content_parts), in_thinking, ""
conversation_blueprint = Blueprint(
"conversation_api", __name__, url_prefix="/api/conversation"
)
@@ -222,9 +271,10 @@ async def stream_query():
async def event_generator():
final_message = None
streamed = ""
# Reasoning models emit empty-content deltas while thinking; surface a
# throttled heartbeat so the UI shows progress (and bytes keep flowing
# on the SSE connection) instead of dead air for the whole think.
# Tag-routing state for <think>…</think> blocks.
in_thinking = False
tag_buf = ""
# Fallback heartbeat for models that emit empty chunks before thinking text.
last_thinking = 0.0
turn_has_content = False
@@ -341,19 +391,23 @@ async def stream_query():
elif event_type == "on_chat_model_stream":
chunk = event.get("data", {}).get("chunk")
if chunk is not None and not first_content_at:
extra = getattr(chunk, "additional_kwargs", {})
if extra:
logging.info("chunk.additional_kwargs sample: %s", extra)
content = getattr(chunk, "content", None)
if isinstance(content, str) and content:
if first_content_at is None:
first_content_at = time.monotonic()
turn_has_content = True
n_tokens += 1
streamed += content
yield f"data: {json.dumps({'type': 'content', 'delta': content})}\n\n"
raw = getattr(chunk, "content", None)
if isinstance(raw, str) and raw:
think_text, content_text, in_thinking, tag_buf = _route_content(
tag_buf, raw, in_thinking
)
if think_text:
yield f"data: {json.dumps({'type': 'thinking_delta', 'delta': think_text})}\n\n"
if content_text:
if first_content_at is None:
first_content_at = time.monotonic()
turn_has_content = True
n_tokens += 1
streamed += content_text
yield f"data: {json.dumps({'type': 'content', 'delta': content_text})}\n\n"
elif not turn_has_content:
# Empty chunk before any content — heartbeat for models
# that don't emit <think> tags at all.
now = time.monotonic()
if now - last_thinking >= 1.0:
last_thinking = now
@@ -367,7 +421,7 @@ async def stream_query():
last_msg = msgs[-1]
content = getattr(last_msg, "content", None)
if isinstance(content, str) and content:
final_message = content
final_message = _THINK_RE.sub("", content).strip()
except asyncio.CancelledError:
cancelled = True
@@ -2,6 +2,7 @@ import { userService } from "./userService";
export type SSEEvent =
| { type: "thinking" }
| { type: "thinking_delta"; delta: string }
| { type: "tool_start"; tool: string }
| { type: "tool_end"; tool: string }
| { type: "content"; delta: string }
+9 -2
View File
@@ -4,6 +4,7 @@ import { PanelLeftOpen, Menu, X } from "lucide-react";
import { QuestionBubble } from "./QuestionBubble";
import { AnswerBubble } from "./AnswerBubble";
import { ToolBubble } from "./ToolBubble";
import { ThinkingBubble } from "./ThinkingBubble";
import { MessageInput } from "./MessageInput";
import { SidebarContent } from "./SidebarContent";
import { AdminPanel } from "./AdminPanel";
@@ -63,6 +64,8 @@ export const ChatScreen = ({ setAuthenticated, isAdmin }: ChatScreenProps) => {
clearError,
stalled,
retryStalled,
thinkingText,
thinkingDurationMs,
} = useChat({
createConversation,
refreshConversations,
@@ -252,8 +255,12 @@ export const ChatScreen = ({ setAuthenticated, isAdmin }: ChatScreenProps) => {
return <QuestionBubble key={index} text={msg.text} image_key={msg.image_key} />;
})}
{thinking && !streaming && (
<ToolBubble text="Thinking..." pulse />
{isLoading && (thinking || thinkingText || thinkingDurationMs !== null) && (
<ThinkingBubble
text={thinkingText}
durationMs={thinkingDurationMs}
active={thinking}
/>
)}
{((isLoading && !streaming && !thinking) || messagesLoading) && (
@@ -0,0 +1,78 @@
import React, { useState, useEffect, useRef } from "react";
import { ChevronDown } from "lucide-react";
import { cn } from "../lib/utils";
type ThinkingBubbleProps = {
text: string;
durationMs: number | null;
active: boolean;
};
export const ThinkingBubble = React.memo(
({ text, durationMs, active }: ThinkingBubbleProps) => {
const [expanded, setExpanded] = useState(true);
const textRef = useRef<HTMLDivElement>(null);
// Auto-scroll thinking text as it streams in.
useEffect(() => {
if (active && textRef.current) {
textRef.current.scrollTop = textRef.current.scrollHeight;
}
}, [text, active]);
// Collapse automatically when thinking finishes.
useEffect(() => {
if (durationMs !== null) {
setExpanded(false);
}
}, [durationMs]);
const label = active
? "Thinking…"
: durationMs !== null
? `Thought for ${(durationMs / 1000).toFixed(1)}s`
: "Thought";
return (
<div className="flex justify-center message-enter">
<div className="w-full max-w-2xl">
<button
onClick={() => setExpanded((v) => !v)}
className={cn(
"w-full flex items-center justify-between gap-2",
"px-3 py-1.5 rounded-xl",
"bg-leaf-pale border border-leaf-light/50",
"text-xs text-leaf-dark italic",
"hover:bg-leaf-light/20 transition-colors cursor-pointer",
active && "animate-pulse",
)}
>
<span>{label}</span>
{text && (
<ChevronDown
size={13}
className={cn(
"text-leaf-dark/60 transition-transform duration-200 shrink-0",
expanded && "rotate-180",
)}
/>
)}
</button>
{expanded && text && (
<div
ref={textRef}
className={cn(
"mt-1 px-3 py-2 rounded-xl max-h-48 overflow-y-auto",
"bg-leaf-pale/40 border border-leaf-light/30",
"text-xs text-leaf-dark/70 font-mono leading-relaxed whitespace-pre-wrap",
)}
>
{text}
</div>
)}
</div>
</div>
);
},
);
+16
View File
@@ -54,8 +54,11 @@ export function useChat({
const [stalled, setStalled] = useState(false);
const [streaming, setStreaming] = useState(false);
const [thinking, setThinking] = useState(false);
const [thinkingText, setThinkingText] = useState("");
const [thinkingDurationMs, setThinkingDurationMs] = useState<number | null>(null);
const stalledQueryRef = useRef<string | null>(null);
const thinkingStartRef = useRef<number | null>(null);
const isMountedRef = useRef(true);
// True once tokens for the current assistant reply have started arriving, so
@@ -149,6 +152,9 @@ export function useChat({
setError(null);
setStalled(false);
stalledQueryRef.current = null;
setThinkingText("");
setThinkingDurationMs(null);
thinkingStartRef.current = null;
streamingActiveRef.current = false;
setStreaming(false);
updateMessages((prev) => prev.concat([{ text: query, speaker: "user" }]));
@@ -196,6 +202,11 @@ export function useChat({
if (!isMountedRef.current) return;
if (event.type === "thinking") {
setThinking(true);
if (thinkingStartRef.current === null) thinkingStartRef.current = Date.now();
} else if (event.type === "thinking_delta") {
setThinking(true);
if (thinkingStartRef.current === null) thinkingStartRef.current = Date.now();
setThinkingText((prev) => prev + event.delta);
} else if (event.type === "tool_start") {
// Any text after a tool call belongs in a fresh bubble.
streamingActiveRef.current = false;
@@ -207,6 +218,9 @@ export function useChat({
);
} else if (event.type === "content") {
setThinking(false);
if (thinkingStartRef.current !== null && thinkingDurationMs === null) {
setThinkingDurationMs(Date.now() - thinkingStartRef.current);
}
if (streamingActiveRef.current) {
updateMessages((prev) => {
const updated = [...prev];
@@ -325,5 +339,7 @@ export function useChat({
clearError,
stalled,
retryStalled,
thinkingText,
thinkingDurationMs,
};
}