diff --git a/lightrag_webui/src/api/lightrag.ts b/lightrag_webui/src/api/lightrag.ts index 265126c7..a86519d4 100644 --- a/lightrag_webui/src/api/lightrag.ts +++ b/lightrag_webui/src/api/lightrag.ts @@ -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 = { diff --git a/lightrag_webui/src/components/retrieval/ChatMessage.tsx b/lightrag_webui/src/components/retrieval/ChatMessage.tsx index 0314fd92..18b4fb25 100644 --- a/lightrag_webui/src/components/retrieval/ChatMessage.tsx +++ b/lightrag_webui/src/components/retrieval/ChatMessage.tsx @@ -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(null) + const [isThinkingExpanded, setIsThinkingExpanded] = useState(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) => ( + + ), + p: ({ children }: { children?: ReactNode }) =>

{children}

, + h1: ({ children }: { children?: ReactNode }) =>

{children}

, + h2: ({ children }: { children?: ReactNode }) =>

{children}

, + h3: ({ children }: { children?: ReactNode }) =>

{children}

, + h4: ({ children }: { children?: ReactNode }) =>

{children}

, + ul: ({ children }: { children?: ReactNode }) =>
    {children}
, + ol: ({ children }: { children?: ReactNode }) =>
    {children}
, + li: ({ children }: { children?: ReactNode }) =>
  • {children}
  • + }), [message.mermaidRendered]); + + const thinkingMarkdownComponents = useMemo(() => ({ + code: (props: any) => () + }), [message.mermaidRendered]); + return (
    { // : 'w-[95%] bg-muted' } rounded-lg px-4 py-2`} > -
    - ({ - code: (props: any) => ( // Add type annotation if needed, e.g., props: CodeProps from 'react-markdown/lib/ast-to-react' - - ), - p: ({ children }: { children?: ReactNode }) =>

    {children}

    , - h1: ({ children }: { children?: ReactNode }) =>

    {children}

    , - h2: ({ children }: { children?: ReactNode }) =>

    {children}

    , - h3: ({ children }: { children?: ReactNode }) =>

    {children}

    , - h4: ({ children }: { children?: ReactNode }) =>

    {children}

    , - ul: ({ children }: { children?: ReactNode }) =>
      {children}
    , - ol: ({ children }: { children?: ReactNode }) =>
      {children}
    , - li: ({ children }: { children?: ReactNode }) =>
  • {children}
  • - }), [message.mermaidRendered])} // Dependency ensures update if mermaid state changes - > - {message.content} -
    - {message.role === 'assistant' && message.content && message.content.length > 0 && ( // Added check for message.content existence - - )} -
    - {message.content === '' && } {/* Check for empty string specifically */} + {isThinking ? ( + <> + + {t('retrievePanel.chatMessage.thinking')} + + ) : ( + typeof thinkingTime === 'number' && {t('retrievePanel.chatMessage.thinkingTime', { time: thinkingTime })} + )} + {finalThinkingContent && } +
    + {isThinkingExpanded && finalThinkingContent && ( +
    + + {finalThinkingContent} + +
    + )} + + )} + {/* Main content display */} + {finalDisplayContent && ( +
    + + {finalDisplayContent} + + {message.role === 'assistant' && finalDisplayContent && finalDisplayContent.length > 0 && ( + + )} +
    + )} + {message.content === '' && !isThinking && !thinkingTime && } {/* Check for empty string specifically */} ) } diff --git a/lightrag_webui/src/features/RetrievalTesting.tsx b/lightrag_webui/src/features/RetrievalTesting.tsx index 0744ae5e..288dfeae 100644 --- a/lightrag_webui/src/features/RetrievalTesting.tsx +++ b/lightrag_webui/src/features/RetrievalTesting.tsx @@ -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(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('') && !thinkingStartTime.current) { + thinkingStartTime.current = Date.now() + } + + // Real-time parsing for streaming + const thinkStartTag = ''; + const thinkEndTag = ''; + 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]) diff --git a/lightrag_webui/src/locales/en.json b/lightrag_webui/src/locales/en.json index 1c5cf6e9..7590f74c 100644 --- a/lightrag_webui/src/locales/en.json +++ b/lightrag_webui/src/locales/en.json @@ -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", diff --git a/lightrag_webui/src/locales/zh.json b/lightrag_webui/src/locales/zh.json index 951090f7..6d6323aa 100644 --- a/lightrag_webui/src/locales/zh.json +++ b/lightrag_webui/src/locales/zh.json @@ -337,7 +337,9 @@ "retrievePanel": { "chatMessage": { "copyTooltip": "复制到剪贴板", - "copyError": "复制文本到剪贴板失败" + "copyError": "复制文本到剪贴板失败", + "thinking": "正在思考...", + "thinkingTime": "思考用时 {{time}} 秒" }, "retrieval": { "startPrompt": "输入查询开始检索",