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:
parent
de4fe8bc7d
commit
2ba10dbb5b
5 changed files with 164 additions and 54 deletions
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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])
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -337,7 +337,9 @@
|
|||
"retrievePanel": {
|
||||
"chatMessage": {
|
||||
"copyTooltip": "复制到剪贴板",
|
||||
"copyError": "复制文本到剪贴板失败"
|
||||
"copyError": "复制文本到剪贴板失败",
|
||||
"thinking": "正在思考...",
|
||||
"thinkingTime": "思考用时 {{time}} 秒"
|
||||
},
|
||||
"retrieval": {
|
||||
"startPrompt": "输入查询开始检索",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue