Add AI thinking process display with expandable UI and timing

- Add thinking fields to Message type
- Create expandable thinking section
- Track thinking time during streaming
This commit is contained in:
yangdx 2025-09-08 18:09:04 +08:00
parent de4fe8bc7d
commit 2ba10dbb5b
5 changed files with 164 additions and 54 deletions

View file

@ -99,6 +99,9 @@ export type QueryMode = 'naive' | 'local' | 'global' | 'hybrid' | 'mix' | 'bypas
export type Message = {
role: 'user' | 'assistant' | 'system'
content: string
thinkingContent?: string
displayContent?: string
thinkingTime?: number
}
export type QueryRequest = {

View file

@ -15,12 +15,13 @@ import type { Element } from 'hast'
import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter'
import { oneLight, oneDark } from 'react-syntax-highlighter/dist/cjs/styles/prism'
import { LoaderIcon, CopyIcon } from 'lucide-react'
import { LoaderIcon, CopyIcon, ChevronDownIcon } from 'lucide-react'
import { useTranslation } from 'react-i18next'
export type MessageWithError = Message & {
id: string // Unique identifier for stable React keys
isError?: boolean
isThinking?: boolean // Flag to indicate if the message is in a "thinking" state
/**
* Indicates if the mermaid diagram in this message has been rendered.
* Used to persist the rendering state across updates and prevent flickering.
@ -33,6 +34,16 @@ export const ChatMessage = ({ message }: { message: MessageWithError }) => { //
const { t } = useTranslation()
const { theme } = useTheme()
const [katexPlugin, setKatexPlugin] = useState<any>(null)
const [isThinkingExpanded, setIsThinkingExpanded] = useState<boolean>(false)
// Directly use props passed from the parent.
const { thinkingContent, displayContent, thinkingTime, isThinking } = message
// The content to display is now non-ambiguous.
const finalThinkingContent = thinkingContent
// For user messages, displayContent will be undefined, so we fall back to content.
// For assistant messages, we prefer displayContent.
const finalDisplayContent = message.role === 'user' ? message.content : displayContent
// Load KaTeX dynamically
useEffect(() => {
@ -59,6 +70,27 @@ export const ChatMessage = ({ message }: { message: MessageWithError }) => { //
}
}, [message, t]) // Added t to dependency array
const mainMarkdownComponents = useMemo(() => ({
code: (props: any) => (
<CodeHighlight
{...props}
renderAsDiagram={message.mermaidRendered ?? false}
/>
),
p: ({ children }: { children?: ReactNode }) => <p className="my-2">{children}</p>,
h1: ({ children }: { children?: ReactNode }) => <h1 className="text-xl font-bold mt-4 mb-2">{children}</h1>,
h2: ({ children }: { children?: ReactNode }) => <h2 className="text-lg font-bold mt-4 mb-2">{children}</h2>,
h3: ({ children }: { children?: ReactNode }) => <h3 className="text-base font-bold mt-3 mb-2">{children}</h3>,
h4: ({ children }: { children?: ReactNode }) => <h4 className="text-base font-semibold mt-3 mb-2">{children}</h4>,
ul: ({ children }: { children?: ReactNode }) => <ul className="list-disc pl-5 my-2">{children}</ul>,
ol: ({ children }: { children?: ReactNode }) => <ol className="list-decimal pl-5 my-2">{children}</ol>,
li: ({ children }: { children?: ReactNode }) => <li className="my-1">{children}</li>
}), [message.mermaidRendered]);
const thinkingMarkdownComponents = useMemo(() => ({
code: (props: any) => (<CodeHighlight {...props} renderAsDiagram={message.mermaidRendered ?? false} />)
}), [message.mermaidRendered]);
return (
<div
className={`${
@ -69,55 +101,76 @@ export const ChatMessage = ({ message }: { message: MessageWithError }) => { //
: 'w-[95%] bg-muted'
} rounded-lg px-4 py-2`}
>
<div className="relative">
<ReactMarkdown
className="prose dark:prose-invert max-w-none text-sm break-words prose-headings:mt-4 prose-headings:mb-2 prose-p:my-2 prose-ul:my-2 prose-ol:my-2 prose-li:my-1 [&_.katex]:text-current [&_.katex-display]:my-4 [&_.katex-display]:overflow-x-auto"
remarkPlugins={[remarkGfm, remarkMath]}
rehypePlugins={[
...(katexPlugin ? [[
katexPlugin,
{
errorColor: theme === 'dark' ? '#ef4444' : '#dc2626',
throwOnError: false,
displayMode: false
}
] as any] : []),
rehypeReact
]}
skipHtml={false}
// Memoize the components object to prevent unnecessary re-renders of ReactMarkdown children
components={useMemo(() => ({
code: (props: any) => ( // Add type annotation if needed, e.g., props: CodeProps from 'react-markdown/lib/ast-to-react'
<CodeHighlight
{...props}
renderAsDiagram={message.mermaidRendered ?? false}
/>
),
p: ({ children }: { children?: ReactNode }) => <p className="my-2">{children}</p>,
h1: ({ children }: { children?: ReactNode }) => <h1 className="text-xl font-bold mt-4 mb-2">{children}</h1>,
h2: ({ children }: { children?: ReactNode }) => <h2 className="text-lg font-bold mt-4 mb-2">{children}</h2>,
h3: ({ children }: { children?: ReactNode }) => <h3 className="text-base font-bold mt-3 mb-2">{children}</h3>,
h4: ({ children }: { children?: ReactNode }) => <h4 className="text-base font-semibold mt-3 mb-2">{children}</h4>,
ul: ({ children }: { children?: ReactNode }) => <ul className="list-disc pl-5 my-2">{children}</ul>,
ol: ({ children }: { children?: ReactNode }) => <ol className="list-decimal pl-5 my-2">{children}</ol>,
li: ({ children }: { children?: ReactNode }) => <li className="my-1">{children}</li>
}), [message.mermaidRendered])} // Dependency ensures update if mermaid state changes
>
{message.content}
</ReactMarkdown>
{message.role === 'assistant' && message.content && message.content.length > 0 && ( // Added check for message.content existence
<Button
onClick={handleCopyMarkdown}
className="absolute right-0 bottom-0 size-6 rounded-md opacity-20 transition-opacity hover:opacity-100"
tooltip={t('retrievePanel.chatMessage.copyTooltip')}
variant="default"
size="icon"
{/* Thinking process display - only for assistant messages */}
{message.role === 'assistant' && (isThinking || thinkingTime !== null) && (
<div className="mb-2">
<div
className="flex items-center text-gray-500 dark:text-gray-400 hover:text-gray-700 dark:hover:text-gray-200 transition-colors duration-200 text-sm cursor-pointer select-none"
onClick={() => setIsThinkingExpanded(!isThinkingExpanded)}
>
<CopyIcon className="size-4" /> {/* Explicit size */}
</Button>
)}
</div>
{message.content === '' && <LoaderIcon className="animate-spin duration-2000" />} {/* Check for empty string specifically */}
{isThinking ? (
<>
<LoaderIcon className="mr-2 size-4 animate-spin" />
<span>{t('retrievePanel.chatMessage.thinking')}</span>
</>
) : (
typeof thinkingTime === 'number' && <span>{t('retrievePanel.chatMessage.thinkingTime', { time: thinkingTime })}</span>
)}
{finalThinkingContent && <ChevronDownIcon className={`ml-2 size-4 shrink-0 transition-transform ${isThinkingExpanded ? 'rotate-180' : ''}`} />}
</div>
{isThinkingExpanded && finalThinkingContent && (
<div className="mt-2 pl-4 border-l-2 border-primary/20 text-sm prose dark:prose-invert max-w-none break-words prose-p:my-1 prose-headings:my-2">
<ReactMarkdown
remarkPlugins={[remarkGfm, remarkMath]}
rehypePlugins={[
...(katexPlugin ? [[katexPlugin, { errorColor: theme === 'dark' ? '#ef4444' : '#dc2626', throwOnError: false, displayMode: false }] as any] : []),
rehypeReact
]}
skipHtml={false}
components={thinkingMarkdownComponents}
>
{finalThinkingContent}
</ReactMarkdown>
</div>
)}
</div>
)}
{/* Main content display */}
{finalDisplayContent && (
<div className="relative">
<ReactMarkdown
className="prose dark:prose-invert max-w-none text-sm break-words prose-headings:mt-4 prose-headings:mb-2 prose-p:my-2 prose-ul:my-2 prose-ol:my-2 prose-li:my-1 [&_.katex]:text-current [&_.katex-display]:my-4 [&_.katex-display]:overflow-x-auto"
remarkPlugins={[remarkGfm, remarkMath]}
rehypePlugins={[
...(katexPlugin ? [[
katexPlugin,
{
errorColor: theme === 'dark' ? '#ef4444' : '#dc2626',
throwOnError: false,
displayMode: false
}
] as any] : []),
rehypeReact
]}
skipHtml={false}
components={mainMarkdownComponents}
>
{finalDisplayContent}
</ReactMarkdown>
{message.role === 'assistant' && finalDisplayContent && finalDisplayContent.length > 0 && (
<Button
onClick={handleCopyMarkdown}
className="absolute right-0 bottom-0 size-6 rounded-md opacity-20 transition-opacity hover:opacity-100"
tooltip={t('retrievePanel.chatMessage.copyTooltip')}
variant="default"
size="icon"
>
<CopyIcon className="size-4" /> {/* Explicit size */}
</Button>
)}
</div>
)}
{message.content === '' && !isThinking && !thinkingTime && <LoaderIcon className="animate-spin duration-2000" />} {/* Check for empty string specifically */}
</div>
)
}

