ArnoChen commited on
Commit
ee3e1a1
·
1 Parent(s): 2a9f4c1

improve streaming error handling

Browse files
lightrag_webui/bun.lock CHANGED
@@ -419,9 +419,9 @@
419
 
420
  "@types/prismjs": ["@types/[email protected]", "", {}, "sha512-AUZTa7hQ2KY5L7AmtSiqxlhWxb4ina0yd8hNbl4TWuqnv/pFP0nDMb3YrfSBf4hJVGLh2YEIBfKaBW/9UEl6IQ=="],
421
 
422
- "@types/react": ["@types/[email protected].8", "", { "dependencies": { "csstype": "^3.0.2" } }, "sha512-9P/o1IGdfmQxrujGbIMDyYaaCykhLKc0NGCtYcECNUr9UAaDe4gwvV9bR6tvd5Br1SG0j+PBpbKr2UYY8CwqSw=="],
423
 
424
- "@types/react-dom": ["@types/[email protected].3", "", { "peerDependencies": { "@types/react": "^19.0.0" } }, "sha512-0Knk+HJiMP/qOZgMyNFamlIjw9OFCsyC2ZbigmEEyXXixgre6IQpm/4V+r3qH4GC1JPvRJKInw+on2rV6YZLeA=="],
425
 
426
  "@types/react-transition-group": ["@types/[email protected]", "", { "peerDependencies": { "@types/react": "*" } }, "sha512-8TV6R3h2j7a91c+1DXdJi3Syo69zzIZbz7Lg5tORM5LEJG7X/E6a1V3drRyBRZq7/utz7A+c4OgYLiLcYGHG6w=="],
427
 
 
419
 
420
  "@types/prismjs": ["@types/[email protected]", "", {}, "sha512-AUZTa7hQ2KY5L7AmtSiqxlhWxb4ina0yd8hNbl4TWuqnv/pFP0nDMb3YrfSBf4hJVGLh2YEIBfKaBW/9UEl6IQ=="],
421
 
422
+ "@types/react": ["@types/[email protected].10", "", { "dependencies": { "csstype": "^3.0.2" } }, "sha512-JuRQ9KXLEjaUNjTWpzuR231Z2WpIwczOkBEIvbHNCzQefFIT0L8IqE6NV6ULLyC1SI/i234JnDoMkfg+RjQj2g=="],
423
 
424
+ "@types/react-dom": ["@types/[email protected].4", "", { "peerDependencies": { "@types/react": "^19.0.0" } }, "sha512-4fSQ8vWFkg+TGhePfUzVmat3eC14TXYSsiiDSLI0dVLsrm9gZFABjPy/Qu6TKgl1tq1Bu1yDsuQgY3A3DOjCcg=="],
425
 
426
  "@types/react-transition-group": ["@types/[email protected]", "", { "peerDependencies": { "@types/react": "*" } }, "sha512-8TV6R3h2j7a91c+1DXdJi3Syo69zzIZbz7Lg5tORM5LEJG7X/E6a1V3drRyBRZq7/utz7A+c4OgYLiLcYGHG6w=="],
427
 
lightrag_webui/package.json CHANGED
@@ -59,8 +59,8 @@
59
  "@tailwindcss/vite": "^4.0.6",
60
  "@types/bun": "^1.2.2",
61
  "@types/node": "^22.13.4",
62
- "@types/react": "^19.0.8",
63
- "@types/react-dom": "^19.0.3",
64
  "@types/seedrandom": "^3.0.8",
65
  "@vitejs/plugin-react-swc": "^3.8.0",
66
  "eslint": "^9.20.1",
 
59
  "@tailwindcss/vite": "^4.0.6",
60
  "@types/bun": "^1.2.2",
61
  "@types/node": "^22.13.4",
62
+ "@types/react": "^19.0.10",
63
+ "@types/react-dom": "^19.0.4",
64
  "@types/seedrandom": "^3.0.8",
65
  "@vitejs/plugin-react-swc": "^3.8.0",
66
  "eslint": "^9.20.1",
