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:
@@ -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 }
|
||||
|
||||
@@ -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>
|
||||
);
|
||||
},
|
||||
);
|
||||
@@ -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,
|
||||
};
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user