View file

@ -58,6 +58,7 @@ export default function RetrievalTesting() {
const [inputError, setInputError] = useState('') // Error message for input
// Reference to track if we should follow scroll during streaming (using ref for synchronous updates)
const shouldFollowScrollRef = useRef(true)
const thinkingStartTime = useRef<number | null>(null)
// Reference to track if user interaction is from the form area
const isFormInteractionRef = useRef(false)
// Reference to track if scroll was triggered programmatically
@ -153,6 +154,39 @@ export default function RetrievalTesting() {
const updateAssistantMessage = (chunk: string, isError?: boolean) => {
assistantMessage.content += chunk
// Start thinking timer on first sight of think tag
if (assistantMessage.content.includes('<think>') && !thinkingStartTime.current) {
thinkingStartTime.current = Date.now()
}
// Real-time parsing for streaming
const thinkStartTag = '<think>';
const thinkEndTag = '</think>';
const thinkStartIndex = assistantMessage.content.indexOf(thinkStartTag);
const thinkEndIndex = assistantMessage.content.indexOf(thinkEndTag);
if (thinkStartIndex !== -1) {
if (thinkEndIndex !== -1) {
// Thinking has finished for this chunk, calculate time now
assistantMessage.isThinking = false
if (thinkingStartTime.current && !assistantMessage.thinkingTime) {
const duration = (Date.now() - thinkingStartTime.current) / 1000
assistantMessage.thinkingTime = parseFloat(duration.toFixed(2))
}
assistantMessage.thinkingContent = assistantMessage.content.substring(thinkStartIndex + thinkStartTag.length, thinkEndIndex).trim()
assistantMessage.displayContent = assistantMessage.content.substring(thinkEndIndex + thinkEndTag.length).trim()
} else {
// Still thinking
assistantMessage.isThinking = true;
assistantMessage.thinkingContent = assistantMessage.content.substring(thinkStartIndex + thinkStartTag.length);
assistantMessage.displayContent = '';
}
} else {
assistantMessage.isThinking = false;
assistantMessage.displayContent = assistantMessage.content;
}
// Detect if the assistant message contains a complete mermaid code block
// Simple heuristic: look for ```mermaid ... ```
const mermaidBlockRegex = /```mermaid\s+([\s\S]+?)```/g
@ -171,9 +205,12 @@ export default function RetrievalTesting() {
const newMessages = [...prev]
const lastMessage = newMessages[newMessages.length - 1]
if (lastMessage.role === 'assistant') {
lastMessage.content = assistantMessage.content
lastMessage.isError = isError
lastMessage.mermaidRendered = assistantMessage.mermaidRendered
lastMessage.content = assistantMessage.content;
lastMessage.thinkingContent = assistantMessage.thinkingContent;
lastMessage.displayContent = assistantMessage.displayContent;
lastMessage.isThinking = assistantMessage.isThinking;
lastMessage.isError = isError;
lastMessage.mermaidRendered = assistantMessage.mermaidRendered;
}
return newMessages
})
@ -223,6 +260,19 @@ export default function RetrievalTesting() {
// Clear loading and add messages to state
setIsLoading(false)
isReceivingResponseRef.current = false
// Final calculation for thinking time, only if not already calculated
if (assistantMessage.thinkingContent && thinkingStartTime.current && !assistantMessage.thinkingTime) {
const duration = (Date.now() - thinkingStartTime.current) / 1000
assistantMessage.thinkingTime = parseFloat(duration.toFixed(2))
}
// Ensure isThinking is false at the very end
assistantMessage.isThinking = false;
// Always reset the timer at the end of a query
if (thinkingStartTime.current) {
thinkingStartTime.current = null;
}
useSettingsStore
.getState()
.setRetrievalHistory([...prevMessages, userMessage, assistantMessage])

View file

@ -337,7 +337,9 @@
"retrievePanel": {
"chatMessage": {
"copyTooltip": "Copy to clipboard",
"copyError": "Failed to copy text to clipboard"
"copyError": "Failed to copy text to clipboard",
"thinking": "Thinking...",
"thinkingTime": "Thinking time {{time}}s"
},
"retrieval": {
"startPrompt": "Start a retrieval by typing your query below",

View file

@ -337,7 +337,9 @@
"retrievePanel": {
"chatMessage": {
"copyTooltip": "复制到剪贴板",
"copyError": "复制文本到剪贴板失败"
"copyError": "复制文本到剪贴板失败",
"thinking": "正在思考...",
"thinkingTime": "思考用时 {{time}} 秒"
},
"retrieval": {
"startPrompt": "输入查询开始检索",