lightrag_webui/src/api/lightrag.ts CHANGED
@@ -212,38 +212,42 @@ export const queryTextStream = async (
212
  ) => {
213
  try {
214
  let buffer = ''
215
- await axiosInstance.post('/query/stream', request, {
216
- responseType: 'text',
217
- headers: {
218
- Accept: 'application/x-ndjson'
219
- },
220
- transformResponse: [
221
- (data: string) => {
222
- // Accumulate the data and process complete lines
223
- buffer += data
224
- const lines = buffer.split('\n')
225
- // Keep the last potentially incomplete line in the buffer
226
- buffer = lines.pop() || ''
227
-
228
- for (const line of lines) {
229
- if (line.trim()) {
230
- try {
231
- const parsed = JSON.parse(line)
232
- if (parsed.response) {
233
- onChunk(parsed.response)
234
- } else if (parsed.error && onError) {
235
- onError(parsed.error)
 
 
 
 
 
236
  }
237
- } catch (e) {
238
- console.error('Error parsing stream chunk:', e)
239
- if (onError) onError('Error parsing server response')
240
  }
241
  }
 
242
  }
243
- return data
244
- }
245
- ]
246
- })
 
247
 
248
  // Process any remaining data in the buffer
249
  if (buffer.trim()) {
@@ -266,11 +270,13 @@ export const queryTextStream = async (
266
  }
267
  }
268
 
269
- export const insertText = async (
270
- text: string,
271
- description?: string
272
- ): Promise<DocActionResponse> => {
273
- const response = await axiosInstance.post('/documents/text', { text, description })
 
 
274
  return response.data
275
  }
276
 
 
212
  ) => {
213
  try {
214
  let buffer = ''
215
+ await axiosInstance
216
+ .post('/query/stream', request, {
217
+ responseType: 'text',
218
+ headers: {
219
+ Accept: 'application/x-ndjson'
220
+ },
221
+ transformResponse: [
222
+ (data: string) => {
223
+ // Accumulate the data and process complete lines
224
+ buffer += data
225
+ const lines = buffer.split('\n')
226
+ // Keep the last potentially incomplete line in the buffer
227
+ buffer = lines.pop() || ''
228
+
229
+ for (const line of lines) {
230
+ if (line.trim()) {
231
+ try {
232
+ const parsed = JSON.parse(line)
233
+ if (parsed.response) {
234
+ onChunk(parsed.response)
235
+ } else if (parsed.error && onError) {
236
+ onError(parsed.error)
237
+ }
238
+ } catch (e) {
239
+ console.error('Error parsing stream chunk:', e)
240
+ if (onError) onError('Error parsing server response')
241
  }
 
 
 
242
  }
243
  }
244
+ return data
245
  }
246
+ ]
247
+ })
248
+ .catch((error) => {
249
+ if (onError) onError(errorMessage(error))
250
+ })
251
 
252
  // Process any remaining data in the buffer
253
  if (buffer.trim()) {
 
270
  }
271
  }
272
 
273
+ export const insertText = async (text: string): Promise<DocActionResponse> => {
274
+ const response = await axiosInstance.post('/documents/text', { text })
275
+ return response.data
276
+ }
277
+
278
+ export const insertTexts = async (texts: string[]): Promise<DocActionResponse> => {
279
+ const response = await axiosInstance.post('/documents/texts', { texts })
280
  return response.data
281
  }
282
 
lightrag_webui/src/features/DocumentManager.tsx CHANGED
@@ -29,7 +29,16 @@ export default function DocumentManager() {
29
  try {
30
  const docs = await getDocuments()
31
  if (docs && docs.statuses) {
32
- setDocs(docs)
 
 
 
 
 
 
 
 
 
33
  // console.log(docs)
34
  } else {
35
  setDocs(null)
 
29
  try {
30
  const docs = await getDocuments()
31
  if (docs && docs.statuses) {
32
+ // compose all documents count
33
+ const numDocuments = Object.values(docs.statuses).reduce(
34
+ (acc, status) => acc + status.length,
35
+ 0
36
+ )
37
+ if (numDocuments > 0) {
38
+ setDocs(docs)
39
+ } else {
40
+ setDocs(null)
41
+ }
42
  // console.log(docs)
43
  } else {
44
  setDocs(null)
lightrag_webui/src/features/RetrievalTesting.tsx CHANGED
@@ -1,7 +1,7 @@
1
  import Input from '@/components/ui/Input'
2
  import Button from '@/components/ui/Button'
3
  import { useCallback, useEffect, useRef, useState } from 'react'
4
- import { queryText, queryTextStream, Message } from '@/api/lightrag'
5
  import { errorMessage } from '@/lib/utils'
6
  import { useSettingsStore } from '@/stores/settings'
7
  import { useDebounce } from '@/hooks/useDebounce'
@@ -9,6 +9,10 @@ import QuerySettings from '@/components/retrieval/QuerySettings'
9
 
10
  import { EraserIcon, SendIcon, LoaderIcon } from 'lucide-react'
11
 
 
 
 
 
12
  export default function RetrievalTesting() {
13
  const [messages, setMessages] = useState<Message[]>(
14
  () => useSettingsStore.getState().retrievalHistory || []
@@ -47,13 +51,14 @@ export default function RetrievalTesting() {
47
  setIsLoading(true)
48
 
49
  // Create a function to update the assistant's message
50
- const updateAssistantMessage = (chunk: string) => {
51
  assistantMessage.content += chunk
52
  setMessages((prev) => {
53
  const newMessages = [...prev]
54
  const lastMessage = newMessages[newMessages.length - 1]
55
  if (lastMessage.role === 'assistant') {
56
  lastMessage.content = assistantMessage.content
 
57
  }
58
  return newMessages
59
  })
@@ -65,19 +70,30 @@ export default function RetrievalTesting() {
65
  ...state.querySettings,
66
  query: userMessage.content,
67
  conversation_history: prevMessages
 
 
68
  }
69
 
70
  try {
71
  // Run query
72
  if (state.querySettings.stream) {
73
- await queryTextStream(queryParams, updateAssistantMessage)
 
 
 
 
 
 
 
 
 
74
  } else {
75
  const response = await queryText(queryParams)
76
  updateAssistantMessage(response.response)
77
  }
78
  } catch (err) {
79
  // Handle error
80
- updateAssistantMessage(`Error: Failed to get response\n${errorMessage(err)}`)
81
  } finally {
82
  // Clear loading and add messages to state
83
  setIsLoading(false)
@@ -115,7 +131,11 @@ export default function RetrievalTesting() {
115
  >
116
  <div
117
  className={`max-w-[80%] rounded-lg px-4 py-2 ${
118
- message.role === 'user' ? 'bg-primary text-primary-foreground' : 'bg-muted'
 
 
 
 
119
  }`}
120
  >
121
  <pre className="break-words whitespace-pre-wrap">{message.content}</pre>
 
1
  import Input from '@/components/ui/Input'
2
  import Button from '@/components/ui/Button'
3
  import { useCallback, useEffect, useRef, useState } from 'react'
4
+ import { queryText, queryTextStream, Message as ChatMessage } from '@/api/lightrag'
5
  import { errorMessage } from '@/lib/utils'
6
  import { useSettingsStore } from '@/stores/settings'
7
  import { useDebounce } from '@/hooks/useDebounce'
 
9
 
10
  import { EraserIcon, SendIcon, LoaderIcon } from 'lucide-react'
11
 
12
+ type Message = ChatMessage & {
13
+ isError?: boolean
14
+ }
15
+
16
  export default function RetrievalTesting() {
17
  const [messages, setMessages] = useState<Message[]>(
18
  () => useSettingsStore.getState().retrievalHistory || []
 
51
  setIsLoading(true)
52
 
53
  // Create a function to update the assistant's message
54
+ const updateAssistantMessage = (chunk: string, isError?: boolean) => {
55
  assistantMessage.content += chunk
56
  setMessages((prev) => {
57
  const newMessages = [...prev]
58
  const lastMessage = newMessages[newMessages.length - 1]
59
  if (lastMessage.role === 'assistant') {
60
  lastMessage.content = assistantMessage.content
61
+ lastMessage.isError = isError
62
  }
63
  return newMessages
64
  })
 
70
  ...state.querySettings,
71
  query: userMessage.content,
72
  conversation_history: prevMessages
73
+ .filter((m) => m.isError !== true)
74
+ .map((m) => ({ role: m.role, content: m.content }))
75
  }
76
 
77
  try {
78
  // Run query
79
  if (state.querySettings.stream) {
80
+ let errorMessage = ''
81
+ await queryTextStream(queryParams, updateAssistantMessage, (error) => {
82
+ errorMessage += error
83
+ })
84
+ if (errorMessage) {
85
+ if (assistantMessage.content) {
86
+ errorMessage = assistantMessage.content + '\n' + errorMessage
87
+ }
88
+ updateAssistantMessage(errorMessage, true)
89
+ }
90
  } else {
91
  const response = await queryText(queryParams)
92
  updateAssistantMessage(response.response)
93
  }
94
  } catch (err) {
95
  // Handle error
96
+ updateAssistantMessage(`Error: Failed to get response\n${errorMessage(err)}`, true)
97
  } finally {
98
  // Clear loading and add messages to state
99
  setIsLoading(false)
 
131
  >
132
  <div
133
  className={`max-w-[80%] rounded-lg px-4 py-2 ${
134
+ message.role === 'user'
135
+ ? 'bg-primary text-primary-foreground'
136
+ : message.isError
137
+ ? 'bg-red-100 text-red-600 dark:bg-red-950 dark:text-red-400'
138
+ : 'bg-muted'
139
  }`}
140
  >
141
  <pre className="break-words whitespace-pre-wrap">{message.content}</pre>