ocr improved

LightRAG-main/lightrag_webui/src/features/RetrievalTesting.tsx

@@ -8,7 +8,7 @@ import { useSettingsStore } from '@/stores/settings'
 import { useDebounce } from '@/hooks/useDebounce'
 import QuerySettings from '@/components/retrieval/QuerySettings'
 import { ChatMessage, MessageWithError } from '@/components/retrieval/ChatMessage'
-import { EraserIcon, SendIcon } from 'lucide-react'
+import { EraserIcon, SendIcon, Square } from 'lucide-react'
 import { useTranslation } from 'react-i18next'
 import type { QueryMode } from '@/api/lightrag'
 
@@ -68,17 +68,30 @@ export default function RetrievalTesting() {
   const isReceivingResponseRef = useRef(false)
   const messagesEndRef = useRef<HTMLDivElement>(null)
   const messagesContainerRef = useRef<HTMLDivElement>(null)
+  const abortControllerRef = useRef<AbortController | null>(null)
 
   // Add cleanup effect for memory leak prevention
   useEffect(() => {
-    // Component cleanup - reset timer state to prevent memory leaks
+    // Component cleanup - reset timer state and abort any ongoing request
     return () => {
       if (thinkingStartTime.current) {
         thinkingStartTime.current = null;
       }
+      if (abortControllerRef.current) {
+        abortControllerRef.current.abort();
+        abortControllerRef.current = null;
+      }
     };
   }, []);
 
+  // Stop retrieval function
+  const stopRetrieval = useCallback(() => {
+    if (abortControllerRef.current) {
+      abortControllerRef.current.abort();
+      abortControllerRef.current = null;
+    }
+  }, []);
+
   // Scroll to bottom function - restored smooth scrolling with better handling
   const scrollToBottom = useCallback(() => {
     // Set flag to indicate this is a programmatic scroll
@@ -268,13 +281,23 @@ export default function RetrievalTesting() {
         ...(modeOverride ? { mode: modeOverride } : {})
       }
 
+      // Create abort controller for streaming cancellation
+      if (state.querySettings.stream) {
+        abortControllerRef.current = new AbortController();
+      }
+
       try {
         // Run query
         if (state.querySettings.stream) {
           let errorMessage = ''
-          await queryTextStream(queryParams, updateAssistantMessage, (error) => {
-            errorMessage += error
-          })
+          await queryTextStream(
+            queryParams,
+            updateAssistantMessage,
+            (error) => {
+              errorMessage += error
+            },
+            abortControllerRef.current?.signal
+          )
           if (errorMessage) {
             if (assistantMessage.content) {
               errorMessage = assistantMessage.content + '\n' + errorMessage
@@ -292,6 +315,8 @@ export default function RetrievalTesting() {
         // Clear loading and add messages to state
         setIsLoading(false)
         isReceivingResponseRef.current = false
+        // Clean up abort controller
+        abortControllerRef.current = null
 
         // Enhanced cleanup with error handling to prevent memory leaks
         try {
@@ -472,10 +497,22 @@ export default function RetrievalTesting() {
               <div className="absolute left-0 top-full mt-1 text-xs text-red-500">{inputError}</div>
             )}
           </div>
-          <Button type="submit" variant="default" disabled={isLoading} size="sm">
-            <SendIcon />
-            {t('retrievePanel.retrieval.send')}
-          </Button>
+          {isLoading ? (
+            <Button
+              type="button"
+              variant="destructive"
+              onClick={stopRetrieval}
+              size="sm"
+            >
+              <Square />
+              {t('retrievePanel.retrieval.stop')}
+            </Button>
+          ) : (
+            <Button type="submit" variant="default" disabled={isLoading} size="sm">
+              <SendIcon />
+              {t('retrievePanel.retrieval.send')}
+            </Button>
+          )}
         </form>
       </div>
       <QuerySettings />