enhance web UI with retrieval testing and UI improvements

lightrag_webui/eslint.config.js

@@ -10,7 +10,7 @@ import react from 'eslint-plugin-react'
 export default tseslint.config(
   { ignores: ['dist'] },
   {
-    extends: [js.configs.recommended, ...tseslint.configs.recommended],
+    extends: [js.configs.recommended, ...tseslint.configs.recommended, prettier],
     files: ['**/*.{ts,tsx,js,jsx}'],
     languageOptions: {
       ecmaVersion: 2020,
@@ -31,7 +31,6 @@ export default tseslint.config(
       '@stylistic/js/indent': ['error', 2],
       '@stylistic/js/quotes': ['error', 'single'],
       '@typescript-eslint/no-explicit-any': ['off']
-    },
-    prettier
+    }
   }
 )

lightrag_webui/src/App.tsx

@@ -1,3 +1,4 @@
+import { useState, useCallback } from 'react'
 import ThemeProvider from '@/components/ThemeProvider'
 import MessageAlert from '@/components/MessageAlert'
 import StatusIndicator from '@/components/StatusIndicator'
@@ -10,14 +11,16 @@ import SiteHeader from '@/features/SiteHeader'
 
 import GraphViewer from '@/features/GraphViewer'
 import DocumentManager from '@/features/DocumentManager'
+import RetrievalTesting from '@/features/RetrievalTesting'
 
 import { Tabs, TabsContent } from '@/components/ui/Tabs'
 
 function App() {
   const message = useBackendState.use.message()
   const enableHealthCheck = useSettingsStore.use.enableHealthCheck()
+  const [currentTab] = useState(() => useSettingsStore.getState().currentTab)
 
-  // health check
+  // Health check
   useEffect(() => {
     if (!enableHealthCheck) return
 
@@ -30,25 +33,36 @@ function App() {
     return () => clearInterval(interval)
   }, [enableHealthCheck])
 
+  const handleTabChange = useCallback(
+    (tab: string) => useSettingsStore.getState().setCurrentTab(tab as any),
+    []
+  )
+
   return (
     <ThemeProvider>
-      <div className="flex h-screen w-screen">
-        <Tabs defaultValue="knowledge-graph" className="flex size-full flex-col">
+      <main className="flex h-screen w-screen overflow-x-hidden">
+        <Tabs
+          defaultValue={currentTab}
+          className="!m-0 flex grow flex-col !p-0"
+          onValueChange={handleTabChange}
+        >
           <SiteHeader />
-          <TabsContent value="documents" className="flex-1">
-            <DocumentManager />
-          </TabsContent>
-          <TabsContent value="knowledge-graph" className="flex-1">
-            <GraphViewer />
-          </TabsContent>
-          <TabsContent value="settings" className="size-full">
-            <h1> Settings </h1>
-          </TabsContent>
+          <div className="relative grow">
+            <TabsContent value="documents" className="absolute top-0 right-0 bottom-0 left-0">
+              <DocumentManager />
+            </TabsContent>
+            <TabsContent value="knowledge-graph" className="absolute top-0 right-0 bottom-0 left-0">
+              <GraphViewer />
+            </TabsContent>
+            <TabsContent value="retrieval" className="absolute top-0 right-0 bottom-0 left-0">
+              <RetrievalTesting />
+            </TabsContent>
+          </div>
         </Tabs>
-      </div>
-      {enableHealthCheck && <StatusIndicator />}
-      {message !== null && <MessageAlert />}
-      <Toaster />
+        {enableHealthCheck && <StatusIndicator />}
+        {message !== null && <MessageAlert />}
+        <Toaster />
+      </main>
     </ThemeProvider>
   )
 }

lightrag_webui/src/api/lightrag.ts

@@ -151,32 +151,64 @@ export const queryText = async (request: QueryRequest): Promise<QueryResponse> =
   return response.data
 }
 
