feat: retrieval message code highlight, message copy button

lightrag_webui/src/components/retrieval/ChatMessage.tsx

@@ -0,0 +1,93 @@
+import { ReactNode, useCallback } from 'react'
+import { Message } from '@/api/lightrag'
+import useTheme from '@/hooks/useTheme'
+import Button from '@/components/ui/Button'
+
+import ReactMarkdown from 'react-markdown'
+import remarkGfm from 'remark-gfm'
+import rehypeReact from 'rehype-react'
+import remarkMath from 'remark-math'
+import ShikiHighlighter, { isInlineCode, type Element } from 'react-shiki'
+
+import { LoaderIcon, CopyIcon } from 'lucide-react'
+
+export type MessageWithError = Message & {
+  isError?: boolean
+}
+
+export const ChatMessage = ({ message }: { message: MessageWithError }) => {
+  const handleCopyMarkdown = useCallback(async () => {
+    if (message.content) {
+      try {
+        await navigator.clipboard.writeText(message.content)
+      } catch (err) {
+        console.error('Failed to copy:', err)
+      }
+    }
+  }, [message])
+
+  return (
+    <div
+      className={`max-w-[80%] rounded-lg px-4 py-2 ${
+        message.role === 'user'
+          ? 'bg-primary text-primary-foreground'
+          : message.isError
+            ? 'bg-red-100 text-red-600 dark:bg-red-950 dark:text-red-400'
+            : 'bg-muted'
+      }`}
+    >
+      <pre className="relative break-words whitespace-pre-wrap">
+        <ReactMarkdown
+          className="dark:prose-invert max-w-none text-base text-sm"
+          remarkPlugins={[remarkGfm, remarkMath]}
+          rehypePlugins={[rehypeReact]}
+          skipHtml={false}
+          components={{
+            code: CodeHighlight
+          }}
+        >
+          {message.content}
+        </ReactMarkdown>
+        {message.role === 'assistant' && message.content.length > 0 && (
+          <Button
+            onClick={handleCopyMarkdown}
+            className="absolute right-0 bottom-0 size-6 rounded-md opacity-20 transition-opacity hover:opacity-100"
+            tooltip="Copy to clipboard"
+            variant="default"
+            size="icon"
+          >
+            <CopyIcon />
+          </Button>
+        )}
+      </pre>
+      {message.content.length === 0 && <LoaderIcon className="animate-spin duration-2000" />}
+    </div>
+  )
+}
+
+interface CodeHighlightProps {
+  className?: string | undefined
+  children?: ReactNode | undefined
+  node?: Element | undefined
+}
+
+const CodeHighlight = ({ className, children, node, ...props }: CodeHighlightProps) => {
+  const { theme } = useTheme()
+  const match = className?.match(/language-(\w+)/)
+  const language = match ? match[1] : undefined
+  const inline: boolean | undefined = node ? isInlineCode(node) : undefined
+
+  return !inline ? (
+    <ShikiHighlighter
+      language={language}
+      theme={theme === 'dark' ? 'houston' : 'github-light'}
+      {...props}
+    >
+      {String(children)}
+    </ShikiHighlighter>
+  ) : (
+    <code className={className} {...props}>
+      {children}
+    </code>
+  )
+}

lightrag_webui/src/features/RetrievalTesting.tsx

@@ -1,38 +1,16 @@
 import Input from '@/components/ui/Input'
 import Button from '@/components/ui/Button'
 import { useCallback, useEffect, useRef, useState } from 'react'
-import { queryText, queryTextStream, Message as ChatMessage } from '@/api/lightrag'
+import { queryText, queryTextStream, Message } from '@/api/lightrag'
 import { errorMessage } from '@/lib/utils'
 import { useSettingsStore } from '@/stores/settings'
 import { useDebounce } from '@/hooks/useDebounce'
 import QuerySettings from '@/components/retrieval/QuerySettings'
