improve streaming error handling

lightrag_webui/src/api/lightrag.ts

@@ -212,38 +212,42 @@ export const queryTextStream = async (
 ) => {
   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() || ''
+    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() || ''
 
-          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)
+            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)
-                if (onError) onError('Error parsing server response')
               }
             }
+            return data
           }
-          return data
-        }
-      ]
-    })
+        ]
+      })
+      .catch((error) => {
+        if (onError) onError(errorMessage(error))
+      })
 
     // Process any remaining data in the buffer
     if (buffer.trim()) {
@@ -266,11 +270,13 @@ export const queryTextStream = async (
   }
 }
 
-export const insertText = async (
-  text: string,
-  description?: string
-): Promise<DocActionResponse> => {
-  const response = await axiosInstance.post('/documents/text', { text, description })
+export const insertText = async (text: string): Promise<DocActionResponse> => {
+  const response = await axiosInstance.post('/documents/text', { text })
+  return response.data
+}
+
+export const insertTexts = async (texts: string[]): Promise<DocActionResponse> => {
+  const response = await axiosInstance.post('/documents/texts', { texts })
   return response.data
 }
 

lightrag_webui/src/features/DocumentManager.tsx

@@ -29,7 +29,16 @@ export default function DocumentManager() {
     try {
       const docs = await getDocuments()
       if (docs && docs.statuses) {
-        setDocs(docs)
+        // compose all documents count
+        const numDocuments = Object.values(docs.statuses).reduce(
+          (acc, status) => acc + status.length,
+          0
+        )
+        if (numDocuments > 0) {
+          setDocs(docs)
+        } else {
+          setDocs(null)
+        }
         // console.log(docs)
       } else {
         setDocs(null)

lightrag_webui/src/features/RetrievalTesting.tsx

@@ -1,7 +1,7 @@
 import Input from '@/components/ui/Input'
 import Button from '@/components/ui/Button'
 import { useCallback, useEffect, useRef, useState } from 'react'
-import { queryText, queryTextStream, Message } from '@/api/lightrag'
+import { queryText, queryTextStream, Message as ChatMessage } from '@/api/lightrag'
 import { errorMessage } from '@/lib/utils'
 import { useSettingsStore } from '@/stores/settings'
 import { useDebounce } from '@/hooks/useDebounce'
@@ -9,6 +9,10 @@ import QuerySettings from '@/components/retrieval/QuerySettings'
 
 import { EraserIcon, SendIcon, LoaderIcon } from 'lucide-react'
 
+type Message = ChatMessage & {
+  isError?: boolean
+}
+
 export default function RetrievalTesting() {
   const [messages, setMessages] = useState<Message[]>(
     () => useSettingsStore.getState().retrievalHistory || []
@@ -47,13 +51,14 @@ export default function RetrievalTesting() {
       setIsLoading(true)
 
       // Create a function to update the assistant's message
-      const updateAssistantMessage = (chunk: string) => {
+      const updateAssistantMessage = (chunk: string, isError?: boolean) => {
         assistantMessage.content += chunk
         setMessages((prev) => {
           const newMessages = [...prev]
           const lastMessage = newMessages[newMessages.length - 1]
           if (lastMessage.role === 'assistant') {
             lastMessage.content = assistantMessage.content
+            lastMessage.isError = isError
           }
           return newMessages
         })
@@ -65,19 +70,30 @@ export default function RetrievalTesting() {
         ...state.querySettings,
         query: userMessage.content,
         conversation_history: prevMessages
+          .filter((m) => m.isError !== true)
+          .map((m) => ({ role: m.role, content: m.content }))
       }
 
       try {
         // Run query
         if (state.querySettings.stream) {
-          await queryTextStream(queryParams, updateAssistantMessage)
+          let errorMessage = ''
+          await queryTextStream(queryParams, updateAssistantMessage, (error) => {
+            errorMessage += error
+          })
+          if (errorMessage) {
+            if (assistantMessage.content) {
+              errorMessage = assistantMessage.content + '\n' + errorMessage
+            }
+            updateAssistantMessage(errorMessage, true)
+          }
         } else {
           const response = await queryText(queryParams)
           updateAssistantMessage(response.response)
         }
       } catch (err) {
         // Handle error
-        updateAssistantMessage(`Error: Failed to get response\n${errorMessage(err)}`)
+        updateAssistantMessage(`Error: Failed to get response\n${errorMessage(err)}`, true)
       } finally {
         // Clear loading and add messages to state
         setIsLoading(false)
@@ -115,7 +131,11 @@ export default function RetrievalTesting() {
                   >
                     <div
                       className={`max-w-[80%] rounded-lg px-4 py-2 ${
-                        message.role === 'user' ? 'bg-primary text-primary-foreground' : 'bg-muted'
+                        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">{message.content}</pre>