-export const queryTextStream = async (request: QueryRequest, onChunk: (chunk: string) => void) => {
-  const response = await axiosInstance.post('/query/stream', request, {
-    responseType: 'stream'
-  })
+export const queryTextStream = async (
+  request: QueryRequest,
+  onChunk: (chunk: string) => void,
+  onError?: (error: string) => void
+) => {
+  try {
+    let buffer = ''
+    await axiosInstance.post('/query/stream', request, {
+      responseType: 'text',
+      headers: {
+        Accept: 'application/x-ndjson'
+      },
+      transformResponse: [
+        (data: string) => {
+          // Accumulate the data and process complete lines
+          buffer += data
+          const lines = buffer.split('\n')
+          // Keep the last potentially incomplete line in the buffer
+          buffer = lines.pop() || ''
 
-  const reader = response.data.getReader()
-  const decoder = new TextDecoder()
-
-  while (true) {
-    const { done, value } = await reader.read()
-    if (done) break
-
-    const chunk = decoder.decode(value)
-    const lines = chunk.split('\n')
-    for (const line of lines) {
-      if (line) {
-        try {
-          const data = JSON.parse(line)
-          if (data.response) {
-            onChunk(data.response)
+          for (const line of lines) {
+            if (line.trim()) {
+              try {
+                const parsed = JSON.parse(line)
+                if (parsed.response) {
+                  onChunk(parsed.response)
+                } else if (parsed.error && onError) {
+                  onError(parsed.error)
+                }
+              } catch (e) {
+                console.error('Error parsing stream chunk:', e)
+                if (onError) onError('Error parsing server response')
+              }
+            }
           }
-        } catch (e) {
-          console.error('Error parsing stream chunk:', e)
+          return data
+        }
+      ]
+    })
+
+    // Process any remaining data in the buffer
+    if (buffer.trim()) {
+      try {
+        const parsed = JSON.parse(buffer)
+        if (parsed.response) {
+          onChunk(parsed.response)
+        } else if (parsed.error && onError) {
+          onError(parsed.error)
         }
+      } catch (e) {
+        console.error('Error parsing final chunk:', e)
+        if (onError) onError('Error parsing server response')
       }
     }
+  } catch (error) {
+    const message = errorMessage(error)
+    console.error('Stream request failed:', message)
+    if (onError) onError(message)
   }
 }
 
@@ -199,6 +231,7 @@ export const uploadDocument = async (
     headers: {
       'Content-Type': 'multipart/form-data'
     },
+    // prettier-ignore
     onUploadProgress:
       onUploadProgress !== undefined
         ? (progressEvent) => {

lightrag_webui/src/components/ThemeToggle.tsx

@@ -19,6 +19,7 @@ export default function ThemeToggle() {
         variant={controlButtonVariant}
         tooltip="Switch to light theme"
         size="icon"
+        side="bottom"
       >
         <MoonIcon />
       </Button>
@@ -30,6 +31,7 @@ export default function ThemeToggle() {
       variant={controlButtonVariant}
       tooltip="Switch to dark theme"
       size="icon"
+      side="bottom"
     >
       <SunIcon />
     </Button>

lightrag_webui/src/components/document/ClearDocumentsDialog.tsx

@@ -15,7 +15,7 @@ import { clearDocuments } from '@/api/lightrag'
 import { EraserIcon } from 'lucide-react'
 
 export default function ClearDocumentsDialog() {
-  const [open, setOpen] = useState(false) // 添加状态控制
+  const [open, setOpen] = useState(false)
 
   const handleClear = useCallback(async () => {
     try {
@@ -34,8 +34,8 @@ export default function ClearDocumentsDialog() {
   return (
     <Dialog open={open} onOpenChange={setOpen}>
       <DialogTrigger asChild>
-        <Button variant="outline" tooltip="Clear documents" side="bottom" size="icon">
-          <EraserIcon />
+        <Button variant="outline" side="bottom" tooltip='Clear documents' size="sm">
+          <EraserIcon/> Clear
         </Button>
       </DialogTrigger>
       <DialogContent className="sm:max-w-xl" onCloseAutoFocus={(e) => e.preventDefault()}>

lightrag_webui/src/components/document/UploadDocumentsDialog.tsx

@@ -66,8 +66,8 @@ export default function UploadDocumentsDialog() {
       }}
     >
       <DialogTrigger asChild>
-        <Button variant="outline" tooltip="Upload documents" side="bottom" size="icon">
-          <UploadIcon />
+        <Button variant="default" side="bottom" tooltip='Upload documents' size="sm">
+          <UploadIcon /> Upload
         </Button>
       </DialogTrigger>
       <DialogContent className="sm:max-w-xl" onCloseAutoFocus={(e) => e.preventDefault()}>

lightrag_webui/src/components/ui/Table.tsx

@@ -56,6 +56,7 @@ TableRow.displayName = 'TableRow'
 const TableHead = React.forwardRef<
   HTMLTableCellElement,
   React.ThHTMLAttributes<HTMLTableCellElement>
+// eslint-disable-next-line react/prop-types
 >(({ className, ...props }, ref) => (
   <th
     ref={ref}
@@ -71,6 +72,7 @@ TableHead.displayName = 'TableHead'
 const TableCell = React.forwardRef<
   HTMLTableCellElement,
   React.TdHTMLAttributes<HTMLTableCellElement>
+  // eslint-disable-next-line react/prop-types
 >(({ className, ...props }, ref) => (
   <td
     ref={ref}

lightrag_webui/src/features/DocumentManager.tsx

@@ -16,23 +16,23 @@ import ClearDocumentsDialog from '@/components/document/ClearDocumentsDialog'
 
 import {
   getDocuments,
-  getDocumentsScanProgress,
-  scanNewDocuments,
-  LightragDocumentsScanProgress
+  // getDocumentsScanProgress,
+  scanNewDocuments
+  // LightragDocumentsScanProgress
 } from '@/api/lightrag'
 import { errorMessage } from '@/lib/utils'
 import { toast } from 'sonner'
-import { useBackendState } from '@/stores/state'
+// import { useBackendState } from '@/stores/state'
 
 import { RefreshCwIcon, TrashIcon } from 'lucide-react'
 
 // type DocumentStatus = 'indexed' | 'pending' | 'indexing' | 'error'
 
 export default function DocumentManager() {
-  const health = useBackendState.use.health()
+  // const health = useBackendState.use.health()
   const [files, setFiles] = useState<string[]>([])
   const [indexedFiles, setIndexedFiles] = useState<string[]>([])
-  const [scanProgress, setScanProgress] = useState<LightragDocumentsScanProgress | null>(null)
+  // const [scanProgress, setScanProgress] = useState<LightragDocumentsScanProgress | null>(null)
 
   const fetchDocuments = useCallback(async () => {
     try {
@@ -45,7 +45,7 @@ export default function DocumentManager() {
 
   useEffect(() => {
     fetchDocuments()
-  }, [])
+  }, []) // eslint-disable-line react-hooks/exhaustive-deps
 
   const scanDocuments = useCallback(async () => {
     try {
@@ -54,26 +54,26 @@ export default function DocumentManager() {
     } catch (err) {
       toast.error('Failed to load documents\n' + errorMessage(err))
     }
-  }, [setFiles])
+  }, [])
 
-  useEffect(() => {
-    const interval = setInterval(async () => {
-      try {
-        if (!health) return
-        const progress = await getDocumentsScanProgress()
-        setScanProgress((pre) => {
-          if (pre?.is_scanning === progress.is_scanning && progress.is_scanning === false) {
-            return pre
-          }
-          return progress
-        })
-        console.log(progress)
-      } catch (err) {
-        toast.error('Failed to get scan progress\n' + errorMessage(err))
-      }
-    }, 2000)
-    return () => clearInterval(interval)
-  }, [health])
+  // useEffect(() => {
+  //   const interval = setInterval(async () => {
+  //     try {
+  //       if (!health) return
+  //       const progress = await getDocumentsScanProgress()
+  //       setScanProgress((pre) => {
+  //         if (pre?.is_scanning === progress.is_scanning && progress.is_scanning === false) {
+  //           return pre
+  //         }
+  //         return progress
+  //       })
+  //       console.log(progress)
+  //     } catch (err) {
+  //       toast.error('Failed to get scan progress\n' + errorMessage(err))
+  //     }
+  //   }, 2000)
+  //   return () => clearInterval(interval)
+  // }, [health])
 
   const handleDelete = async (fileName: string) => {
     console.log(`deleting ${fileName}`)
@@ -88,19 +88,19 @@ export default function DocumentManager() {
         <div className="flex gap-2">
           <Button
             variant="outline"
-            size="icon"
-            tooltip="Scan Documents"
             onClick={scanDocuments}
             side="bottom"
+            tooltip="Scan documents"
+            size="sm"
           >
-            <RefreshCwIcon />
+            <RefreshCwIcon /> Scan
           </Button>
           <div className="flex-1" />
           <ClearDocumentsDialog />
           <UploadDocumentsDialog />
         </div>
 
-        {scanProgress?.is_scanning && (
+        {/* {scanProgress?.is_scanning && (
           <div className="space-y-2">
             <div className="flex justify-between text-sm">
               <span>Indexing {scanProgress.current_file}</span>
@@ -108,7 +108,7 @@ export default function DocumentManager() {
             </div>
             <Progress value={scanProgress.progress} />
           </div>
-        )}
+        )} */}
 
         <Card>
           <CardHeader>

lightrag_webui/src/features/RetrievalTesting.tsx

@@ -0,0 +1,170 @@
+import Input from '@/components/ui/Input'
+import Button from '@/components/ui/Button'
+import { useCallback, useEffect, useRef, useState } from 'react'
+import { queryTextStream, QueryMode } from '@/api/lightrag'
+import { errorMessage } from '@/lib/utils'
+import { useSettingsStore } from '@/stores/settings'
+import { useDebounce } from '@/hooks/useDebounce'
+import { EraserIcon, SendIcon, LoaderIcon } from 'lucide-react'
+
+type Message = {
+  id: string
+  content: string
+  role: 'User' | 'LightRAG'
+}
+
+export default function RetrievalTesting() {
+  const [messages, setMessages] = useState<Message[]>(
+    () => useSettingsStore.getState().retrievalHistory || []
+  )
+  const [inputValue, setInputValue] = useState('')
+  const [isLoading, setIsLoading] = useState(false)
+  const [mode, setMode] = useState<QueryMode>('mix')
+  const messagesEndRef = useRef<HTMLDivElement>(null)
+
+  const scrollToBottom = useCallback(() => {
+    messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' })
+  }, [])
+
+  const handleSubmit = useCallback(
+    async (e: React.FormEvent) => {
+      e.preventDefault()
+      if (!inputValue.trim() || isLoading) return
+
+      const userMessage: Message = {
+        id: Date.now().toString(),
+        content: inputValue,
+        role: 'User'
+      }
+
+      const assistantMessage: Message = {
+        id: (Date.now() + 1).toString(),
+        content: '',
+        role: 'LightRAG'
+      }
+
+      setMessages((prev) => {
+        const newMessages = [...prev, userMessage, assistantMessage]
+        return newMessages
+      })
+
+      setInputValue('')
+      setIsLoading(true)
+
+      // Create a function to update the assistant's message
+      const updateAssistantMessage = (chunk: string) => {
+        assistantMessage.content += chunk
+        setMessages((prev) => {
+          const newMessages = [...prev]
+          const lastMessage = newMessages[newMessages.length - 1]
+          if (lastMessage.role === 'LightRAG') {
+            lastMessage.content = assistantMessage.content
+          }
+          return newMessages
+        })
+      }
+
+      try {
+        await queryTextStream(
+          {
+            query: userMessage.content,
+            mode: mode,
+            stream: true
+          },
+          updateAssistantMessage
+        )
+      } catch (err) {
+        updateAssistantMessage(`Error: Failed to get response\n${errorMessage(err)}`)
+      } finally {
+        setIsLoading(false)
+        useSettingsStore
+          .getState()
+          .setRetrievalHistory([
+            ...useSettingsStore.getState().retrievalHistory,
+            userMessage,
+            assistantMessage
+          ])
+      }
+    },
+    [inputValue, isLoading, mode, setMessages]
+  )
+
+  const debouncedMessages = useDebounce(messages, 100)
+  useEffect(() => scrollToBottom(), [debouncedMessages, scrollToBottom])
+
+  const clearMessages = useCallback(() => {
+    setMessages([])
+    useSettingsStore.getState().setRetrievalHistory([])
+  }, [setMessages])
+
+  return (
+    <div className="flex size-full flex-col gap-4 px-32 py-6">
+      <div className="relative grow">
+        <div className="bg-primary-foreground/60 absolute inset-0 flex flex-col overflow-auto rounded-lg border p-2">
+          <div className="flex min-h-0 flex-1 flex-col gap-2">
+            {messages.length === 0 ? (
+              <div className="text-muted-foreground flex h-full items-center justify-center text-lg">
+                Start a retrieval by typing your query below
+              </div>
+            ) : (
+              messages.map((message) => (
+                <div
+                  key={message.id}
+                  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' : 'bg-muted'
+                    }`}
+                  >
+                    <pre className="break-words whitespace-pre-wrap">{message.content}</pre>
+                    {message.content.length === 0 && (
+                      <LoaderIcon className="animate-spin duration-2000" />
+                    )}
+                  </div>
+                </div>
+              ))
+            )}
+            <div ref={messagesEndRef} className="pb-1" />
+          </div>
+        </div>
+      </div>
+
+      <form onSubmit={handleSubmit} className="flex shrink-0 items-center gap-2 pb-2">
+        <Button
+          type="button"
+          variant="outline"
+          onClick={clearMessages}
+          disabled={isLoading}
+          size="sm"
+        >
+          <EraserIcon />
+          Clear
+        </Button>
+        <select
+          className="border-input bg-background ring-offset-background h-9 rounded-md border px-3 py-1 text-sm"
+          value={mode}
+          onChange={(e) => setMode(e.target.value as QueryMode)}
+          disabled={isLoading}
+        >
+          <option value="naive">Naive</option>
+          <option value="local">Local</option>
+          <option value="global">Global</option>
+          <option value="hybrid">Hybrid</option>
+          <option value="mix">Mix</option>
+        </select>
+        <Input
+          className="flex-1"
+          value={inputValue}
+          onChange={(e) => setInputValue(e.target.value)}
+          placeholder="Type your query..."
+          disabled={isLoading}
+        />
+        <Button type="submit" variant="default" disabled={isLoading} size="sm">
+          <SendIcon />
+          Send
+        </Button>
+      </form>
+    </div>
+  )
+}

lightrag_webui/src/features/SiteHeader.tsx

@@ -2,14 +2,18 @@ import Button from '@/components/ui/Button'
 import { SiteInfo } from '@/lib/constants'
 import ThemeToggle from '@/components/ThemeToggle'
 import { TabsList, TabsTrigger } from '@/components/ui/Tabs'
+import { useSettingsStore } from '@/stores/settings'
+import { cn } from '@/lib/utils'
 
 import { ZapIcon, GithubIcon } from 'lucide-react'
 
 export default function SiteHeader() {
+  const currentTab = useSettingsStore.use.currentTab()
+
   return (
     <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-teal-400" aria-hidden="true" />
+        <ZapIcon className="size-4 text-emerald-400" aria-hidden="true" />
         <span className="font-bold md:inline-block">{SiteInfo.name}</span>
       </a>
 
@@ -18,28 +22,43 @@ export default function SiteHeader() {
           <TabsList className="h-full gap-2">
             <TabsTrigger
               value="documents"
-              className="hover:bg-background/60 cursor-pointer px-2 py-1 transition-all"
+              className={cn(
+                'cursor-pointer px-2 py-1 transition-all',
+                currentTab === 'documents'
+                  ? '!bg-emerald-400 !text-zinc-50'
+                  : 'hover:bg-background/60'
+              )}
             >
               Documents
             </TabsTrigger>
             <TabsTrigger
               value="knowledge-graph"
-              className="hover:bg-background/60 cursor-pointer px-2 py-1 transition-all"
+              className={cn(
+                'cursor-pointer px-2 py-1 transition-all',
+                currentTab === 'knowledge-graph'
+                  ? '!bg-emerald-400 !text-zinc-50'
+                  : 'hover:bg-background/60'
+              )}
             >
               Knowledge Graph
             </TabsTrigger>
-            {/* <TabsTrigger
-              value="settings"
-              className="hover:bg-background/60 cursor-pointer px-2 py-1 transition-all"
+            <TabsTrigger
+              value="retrieval"
+              className={cn(
+                'cursor-pointer px-2 py-1 transition-all',
+                currentTab === 'retrieval'
+                  ? '!bg-emerald-400 !text-zinc-50'
+                  : 'hover:bg-background/60'
+              )}
             >
-              Settings
-            </TabsTrigger> */}
+              Retrieval
+            </TabsTrigger>
           </TabsList>
         </div>
       </div>
 
       <nav className="flex items-center">
-        <Button variant="ghost" size="icon">
+        <Button variant="ghost" size="icon" side="bottom" tooltip="Project Repository">
           <a href={SiteInfo.github} target="_blank" rel="noopener noreferrer">
             <GithubIcon className="size-4" aria-hidden="true" />
           </a>

lightrag_webui/src/stores/settings.ts

@@ -4,6 +4,7 @@ import { createSelectors } from '@/lib/utils'
 import { defaultQueryLabel } from '@/lib/constants'
 
 type Theme = 'dark' | 'light' | 'system'
+type Tab = 'documents' | 'knowledge-graph' | 'retrieval'
 
 interface SettingsState {
   theme: Theme
@@ -27,6 +28,12 @@ interface SettingsState {
 
   apiKey: string | null
   setApiKey: (key: string | null) => void
+
+  currentTab: Tab
+  setCurrentTab: (tab: Tab) => void
+
+  retrievalHistory: any[]
+  setRetrievalHistory: (history: any[]) => void
 }
 
 const useSettingsStoreBase = create<SettingsState>()(
@@ -49,6 +56,10 @@ const useSettingsStoreBase = create<SettingsState>()(
 
       apiKey: null,
 
+      currentTab: 'documents',
+
+      retrievalHistory: [],
+
       setTheme: (theme: Theme) => set({ theme }),
 
       setQueryLabel: (queryLabel: string) =>
@@ -58,12 +69,19 @@ const useSettingsStoreBase = create<SettingsState>()(
 
       setEnableHealthCheck: (enable: boolean) => set({ enableHealthCheck: enable }),
 
-      setApiKey: (apiKey: string | null) => set({ apiKey })
+      setApiKey: (apiKey: string | null) => set({ apiKey }),
+
+      setCurrentTab: (tab: Tab) => set({ currentTab: tab }),
+
+      setRetrievalHistory: (history: any[]) => set({ retrievalHistory: history })
     }),
     {
       name: 'settings-storage',
       storage: createJSONStorage(() => localStorage),
-      version: 4,
+      version: 5,
+      partialize(state) {
+        return { ...state, retrievalHistory: undefined }
+      },
       migrate: (state: any, version: number) => {
         if (version < 2) {
           state.showEdgeLabel = false
@@ -78,6 +96,9 @@ const useSettingsStoreBase = create<SettingsState>()(
           state.enableHealthCheck = true
           state.apiKey = null
         }
+        if (version < 5) {
+          state.currentTab = 'documents'
+        }
       }
     }
   )