-
-import ReactMarkdown from 'react-markdown'
-import remarkGfm from 'remark-gfm'
-import rehypeReact from 'rehype-react'
-import remarkMath from 'remark-math'
-
-import { EraserIcon, SendIcon, LoaderIcon } from 'lucide-react'
-
-type Message = ChatMessage & {
-  isError?: boolean
-}
-
-const ChatMessageComponent = ({ message }: { message: Message }) => {
-  return (
-    <ReactMarkdown
-      className="prose lg:prose-xs dark:prose-invert max-w-none text-base"
-      remarkPlugins={[remarkGfm, remarkMath]}
-      rehypePlugins={[rehypeReact]}
-      skipHtml={false}
-    >
-      {message.content}
-    </ReactMarkdown>
-  )
-}
+import { ChatMessage, MessageWithError } from '@/components/retrieval/ChatMessage'
+import { EraserIcon, SendIcon } from 'lucide-react'
 
 export default function RetrievalTesting() {
-  const [messages, setMessages] = useState<Message[]>(
+  const [messages, setMessages] = useState<MessageWithError[]>(
     () => useSettingsStore.getState().retrievalHistory || []
   )
   const [inputValue, setInputValue] = useState('')
@@ -147,22 +125,7 @@ export default function RetrievalTesting() {
                     key={idx}
                     className={`flex ${message.role === 'user' ? 'justify-end' : 'justify-start'}`}
                   >
-                    <div
-                      className={`max-w-[80%] rounded-lg px-4 py-2 ${
-                        message.role === 'user'
-                          ? 'bg-primary text-primary-foreground'
-                          : message.isError
-                            ? 'bg-red-100 text-red-600 dark:bg-red-950 dark:text-red-400'
-                            : 'bg-muted'
-                      }`}
-                    >
-                      <pre className="break-words whitespace-pre-wrap">
-                        {<ChatMessageComponent message={message} />}
-                      </pre>
-                      {message.content.length === 0 && (
-                        <LoaderIcon className="animate-spin duration-2000" />
-                      )}
-                    </div>
+                    {<ChatMessage message={message} />}
                   </div>
                 ))
               )}

lightrag_webui/src/features/SiteHeader.tsx

@@ -55,6 +55,7 @@ export default function SiteHeader() {
     <header className="border-border/40 bg-background/95 supports-[backdrop-filter]:bg-background/60 sticky top-0 z-50 flex h-10 w-full border-b px-4 backdrop-blur">
       <a href="/" className="mr-6 flex items-center gap-2">
         <ZapIcon className="size-4 text-emerald-400" aria-hidden="true" />
+        {/* <img src='/logo.png' className="size-4" /> */}
         <span className="font-bold md:inline-block">{SiteInfo.name}</span>
       </a>
 

lightrag_webui/src/lib/constants.ts

@@ -27,10 +27,39 @@ export const defaultQueryLabel = '*'
 // reference: https://developer.mozilla.org/en-US/docs/Web/HTTP/MIME_types/Common_types
 export const supportedFileTypes = {
   'text/plain': [
-      '.txt', '.md', '.html', '.htm', '.tex', '.json', '.xml', '.yaml', '.yml',
-      '.rtf', '.odt', '.epub', '.csv', '.log', '.conf', '.ini', '.properties',
-      '.sql', '.bat', '.sh', '.c', '.cpp', '.py', '.java', '.js', '.ts',
-      '.swift', '.go', '.rb', '.php', '.css', '.scss', '.less'
+    '.txt',
+    '.md',
+    '.html',
+    '.htm',
+    '.tex',
+    '.json',
+    '.xml',
+    '.yaml',
+    '.yml',
+    '.rtf',
+    '.odt',
+    '.epub',
+    '.csv',
+    '.log',
+    '.conf',
+    '.ini',
+    '.properties',
+    '.sql',
+    '.bat',
+    '.sh',
+    '.c',
+    '.cpp',
+    '.py',
+    '.java',
+    '.js',
+    '.ts',
+    '.swift',
+    '.go',
+    '.rb',
+    '.php',
+    '.css',
+    '.scss',
+    '.less'
   ],
   'application/pdf': ['.pdf'],
   'application/msword': ['.doc'],