From ee0700d288c41fb6ee76424705b9ff06ba7ef945 Mon Sep 17 00:00:00 2001
From: yangdx <yangdx@znipower.com>
Date: Wed, 12 Mar 2025 08:52:28 +0800
Subject: [PATCH] Ensure seamless toggling between selected labels and the
 complete graph

---
 .../src/components/graph/GraphLabels.tsx      | 15 +++++--
 lightrag_webui/src/hooks/useLightragGraph.tsx | 43 +++++++++++--------
 lightrag_webui/src/stores/graph.ts            |  4 +-
 3 files changed, 39 insertions(+), 23 deletions(-)

diff --git a/lightrag_webui/src/components/graph/GraphLabels.tsx b/lightrag_webui/src/components/graph/GraphLabels.tsx
index 474eb410..f2f50ea8 100644
--- a/lightrag_webui/src/components/graph/GraphLabels.tsx
+++ b/lightrag_webui/src/components/graph/GraphLabels.tsx
@@ -50,9 +50,17 @@ const GraphLabels = () => {
     [getSearchEngine]
   )
 
-  const setQueryLabel = useCallback((label: string) => {
-    if (label.startsWith('And ') && label.endsWith(' others')) return
-    useSettingsStore.getState().setQueryLabel(label)
+  const setQueryLabel = useCallback((newLabel: string) => {
+    if (newLabel.startsWith('And ') && newLabel.endsWith(' others')) return
+    
+    const currentLabel = useSettingsStore.getState().queryLabel
+    
+    // When selecting the same label (except '*'), switch to '*'
+    if (newLabel === currentLabel && newLabel !== '*') {
+      useSettingsStore.getState().setQueryLabel('*')
+    } else {
+      useSettingsStore.getState().setQueryLabel(newLabel)
+    }
   }, [])
 
   return (
@@ -70,6 +78,7 @@ const GraphLabels = () => {
       placeholder="Search labels..."
       value={label !== null ? label : ''}
       onChange={setQueryLabel}
+      clearable={false}  // Prevent clearing value on reselect
     />
   )
 }
diff --git a/lightrag_webui/src/hooks/useLightragGraph.tsx b/lightrag_webui/src/hooks/useLightragGraph.tsx
index 3951b783..162842f1 100644
--- a/lightrag_webui/src/hooks/useLightragGraph.tsx
+++ b/lightrag_webui/src/hooks/useLightragGraph.tsx
@@ -1,5 +1,5 @@
 import Graph, { DirectedGraph } from 'graphology'
-import { useCallback, useEffect } from 'react'
+import { useCallback, useEffect, useRef } from 'react'
 import { randomColor, errorMessage } from '@/lib/utils'
 import * as Constants from '@/lib/constants'
 import { useGraphStore, RawGraph } from '@/stores/graph'
@@ -161,9 +161,9 @@ const createSigmaGraph = (rawGraph: RawGraph | null) => {
   return graph
 }
 
-const lastQueryLabel = { label: '', maxQueryDepth: 0, minDegree: 0 }
-
 const useLightrangeGraph = () => {
+  // Use useRef to maintain lastQueryLabel state between renders
+  const lastQueryLabelRef = useRef({ label: '', maxQueryDepth: 0, minDegree: 0 })
   const queryLabel = useSettingsStore.use.queryLabel()
   const rawGraph = useGraphStore.use.rawGraph()
   const sigmaGraph = useGraphStore.use.sigmaGraph()
@@ -188,37 +188,44 @@ const useLightrangeGraph = () => {
     if (queryLabel) {
       // Always fetch data for "*" label
       // For other labels, only fetch when parameters change
-      const shouldUpdate = queryLabel === '*' ||
-        lastQueryLabel.label !== queryLabel ||
-        lastQueryLabel.maxQueryDepth !== maxQueryDepth ||
-        lastQueryLabel.minDegree !== minDegree;
+      const shouldUpdate = true;
 
       if (shouldUpdate) {
-        lastQueryLabel.label = queryLabel
-        lastQueryLabel.maxQueryDepth = maxQueryDepth
-        lastQueryLabel.minDegree = minDegree
+        lastQueryLabelRef.current = {
+          label: queryLabel,
+          maxQueryDepth,
+          minDegree
+        }
 
-        const state = useGraphStore.getState()
-        state.reset()
         fetchGraph(queryLabel, maxQueryDepth, minDegree).then((data) => {
-          // console.debug('Query label: ' + queryLabel)
-          state.setSigmaGraph(createSigmaGraph(data))
+          const state = useGraphStore.getState()
+          const newSigmaGraph = createSigmaGraph(data)
           data?.buildDynamicMap()
+          
+          // Update all graph data at once to minimize UI flicker
+          state.clearSelection()
+          state.setMoveToSelectedNode(false)
+          state.setSigmaGraph(newSigmaGraph)
           state.setRawGraph(data)
           
           // Extract labels from graph data
           if (data) {
-            const labelSet = new Set<string>(['*'])
+            const labelSet = new Set<string>();
             for (const node of data.nodes) {
               if (node.labels && Array.isArray(node.labels)) {
                 for (const label of node.labels) {
-                  labelSet.add(label)
+                  if (label !== '*') {  // filter out label "*"
+                    labelSet.add(label);
+                  }
                 }
               }
             }
-            state.setGraphLabels(Array.from(labelSet).sort())
+            // Put * on top of other labels
+            const sortedLabels = Array.from(labelSet).sort();
+            state.setGraphLabels(['*', ...sortedLabels]);
           } else {
-            state.setGraphLabels(['*'])
+            // Ensure * is there eventhough there is no graph data
+            state.setGraphLabels(['*']);
           }
         })
       }
diff --git a/lightrag_webui/src/stores/graph.ts b/lightrag_webui/src/stores/graph.ts
index 082c3161..8eab9049 100644
--- a/lightrag_webui/src/stores/graph.ts
+++ b/lightrag_webui/src/stores/graph.ts
@@ -93,7 +93,7 @@ const useGraphStoreBase = create<GraphState>()((set) => ({
 
   rawGraph: null,
   sigmaGraph: null,
-  graphLabels: [],
+  graphLabels: ['*'],
 
   setSelectedNode: (nodeId: string | null, moveToSelectedNode?: boolean) =>
     set({ selectedNode: nodeId, moveToSelectedNode }),
@@ -115,7 +115,7 @@ const useGraphStoreBase = create<GraphState>()((set) => ({
       focusedEdge: null,
       rawGraph: null,
       sigmaGraph: null,
-      graphLabels: [],
+      graphLabels: ['*'],
       moveToSelectedNode: false
     }),