From e083ebb95adea18554f746f31f7492ce2dd77f06 Mon Sep 17 00:00:00 2001 From: ArnoChen Date: Tue, 18 Feb 2025 00:30:51 +0800 Subject: [PATCH] improve streaming error handling --- lightrag_webui/bun.lock | 4 +- lightrag_webui/package.json | 4 +- lightrag_webui/src/api/lightrag.ts | 70 ++++++++++--------- .../src/features/DocumentManager.tsx | 11 ++- .../src/features/RetrievalTesting.tsx | 30 ++++++-- 5 files changed, 77 insertions(+), 42 deletions(-) diff --git a/lightrag_webui/bun.lock b/lightrag_webui/bun.lock index 570fa3a6..ca1b0e6a 100644 --- a/lightrag_webui/bun.lock +++ b/lightrag_webui/bun.lock @@ -419,9 +419,9 @@ "@types/prismjs": ["@types/prismjs@1.26.5", "", {}, "sha512-AUZTa7hQ2KY5L7AmtSiqxlhWxb4ina0yd8hNbl4TWuqnv/pFP0nDMb3YrfSBf4hJVGLh2YEIBfKaBW/9UEl6IQ=="], - "@types/react": ["@types/react@19.0.8", "", { "dependencies": { "csstype": "^3.0.2" } }, "sha512-9P/o1IGdfmQxrujGbIMDyYaaCykhLKc0NGCtYcECNUr9UAaDe4gwvV9bR6tvd5Br1SG0j+PBpbKr2UYY8CwqSw=="], + "@types/react": ["@types/react@19.0.10", "", { "dependencies": { "csstype": "^3.0.2" } }, "sha512-JuRQ9KXLEjaUNjTWpzuR231Z2WpIwczOkBEIvbHNCzQefFIT0L8IqE6NV6ULLyC1SI/i234JnDoMkfg+RjQj2g=="], - "@types/react-dom": ["@types/react-dom@19.0.3", "", { "peerDependencies": { "@types/react": "^19.0.0" } }, "sha512-0Knk+HJiMP/qOZgMyNFamlIjw9OFCsyC2ZbigmEEyXXixgre6IQpm/4V+r3qH4GC1JPvRJKInw+on2rV6YZLeA=="], + "@types/react-dom": ["@types/react-dom@19.0.4", "", { "peerDependencies": { "@types/react": "^19.0.0" } }, "sha512-4fSQ8vWFkg+TGhePfUzVmat3eC14TXYSsiiDSLI0dVLsrm9gZFABjPy/Qu6TKgl1tq1Bu1yDsuQgY3A3DOjCcg=="], "@types/react-transition-group": ["@types/react-transition-group@4.4.12", "", { "peerDependencies": { "@types/react": "*" } }, "sha512-8TV6R3h2j7a91c+1DXdJi3Syo69zzIZbz7Lg5tORM5LEJG7X/E6a1V3drRyBRZq7/utz7A+c4OgYLiLcYGHG6w=="], diff --git a/lightrag_webui/package.json b/lightrag_webui/package.json index 26991d9c..3d9d5fe7 100644 --- a/lightrag_webui/package.json +++ b/lightrag_webui/package.json @@ -59,8 +59,8 @@ "@tailwindcss/vite": "^4.0.6", "@types/bun": "^1.2.2", "@types/node": "^22.13.4", - "@types/react": "^19.0.8", - "@types/react-dom": "^19.0.3", + "@types/react": "^19.0.10", + "@types/react-dom": "^19.0.4", "@types/seedrandom": "^3.0.8", "@vitejs/plugin-react-swc": "^3.8.0", "eslint": "^9.20.1", diff --git a/lightrag_webui/src/api/lightrag.ts b/lightrag_webui/src/api/lightrag.ts index 6308ef62..1de79898 100644 --- a/lightrag_webui/src/api/lightrag.ts +++ b/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 => { - const response = await axiosInstance.post('/documents/text', { text, description }) +export const insertText = async (text: string): Promise => { + const response = await axiosInstance.post('/documents/text', { text }) + return response.data +} + +export const insertTexts = async (texts: string[]): Promise => { + const response = await axiosInstance.post('/documents/texts', { texts }) return response.data } diff --git a/lightrag_webui/src/features/DocumentManager.tsx b/lightrag_webui/src/features/DocumentManager.tsx index 0aa068d4..30e55d48 100644 --- a/lightrag_webui/src/features/DocumentManager.tsx +++ b/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) diff --git a/lightrag_webui/src/features/RetrievalTesting.tsx b/lightrag_webui/src/features/RetrievalTesting.tsx index e293322b..7fc46abf 100644 --- a/lightrag_webui/src/features/RetrievalTesting.tsx +++ b/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( () => 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() { > 
{message.content}