From df0594852a7422991472a6375ad5edc4c0fa05f8 Mon Sep 17 00:00:00 2001
From: yangdx <yangdx@znipower.com>
Date: Sun, 6 Apr 2025 14:56:55 +0800
Subject: [PATCH] Set queryLabel after query on page first load

---
 .../src/components/graph/GraphLabels.tsx      | 34 +++++++++++--------
 lightrag_webui/src/hooks/useLightragGraph.tsx | 12 +++++++
 2 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/lightrag_webui/src/components/graph/GraphLabels.tsx b/lightrag_webui/src/components/graph/GraphLabels.tsx
index f43bd955..3bfe0024 100644
--- a/lightrag_webui/src/components/graph/GraphLabels.tsx
+++ b/lightrag_webui/src/components/graph/GraphLabels.tsx
@@ -12,7 +12,8 @@ const GraphLabels = () => {
   const { t } = useTranslation()
   const label = useSettingsStore.use.queryLabel()
   const allDatabaseLabels = useGraphStore.use.allDatabaseLabels()
-
+  const labelsFetchAttempted = useGraphStore.use.labelsFetchAttempted()
+  
   // Remove initial label fetch effect as it's now handled by fetchGraph based on lastSuccessfulQueryLabel
 
   const getSearchEngine = useCallback(() => {
@@ -56,22 +57,25 @@ const GraphLabels = () => {
     [getSearchEngine]
   )
 
-  // Validate if current queryLabel exists in allDatabaseLabels
+  // Show queryLabel validation status
   useEffect(() => {
-    // Only update label when all conditions are met:
-    // 1. allDatabaseLabels is loaded (length > 1, as it has at least '*' by default)
-    // 2. Current label is not the default '*'
-    // 3. Current label doesn't exist in allDatabaseLabels
-    if (
-      allDatabaseLabels.length > 1 &&
-      label &&
-      label !== '*' &&
-      !allDatabaseLabels.includes(label)
-    ) {
-      console.log(`Label "${label}" not found in available labels, resetting to default`);
-      useSettingsStore.getState().setQueryLabel('*');
+    
+    if (labelsFetchAttempted) {      
+      if (allDatabaseLabels.length > 1) {
+        if (label && label !== '*' && !allDatabaseLabels.includes(label)) {
+          console.log(`Label "${label}" not in available labels`);
+          // useSettingsStore.getState().setQueryLabel('*');
+        } else {
+          console.log(`Label "${label}" is valid`);
+        }
+      } else if (allDatabaseLabels.length <= 1 && label && label !== '*') {
+        console.log('Available labels list is empty');
+        // useSettingsStore.getState().setQueryLabel('');
+      }
+      useGraphStore.getState().setLabelsFetchAttempted(false)
     }
-  }, [allDatabaseLabels, label]);
+
+  }, [allDatabaseLabels, label, labelsFetchAttempted]);
 
   const handleRefresh = useCallback(() => {
     // Reset fetch status flags
diff --git a/lightrag_webui/src/hooks/useLightragGraph.tsx b/lightrag_webui/src/hooks/useLightragGraph.tsx
index 39521fff..af6810e7 100644
--- a/lightrag_webui/src/hooks/useLightragGraph.tsx
+++ b/lightrag_webui/src/hooks/useLightragGraph.tsx
@@ -119,6 +119,10 @@ const fetchGraph = async (label: string, maxDepth: number, maxNodes: number) =>
       // Continue with graph fetch even if labels fetch fails
     }
   }
+  
+  // Trigger GraphLabels component to check if the label is valid
+  // console.log('Setting labelsFetchAttempted to true');
+  useGraphStore.getState().setLabelsFetchAttempted(true)
 
   // If label is empty, use default label '*'
   const queryLabel = label || '*';
@@ -339,6 +343,7 @@ const useLightrangeGraph = () => {
     }
 
     // Only fetch data when graphDataFetchAttempted is false (avoids re-fetching on vite dev mode)
+    // GraphDataFetchAttempted must set to false when queryLabel is changed
     if (!isFetching && !useGraphStore.getState().graphDataFetchAttempted) {
       // Set flags
       fetchInProgressRef.current = true
@@ -445,6 +450,13 @@ const useLightrangeGraph = () => {
           state.setRawGraph(data);
           state.setGraphIsEmpty(false);
 
+          // ensusre GraphLabels show the current label on first load
+          const lastSuccessfulQueryLabel = useGraphStore.getState().lastSuccessfulQueryLabel;
+          if (!lastSuccessfulQueryLabel){
+            useSettingsStore.getState().setQueryLabel('');
+            useSettingsStore.getState().setQueryLabel(queryLabel);
+            console.log('Set queryLabel after query on page first load');
+          }
           // Update last successful query label
           state.setLastSuccessfulQueryLabel(currentQueryLabel);