From d418ceee82c4ee70a0e2c78d2b13f3afd96ad1c7 Mon Sep 17 00:00:00 2001
From: ArnoChen <arnochen101@gmail.com>
Date: Tue, 11 Feb 2025 06:48:04 +0800
Subject: [PATCH] add label filter

---
 lightrag/api/graph_viewer_webui/bun.lock      |   1 +
 lightrag/api/graph_viewer_webui/package.json  |   1 +
 lightrag/api/graph_viewer_webui/src/App.tsx   |  18 +-
 .../graph_viewer_webui/src/GraphViewer.tsx    |   5 +-
 .../graph_viewer_webui/src/api/lightrag.ts    |   3 +-
 .../src/components/BackendMessageAlert.tsx    |  22 --
 .../src/components/GraphLabels.tsx            |  87 +++++++
 .../src/components/GraphSearch.tsx            |  78 ++++--
 .../src/components/MessageAlert.tsx           |  38 +++
 .../src/components/ui/AsyncSearch.tsx         | 243 ++++++++++++++++++
 .../src/components/ui/AsyncSelect.tsx         | 160 +++++++-----
 .../src/hooks/useLightragGraph.tsx            |  34 ++-
 .../graph_viewer_webui/src/lib/constants.ts   |   4 +
 .../api/graph_viewer_webui/src/lib/utils.ts   |   4 +
 .../graph_viewer_webui/src/stores/graph.ts    |   9 -
 .../graph_viewer_webui/src/stores/settings.ts |  19 +-
 16 files changed, 583 insertions(+), 143 deletions(-)
 delete mode 100644 lightrag/api/graph_viewer_webui/src/components/BackendMessageAlert.tsx
 create mode 100644 lightrag/api/graph_viewer_webui/src/components/GraphLabels.tsx
 create mode 100644 lightrag/api/graph_viewer_webui/src/components/MessageAlert.tsx
 create mode 100644 lightrag/api/graph_viewer_webui/src/components/ui/AsyncSearch.tsx

diff --git a/lightrag/api/graph_viewer_webui/bun.lock b/lightrag/api/graph_viewer_webui/bun.lock
index fcd1a18d..21073f15 100644
--- a/lightrag/api/graph_viewer_webui/bun.lock
+++ b/lightrag/api/graph_viewer_webui/bun.lock
@@ -27,6 +27,7 @@
         "graphology": "^0.26.0",
         "graphology-generators": "^0.11.2",
         "lucide-react": "^0.475.0",
+        "minisearch": "^7.1.1",
         "react": "^19.0.0",
         "react-dom": "^19.0.0",
         "seedrandom": "^3.0.5",
diff --git a/lightrag/api/graph_viewer_webui/package.json b/lightrag/api/graph_viewer_webui/package.json
index cd488cec..1057f3f6 100644
--- a/lightrag/api/graph_viewer_webui/package.json
+++ b/lightrag/api/graph_viewer_webui/package.json
@@ -33,6 +33,7 @@
     "graphology": "^0.26.0",
     "graphology-generators": "^0.11.2",
     "lucide-react": "^0.475.0",
+    "minisearch": "^7.1.1",
     "react": "^19.0.0",
     "react-dom": "^19.0.0",
     "seedrandom": "^3.0.5",
diff --git a/lightrag/api/graph_viewer_webui/src/App.tsx b/lightrag/api/graph_viewer_webui/src/App.tsx
index d859c239..11a5761a 100644
--- a/lightrag/api/graph_viewer_webui/src/App.tsx
+++ b/lightrag/api/graph_viewer_webui/src/App.tsx
@@ -1,18 +1,28 @@
 import ThemeProvider from '@/components/ThemeProvider'
-import BackendMessageAlert from '@/components/BackendMessageAlert'
+import MessageAlert from '@/components/MessageAlert'
 import { GraphViewer } from '@/GraphViewer'
 import { cn } from '@/lib/utils'
+import { healthCheckInterval } from '@/lib/constants'
 import { useBackendState } from '@/stores/state'
+import { useEffect } from 'react'
 
 function App() {
-  const health = useBackendState.use.health()
+  const message = useBackendState.use.message()
+
+  // health check
+  useEffect(() => {
+    const interval = setInterval(async () => {
+      await useBackendState.getState().check()
+    }, healthCheckInterval * 1000)
+    return () => clearInterval(interval)
+  }, [])
 
   return (
     <ThemeProvider>
-      <div className={cn('h-screen w-screen', !health && 'pointer-events-none')}>
+      <div className={cn('h-screen w-screen', message !== null && 'pointer-events-none')}>
         <GraphViewer />
       </div>
-      {!health && <BackendMessageAlert />}
+      {message !== null && <MessageAlert />}
     </ThemeProvider>
   )
 }
diff --git a/lightrag/api/graph_viewer_webui/src/GraphViewer.tsx b/lightrag/api/graph_viewer_webui/src/GraphViewer.tsx
index 8cd8d561..993b4b3c 100644
--- a/lightrag/api/graph_viewer_webui/src/GraphViewer.tsx
+++ b/lightrag/api/graph_viewer_webui/src/GraphViewer.tsx
@@ -15,6 +15,7 @@ import ZoomControl from '@/components/ZoomControl'
 import FullScreenControl from '@/components/FullScreenControl'
 import Settings from '@/components/Settings'
 import GraphSearch from '@/components/GraphSearch'
+import GraphLabels from '@/components/GraphLabels'
 import PropertiesView from '@/components/PropertiesView'
 
 import { useSettingsStore } from '@/stores/settings'
@@ -144,9 +145,9 @@ export const GraphViewer = () => {
 
       <FocusOnNode node={autoFocusedNode} move={moveToSelectedNode} />
 
-      <div className="absolute top-2 left-2">
+      <div className="absolute top-2 left-2 flex items-start gap-2">
+        <GraphLabels />
         <GraphSearch
-          type="nodes"
           value={searchInitSelectedNode}
           onFocus={onSearchFocus}
           onChange={onSearchSelect}
diff --git a/lightrag/api/graph_viewer_webui/src/api/lightrag.ts b/lightrag/api/graph_viewer_webui/src/api/lightrag.ts
index 648e26d9..58e1b40e 100644
--- a/lightrag/api/graph_viewer_webui/src/api/lightrag.ts
+++ b/lightrag/api/graph_viewer_webui/src/api/lightrag.ts
@@ -1,4 +1,5 @@
 import { backendBaseUrl } from '@/lib/constants'
+import { errorMessage } from '@/lib/utils'
 
 export type LightragNodeType = {
   id: string
@@ -81,7 +82,7 @@ export const checkHealth = async (): Promise<
   } catch (e) {
     return {
       status: 'error',
-      message: e instanceof Error ? e.message : `${e}`
+      message: errorMessage(e)
     }
   }
 }
diff --git a/lightrag/api/graph_viewer_webui/src/components/BackendMessageAlert.tsx b/lightrag/api/graph_viewer_webui/src/components/BackendMessageAlert.tsx
deleted file mode 100644
index f90799f6..00000000
--- a/lightrag/api/graph_viewer_webui/src/components/BackendMessageAlert.tsx
+++ /dev/null
@@ -1,22 +0,0 @@
-import { Alert, AlertDescription, AlertTitle } from '@/components/ui/Alert'
-import { useBackendState } from '@/stores/state'
-import { AlertCircle } from 'lucide-react'
-
-const BackendMessageAlert = () => {
-  const health = useBackendState.use.health()
-  const message = useBackendState.use.message()
-  const messageTitle = useBackendState.use.messageTitle()
-
-  return (
-    <Alert
-      variant={health ? 'default' : 'destructive'}
-      className="absolute top-1/2 left-1/2 w-auto -translate-x-1/2 -translate-y-1/2 transform"
-    >
-      {!health && <AlertCircle className="h-4 w-4" />}
-      <AlertTitle>{messageTitle}</AlertTitle>
-      <AlertDescription>{message}</AlertDescription>
-    </Alert>
-  )
-}
-
-export default BackendMessageAlert
diff --git a/lightrag/api/graph_viewer_webui/src/components/GraphLabels.tsx b/lightrag/api/graph_viewer_webui/src/components/GraphLabels.tsx
new file mode 100644
index 00000000..b3c325ad
--- /dev/null
+++ b/lightrag/api/graph_viewer_webui/src/components/GraphLabels.tsx
@@ -0,0 +1,87 @@
+import { useCallback, useState } from 'react'
+import { AsyncSelect } from '@/components/ui/AsyncSelect'
+import { getGraphLabels } from '@/api/lightrag'
+import { useSettingsStore } from '@/stores/settings'
+import MiniSearch from 'minisearch'
+
+const GraphLabels = () => {
+  const label = useSettingsStore.use.queryLabel()
+  const [labels, setLabels] = useState<{
+    labels: string[]
+    searchEngine: MiniSearch | null
+  }>({
+    labels: [],
+    searchEngine: null
+  })
+  const [fetched, setFetched] = useState(false)
+
+  const fetchData = useCallback(
+    async (query?: string): Promise<string[]> => {
+      let _labels = labels.labels
+      let _searchEngine = labels.searchEngine
+
+      if (!fetched || !_searchEngine) {
+        _labels = ['*'].concat(await getGraphLabels())
+
+        // Ensure query label exists
+        if (!_labels.includes(useSettingsStore.getState().queryLabel)) {
+          useSettingsStore.getState().setQueryLabel(_labels[0])
+        }
+
+        // Create search engine
+        _searchEngine = new MiniSearch({
+          idField: 'id',
+          fields: ['value'],
+          searchOptions: {
+            prefix: true,
+            fuzzy: 0.2,
+            boost: {
+              label: 2
+            }
+          }
+        })
+
+        // Add documents
+        const documents = _labels.map((str, index) => ({ id: index, value: str }))
+        _searchEngine.addAll(documents)
+
+        setLabels({
+          labels: _labels,
+          searchEngine: _searchEngine
+        })
+        setFetched(true)
+      }
+      if (!query) {
+        return _labels
+      }
+
+      // Search labels
+      return _searchEngine.search(query).map((result) => _labels[result.id])
+    },
+    [labels, fetched, setLabels, setFetched]
+  )
+
+  const setQueryLabel = useCallback((label: string) => {
+    useSettingsStore.getState().setQueryLabel(label)
+  }, [])
+
+  return (
+    <AsyncSelect<string>
+      className="ml-2"
+      triggerClassName="max-h-8"
+      searchInputClassName="max-h-8"
+      triggerTooltip="Select query label"
+      fetcher={fetchData}
+      renderOption={(item) => <div>{item}</div>}
+      getOptionValue={(item) => item}
+      getDisplayValue={(item) => <div>{item}</div>}
+      notFound={<div className="py-6 text-center text-sm">No labels found</div>}
+      label="Label"
+      placeholder="Search labels..."
+      value={label !== null ? label : ''}
+      onChange={setQueryLabel}
+    />
+  )
+}
+
+export default GraphLabels
diff --git a/lightrag/api/graph_viewer_webui/src/components/GraphSearch.tsx b/lightrag/api/graph_viewer_webui/src/components/GraphSearch.tsx
index b6267f91..16d4c62d 100644
--- a/lightrag/api/graph_viewer_webui/src/components/GraphSearch.tsx
+++ b/lightrag/api/graph_viewer_webui/src/components/GraphSearch.tsx
@@ -1,14 +1,14 @@
-import { FC, useCallback } from 'react'
+import { FC, useCallback, useMemo } from 'react'
 import {
   EdgeById,
   NodeById,
-  useGraphSearch,
   GraphSearchInputProps,
-  GraphSearchContextProvider,
   GraphSearchContextProviderProps
 } from '@react-sigma/graph-search'
-import { AsyncSelect } from '@/components/ui/AsyncSelect'
+import { AsyncSearch } from '@/components/ui/AsyncSearch'
 import { searchResultLimit } from '@/lib/constants'
+import { useGraphStore } from '@/stores/graph'
+import MiniSearch from 'minisearch'
 
 interface OptionItem {
   id: string
@@ -27,6 +27,10 @@ function OptionComponent(item: OptionItem) {
 }
 
 const messageId = '__message_item'
+const lastGraph: any = {
+  graph: null,
+  searchEngine: null
+}
 
 /**
  * Component thats display the search input.
@@ -34,15 +38,44 @@ const messageId = '__message_item'
 export const GraphSearchInput = ({
   onChange,
   onFocus,
-  type,
   value
 }: {
   onChange: GraphSearchInputProps['onChange']
   onFocus?: GraphSearchInputProps['onFocus']
-  type?: GraphSearchInputProps['type']
   value?: GraphSearchInputProps['value']
 }) => {
-  const { search } = useGraphSearch()
+  const graph = useGraphStore.use.sigmaGraph()
+
+  const search = useMemo(() => {
+    if (lastGraph.graph == graph) {
+      return lastGraph.searchEngine
+    }
+    if (!graph || graph.nodes().length == 0) return
+
+    lastGraph.graph = graph
+
+    const searchEngine = new MiniSearch({
+      idField: 'id',
+      fields: ['label'],
+      searchOptions: {
+        prefix: true,
+        fuzzy: 0.2,
+        boost: {
+          label: 2
+        }
+      }
+    })
+
+    // Add documents
+    const documents = graph.nodes().map((id: string) => ({
+      id: id,
+      label: graph.getNodeAttribute(id, 'label')
+    }))
+    searchEngine.addAll(documents)
+
+    lastGraph.searchEngine = searchEngine
+    return searchEngine
+  }, [graph])
 
   /**
    * Loading the options while the user is typing.
@@ -50,8 +83,11 @@ export const GraphSearchInput = ({
   const loadOptions = useCallback(
     async (query?: string): Promise<OptionItem[]> => {
       if (onFocus) onFocus(null)
-      if (!query) return []
-      const result = (await search(query, type)) as OptionItem[]
+      if (!query || !search) return []
+      const result: OptionItem[] = search.search(query).map((result) => ({
+        id: result.id,
+        type: 'nodes'
+      }))
 
       // prettier-ignore
       return result.length <= searchResultLimit
@@ -65,25 +101,24 @@ export const GraphSearchInput = ({
           }
         ]
     },
-    [type, search, onFocus]
+    [search, onFocus]
   )
 
   return (
-    <AsyncSelect
-      className="bg-background/60 w-52 rounded-xl border-1 opacity-60 backdrop-blur-lg transition-opacity hover:opacity-100"
+    <AsyncSearch
+      className="bg-background/60 w-24 rounded-xl border-1 opacity-60 backdrop-blur-lg transition-all hover:w-fit hover:opacity-100"
       fetcher={loadOptions}
       renderOption={OptionComponent}
       getOptionValue={(item) => item.id}
       value={value && value.type !== 'message' ? value.id : null}
       onChange={(id) => {
-        if (id !== messageId && type) onChange(id ? { id, type } : null)
+        if (id !== messageId) onChange(id ? { id, type: 'nodes' } : null)
       }}
       onFocus={(id) => {
-        if (id !== messageId && onFocus && type) onFocus(id ? { id, type } : null)
+        if (id !== messageId && onFocus) onFocus(id ? { id, type: 'nodes' } : null)
       }}
       label={'item'}
-      preload={false}
-      placeholder="Type search here..."
+      placeholder="Search nodes..."
     />
   )
 }
@@ -91,13 +126,8 @@ export const GraphSearchInput = ({
 /**
  * Component that display the search.
  */
-const GraphSearch: FC<GraphSearchInputProps & GraphSearchContextProviderProps> = ({
-  minisearchOptions,
-  ...props
-}) => (
-  <GraphSearchContextProvider minisearchOptions={minisearchOptions}>
-    <GraphSearchInput {...props} />
-  </GraphSearchContextProvider>
-)
+const GraphSearch: FC<GraphSearchInputProps & GraphSearchContextProviderProps> = ({ ...props }) => {
+  return <GraphSearchInput {...props} />
+}
 
 export default GraphSearch
diff --git a/lightrag/api/graph_viewer_webui/src/components/MessageAlert.tsx b/lightrag/api/graph_viewer_webui/src/components/MessageAlert.tsx
new file mode 100644
index 00000000..392459aa
--- /dev/null
+++ b/lightrag/api/graph_viewer_webui/src/components/MessageAlert.tsx
@@ -0,0 +1,38 @@
+import { Alert, AlertDescription, AlertTitle } from '@/components/ui/Alert'
+import Button from '@/components/ui/Button'
+import { useBackendState } from '@/stores/state'
+import { controlButtonVariant } from '@/lib/constants'
+
+import { AlertCircle } from 'lucide-react'
+
+const MessageAlert = () => {
+  const health = useBackendState.use.health()
+  const message = useBackendState.use.message()
+  const messageTitle = useBackendState.use.messageTitle()
+
+  return (
+    <Alert
+      variant={health ? 'default' : 'destructive'}
+      className="bg-background/90 absolute top-1/2 left-1/2 w-auto -translate-x-1/2 -translate-y-1/2 transform backdrop-blur-lg"
+    >
+      {!health && <AlertCircle className="h-4 w-4" />}
+      <AlertTitle>{messageTitle}</AlertTitle>
+
+      <AlertDescription>{message}</AlertDescription>
+      <div className="h-2" />
+      <div className="flex">
+        <div className="flex-auto" />
+        <Button
+          size="sm"
+          variant={controlButtonVariant}
+          className="text-primary max-h-8 border !p-2 text-xs"
+          onClick={() => useBackendState.getState().clear()}
+        >
+          Continue
+        </Button>
+      </div>
+    </Alert>
+  )
+}
+
+export default MessageAlert
diff --git a/lightrag/api/graph_viewer_webui/src/components/ui/AsyncSearch.tsx b/lightrag/api/graph_viewer_webui/src/components/ui/AsyncSearch.tsx
new file mode 100644
index 00000000..c417fd6b
--- /dev/null
+++ b/lightrag/api/graph_viewer_webui/src/components/ui/AsyncSearch.tsx
@@ -0,0 +1,243 @@
+import { useState, useEffect, useCallback } from 'react'
+import { Loader2 } from 'lucide-react'
+import { useDebounce } from '@/hooks/useDebounce'
+
+import { cn } from '@/lib/utils'
+import {
+  Command,
+  CommandEmpty,
+  CommandGroup,
+  CommandInput,
+  CommandItem,
+  CommandList
+} from '@/components/ui/Command'
+
+export interface Option {
+  value: string
+  label: string
+  disabled?: boolean
+  description?: string
+  icon?: React.ReactNode
+}
+
+export interface AsyncSearchProps<T> {
+  /** Async function to fetch options */
+  fetcher: (query?: string) => Promise<T[]>
+  /** Preload all data ahead of time */
+  preload?: boolean
+  /** Function to filter options */
+  filterFn?: (option: T, query: string) => boolean
+  /** Function to render each option */
+  renderOption: (option: T) => React.ReactNode
+  /** Function to get the value from an option */
+  getOptionValue: (option: T) => string
+  /** Custom not found message */
+  notFound?: React.ReactNode
+  /** Custom loading skeleton */
+  loadingSkeleton?: React.ReactNode
+  /** Currently selected value */
+  value: string | null
+  /** Callback when selection changes */
+  onChange: (value: string) => void
+  /** Callback when focus changes */
+  onFocus: (value: string) => void
+  /** Label for the select field */
+  label: string
+  /** Placeholder text when no selection */
+  placeholder?: string
+  /** Disable the entire select */
+  disabled?: boolean
+  /** Custom width for the popover */
+  width?: string | number
+  /** Custom class names */
+  className?: string
+  /** Custom trigger button class names */
+  triggerClassName?: string
+  /** Custom no results message */
+  noResultsMessage?: string
+  /** Allow clearing the selection */
+  clearable?: boolean
+}
+
+export function AsyncSearch<T>({
+  fetcher,
+  preload,
+  filterFn,
+  renderOption,
+  getOptionValue,
+  notFound,
+  loadingSkeleton,
+  label,
+  placeholder = 'Select...',
+  value,
+  onChange,
+  onFocus,
+  disabled = false,
+  className,
+  noResultsMessage
+}: AsyncSearchProps<T>) {
+  const [mounted, setMounted] = useState(false)
+  const [open, setOpen] = useState(false)
+  const [options, setOptions] = useState<T[]>([])
+  const [loading, setLoading] = useState(false)
+  const [error, setError] = useState<string | null>(null)
+  const [selectedValue, setSelectedValue] = useState(value)
+  const [focusedValue, setFocusedValue] = useState<string | null>(null)
+  const [searchTerm, setSearchTerm] = useState('')
+  const debouncedSearchTerm = useDebounce(searchTerm, preload ? 0 : 150)
+  const [originalOptions, setOriginalOptions] = useState<T[]>([])
+
+  useEffect(() => {
+    setMounted(true)
+    setSelectedValue(value)
+  }, [value])
+
+  // Effect for initial fetch
+  useEffect(() => {
+    const initializeOptions = async () => {
+      try {
+        setLoading(true)
+        setError(null)
+        // If we have a value, use it for the initial search
+        const data = value !== null ? await fetcher(value) : []
+        setOriginalOptions(data)
+        setOptions(data)
+      } catch (err) {
+        setError(err instanceof Error ? err.message : 'Failed to fetch options')
+      } finally {
+        setLoading(false)
+      }
+    }
+
+    if (!mounted) {
+      initializeOptions()
+    }
+  }, [mounted, fetcher, value])
+
+  useEffect(() => {
+    const fetchOptions = async () => {
+      try {
+        setLoading(true)
+        setError(null)
+        const data = await fetcher(debouncedSearchTerm)
+        setOriginalOptions(data)
+        setOptions(data)
+      } catch (err) {
+        setError(err instanceof Error ? err.message : 'Failed to fetch options')
+      } finally {
+        setLoading(false)
+      }
+    }
+
+    if (!mounted) {
+      fetchOptions()
+    } else if (!preload) {
+      fetchOptions()
+    } else if (preload) {
+      if (debouncedSearchTerm) {
+        setOptions(
+          originalOptions.filter((option) =>
+            filterFn ? filterFn(option, debouncedSearchTerm) : true
+          )
+        )
+      } else {
+        setOptions(originalOptions)
+      }
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [fetcher, debouncedSearchTerm, mounted, preload, filterFn])
+
+  const handleSelect = useCallback(
+    (currentValue: string) => {
+      if (currentValue !== selectedValue) {
+        setSelectedValue(currentValue)
+        onChange(currentValue)
+      }
+      setOpen(false)
+    },
+    [selectedValue, setSelectedValue, setOpen, onChange]
+  )
+
+  const handleFocus = useCallback(
+    (currentValue: string) => {
+      if (currentValue !== focusedValue) {
+        setFocusedValue(currentValue)
+        onFocus(currentValue)
+      }
+    },
+    [focusedValue, setFocusedValue, onFocus]
+  )
+
+  return (
+    <div
+      className={cn(disabled && 'cursor-not-allowed opacity-50', className)}
+      onFocus={() => {
+        setOpen(true)
+      }}
+      onBlur={() => setOpen(false)}
+    >
+      <Command shouldFilter={false} className="bg-transparent">
+        <div>
+          <CommandInput
+            placeholder={placeholder}
+            value={searchTerm}
+            className="max-h-8"
+            onValueChange={(value) => {
+              setSearchTerm(value)
+              if (value && !open) setOpen(true)
+            }}
+          />
+          {loading && options.length > 0 && (
+            <div className="absolute top-1/2 right-2 flex -translate-y-1/2 transform items-center">
+              <Loader2 className="h-4 w-4 animate-spin" />
+            </div>
+          )}
+        </div>
+        <CommandList className="max-h-auto" hidden={!open || debouncedSearchTerm.length === 0}>
+          {error && <div className="text-destructive p-4 text-center">{error}</div>}
+          {loading && options.length === 0 && (loadingSkeleton || <DefaultLoadingSkeleton />)}
+          {!loading &&
+            !error &&
+            options.length === 0 &&
+            (notFound || (
+              <CommandEmpty>{noResultsMessage ?? `No ${label.toLowerCase()} found.`}</CommandEmpty>
+            ))}
+          <CommandGroup>
+            {options.map((option, idx) => (
+              <>
+                <CommandItem
+                  key={getOptionValue(option) + `${idx}`}
+                  value={getOptionValue(option)}
+                  onSelect={handleSelect}
+                  onMouseEnter={() => handleFocus(getOptionValue(option))}
+                  className="truncate"
+                >
+                  {renderOption(option)}
+                </CommandItem>
+                {idx !== options.length - 1 && (
+                  <div key={idx} className="bg-foreground/10 h-[1px]" />
+                )}
+              </>
+            ))}
+          </CommandGroup>
+        </CommandList>
+      </Command>
+    </div>
+  )
+}
+
+function DefaultLoadingSkeleton() {
+  return (
+    <CommandGroup>
+      <CommandItem disabled>
+        <div className="flex w-full items-center gap-2">
+          <div className="bg-muted h-6 w-6 animate-pulse rounded-full" />
+          <div className="flex flex-1 flex-col gap-1">
+            <div className="bg-muted h-4 w-24 animate-pulse rounded" />
+            <div className="bg-muted h-3 w-16 animate-pulse rounded" />
+          </div>
+        </div>
+      </CommandItem>
+    </CommandGroup>
+  )
+}
diff --git a/lightrag/api/graph_viewer_webui/src/components/ui/AsyncSelect.tsx b/lightrag/api/graph_viewer_webui/src/components/ui/AsyncSelect.tsx
index 41747e26..37c50185 100644
--- a/lightrag/api/graph_viewer_webui/src/components/ui/AsyncSelect.tsx
+++ b/lightrag/api/graph_viewer_webui/src/components/ui/AsyncSelect.tsx
@@ -1,8 +1,9 @@
 import { useState, useEffect, useCallback } from 'react'
-import { Loader2 } from 'lucide-react'
+import { Check, ChevronsUpDown, Loader2 } from 'lucide-react'
 import { useDebounce } from '@/hooks/useDebounce'
 
 import { cn } from '@/lib/utils'
+import Button from '@/components/ui/Button'
 import {
   Command,
   CommandEmpty,
@@ -11,6 +12,7 @@ import {
   CommandItem,
   CommandList
 } from '@/components/ui/Command'
+import { Popover, PopoverContent, PopoverTrigger } from '@/components/ui/Popover'
 
 export interface Option {
   value: string
@@ -31,30 +33,34 @@ export interface AsyncSelectProps<T> {
   renderOption: (option: T) => React.ReactNode
   /** Function to get the value from an option */
   getOptionValue: (option: T) => string
+  /** Function to get the display value for the selected option */
+  getDisplayValue: (option: T) => React.ReactNode
   /** Custom not found message */
   notFound?: React.ReactNode
   /** Custom loading skeleton */
   loadingSkeleton?: React.ReactNode
   /** Currently selected value */
-  value: string | null
+  value: string
   /** Callback when selection changes */
   onChange: (value: string) => void
-  /** Callback when focus changes */
-  onFocus: (value: string) => void
   /** Label for the select field */
   label: string
   /** Placeholder text when no selection */
   placeholder?: string
   /** Disable the entire select */
   disabled?: boolean
-  /** Custom width for the popover */
+  /** Custom width for the popover *
   width?: string | number
   /** Custom class names */
   className?: string
   /** Custom trigger button class names */
   triggerClassName?: string
+  /** Custom search input class names */
+  searchInputClassName?: string
   /** Custom no results message */
   noResultsMessage?: string
+  /** Custom trigger tooltip */
+  triggerTooltip?: string
   /** Allow clearing the selection */
   clearable?: boolean
 }
@@ -65,16 +71,20 @@ export function AsyncSelect<T>({
   filterFn,
   renderOption,
   getOptionValue,
+  getDisplayValue,
   notFound,
   loadingSkeleton,
   label,
   placeholder = 'Select...',
   value,
   onChange,
-  onFocus,
   disabled = false,
   className,
-  noResultsMessage
+  triggerClassName,
+  searchInputClassName,
+  noResultsMessage,
+  triggerTooltip,
+  clearable = true
 }: AsyncSelectProps<T>) {
   const [mounted, setMounted] = useState(false)
   const [open, setOpen] = useState(false)
@@ -82,7 +92,7 @@ export function AsyncSelect<T>({
   const [loading, setLoading] = useState(false)
   const [error, setError] = useState<string | null>(null)
   const [selectedValue, setSelectedValue] = useState(value)
-  const [focusedValue, setFocusedValue] = useState<string | null>(null)
+  const [selectedOption, setSelectedOption] = useState<T | null>(null)
   const [searchTerm, setSearchTerm] = useState('')
   const debouncedSearchTerm = useDebounce(searchTerm, preload ? 0 : 150)
   const [originalOptions, setOriginalOptions] = useState<T[]>([])
@@ -92,6 +102,16 @@ export function AsyncSelect<T>({
     setSelectedValue(value)
   }, [value])
 
+  // Initialize selectedOption when options are loaded and value exists
+  useEffect(() => {
+    if (value && options.length > 0) {
+      const option = options.find((opt) => getOptionValue(opt) === value)
+      if (option) {
+        setSelectedOption(option)
+      }
+    }
+  }, [value, options, getOptionValue])
+
   // Effect for initial fetch
   useEffect(() => {
     const initializeOptions = async () => {
@@ -99,7 +119,7 @@ export function AsyncSelect<T>({
         setLoading(true)
         setError(null)
         // If we have a value, use it for the initial search
-        const data = value !== null ? await fetcher(value) : []
+        const data = await fetcher(value)
         setOriginalOptions(data)
         setOptions(data)
       } catch (err) {
@@ -149,75 +169,85 @@ export function AsyncSelect<T>({
 
   const handleSelect = useCallback(
     (currentValue: string) => {
-      console.log('handleSelect')
-      if (currentValue !== selectedValue) {
-        setSelectedValue(currentValue)
-        onChange(currentValue)
-      }
+      const newValue = clearable && currentValue === selectedValue ? '' : currentValue
+      setSelectedValue(newValue)
+      setSelectedOption(options.find((option) => getOptionValue(option) === newValue) || null)
+      onChange(newValue)
       setOpen(false)
     },
-    [selectedValue, onChange]
-  )
-
-  const handleFocus = useCallback(
-    (currentValue: string) => {
-      if (currentValue !== focusedValue) {
-        setFocusedValue(currentValue)
-        onFocus(currentValue)
-      }
-    },
-    [focusedValue, onFocus]
+    [selectedValue, onChange, clearable, options, getOptionValue]
   )
 
   return (
-    <div
-      className={cn(disabled && 'cursor-not-allowed opacity-50', className)}
-      onFocus={() => setOpen(true)}
-      onBlur={() => setOpen(false)}
-    >
-      <Command shouldFilter={false} className="bg-transparent">
-        <div className="relative w-full">
-          <CommandInput
-            placeholder={placeholder}
-            value={searchTerm}
-            onValueChange={(value) => {
-              setSearchTerm(value)
-              if (value && !open) setOpen(true)
-            }}
-          />
-          {loading && options.length > 0 && (
-            <div className="absolute top-1/2 right-2 flex -translate-y-1/2 transform items-center">
-              <Loader2 className="h-4 w-4 animate-spin" />
-            </div>
+    <Popover open={open} onOpenChange={setOpen}>
+      <PopoverTrigger asChild>
+        <Button
+          variant="outline"
+          role="combobox"
+          aria-expanded={open}
+          className={cn(
+            'justify-between',
+            disabled && 'cursor-not-allowed opacity-50',
+            triggerClassName
           )}
-        </div>
-        <CommandList className="max-h-auto" hidden={!open || debouncedSearchTerm.length === 0}>
-          {error && <div className="text-destructive p-4 text-center">{error}</div>}
-          {loading && options.length === 0 && (loadingSkeleton || <DefaultLoadingSkeleton />)}
-          {!loading &&
-            !error &&
-            options.length === 0 &&
-            (notFound || (
-              <CommandEmpty>{noResultsMessage ?? `No ${label.toLowerCase()} found.`}</CommandEmpty>
-            ))}
-          <CommandGroup>
-            {options.map((option, idx) => (
-              <>
+          disabled={disabled}
+          tooltip={triggerTooltip}
+          side="bottom"
+        >
+          {selectedOption ? getDisplayValue(selectedOption) : placeholder}
+          <ChevronsUpDown className="opacity-50" size={10} />
+        </Button>
+      </PopoverTrigger>
+      <PopoverContent className={cn('p-0', className)} onCloseAutoFocus={(e) => e.preventDefault()}>
+        <Command shouldFilter={false}>
+          <div className="relative w-full border-b">
+            <CommandInput
+              placeholder={`Search ${label.toLowerCase()}...`}
+              value={searchTerm}
+              onValueChange={(value) => {
+                setSearchTerm(value)
+              }}
+              className={searchInputClassName}
+            />
+            {loading && options.length > 0 && (
+              <div className="absolute top-1/2 right-2 flex -translate-y-1/2 transform items-center">
+                <Loader2 className="h-4 w-4 animate-spin" />
+              </div>
+            )}
+          </div>
+          <CommandList>
+            {error && <div className="text-destructive p-4 text-center">{error}</div>}
+            {loading && options.length === 0 && (loadingSkeleton || <DefaultLoadingSkeleton />)}
+            {!loading &&
+              !error &&
+              options.length === 0 &&
+              (notFound || (
+                <CommandEmpty>
+                  {noResultsMessage ?? `No ${label.toLowerCase()} found.`}
+                </CommandEmpty>
+              ))}
+            <CommandGroup>
+              {options.map((option) => (
                 <CommandItem
                   key={getOptionValue(option)}
                   value={getOptionValue(option)}
                   onSelect={handleSelect}
-                  onMouseEnter={() => handleFocus(getOptionValue(option))}
+                  className="truncate"
                 >
                   {renderOption(option)}
+                  <Check
+                    className={cn(
+                      'ml-auto h-3 w-3',
+                      selectedValue === getOptionValue(option) ? 'opacity-100' : 'opacity-0'
+                    )}
+                  />
                 </CommandItem>
-                {idx !== options.length - 1 && <div className="bg-foreground/10 h-[1px]" />}
-              </>
-            ))}
-          </CommandGroup>
-        </CommandList>
-      </Command>
-    </div>
+              ))}
+            </CommandGroup>
+          </CommandList>
+        </Command>
+      </PopoverContent>
+    </Popover>
   )
 }
 
diff --git a/lightrag/api/graph_viewer_webui/src/hooks/useLightragGraph.tsx b/lightrag/api/graph_viewer_webui/src/hooks/useLightragGraph.tsx
index 3925de0d..1f82760f 100644
--- a/lightrag/api/graph_viewer_webui/src/hooks/useLightragGraph.tsx
+++ b/lightrag/api/graph_viewer_webui/src/hooks/useLightragGraph.tsx
@@ -1,10 +1,13 @@
 import Graph, { DirectedGraph } from 'graphology'
-import { useCallback, useEffect, useState } from 'react'
-import { randomColor } from '@/lib/utils'
+import { useCallback, useEffect } from 'react'
+import { randomColor, errorMessage } from '@/lib/utils'
 import * as Constants from '@/lib/constants'
 import { useGraphStore, RawGraph } from '@/stores/graph'
 import { queryGraphs } from '@/api/lightrag'
 import { useBackendState } from '@/stores/state'
+import { useSettingsStore } from '@/stores/settings'
+
+import seedrandom from 'seedrandom'
 
 const validateGraph = (graph: RawGraph) => {
   if (!graph) {
@@ -53,9 +56,7 @@ const fetchGraph = async (label: string) => {
   try {
     rawData = await queryGraphs(label)
   } catch (e) {
-    useBackendState
-      .getState()
-      .setErrorMessage(e instanceof Error ? e.message : `${e}`, 'Query Graphs Error!')
+    useBackendState.getState().setErrorMessage(errorMessage(e), 'Query Graphs Error!')
     return null
   }
 
@@ -69,9 +70,11 @@ const fetchGraph = async (label: string) => {
       const node = rawData.nodes[i]
       nodeIdMap[node.id] = i
 
+      // const seed = node.labels.length > 0 ? node.labels[0] : node.id
+      seedrandom(node.id, { global: true })
+      node.color = randomColor()
       node.x = Math.random()
       node.y = Math.random()
-      node.color = randomColor()
       node.degree = 0
       node.size = 10
     }
@@ -150,8 +153,10 @@ const createSigmaGraph = (rawGraph: RawGraph | null) => {
   return graph
 }
 
+const lastQueryLabel = { label: '' }
+
 const useLightrangeGraph = () => {
-  const [fetchLabel, setFetchLabel] = useState<string>('*')
+  const queryLabel = useSettingsStore.use.queryLabel()
   const rawGraph = useGraphStore.use.rawGraph()
   const sigmaGraph = useGraphStore.use.sigmaGraph()
 
@@ -170,12 +175,13 @@ const useLightrangeGraph = () => {
   )
 
   useEffect(() => {
-    if (fetchLabel) {
-      const state = useGraphStore.getState()
-      if (state.queryLabel !== fetchLabel) {
+    if (queryLabel) {
+      if (lastQueryLabel.label !== queryLabel) {
+        lastQueryLabel.label = queryLabel
+        const state = useGraphStore.getState()
         state.reset()
-        fetchGraph(fetchLabel).then((data) => {
-          state.setQueryLabel(fetchLabel)
+        fetchGraph(queryLabel).then((data) => {
+          // console.debug('Query label: ' + queryLabel)
           state.setSigmaGraph(createSigmaGraph(data))
           data?.buildDynamicMap()
           state.setRawGraph(data)
@@ -186,7 +192,7 @@ const useLightrangeGraph = () => {
       state.reset()
       state.setSigmaGraph(new DirectedGraph())
     }
-  }, [fetchLabel])
+  }, [queryLabel])
 
   const lightrageGraph = useCallback(() => {
     if (sigmaGraph) {
@@ -197,7 +203,7 @@ const useLightrangeGraph = () => {
     return graph as Graph<NodeType, EdgeType>
   }, [sigmaGraph])
 
-  return { lightrageGraph, fetchLabel, setFetchLabel, getNode, getEdge }
+  return { lightrageGraph, getNode, getEdge }
 }
 
 export default useLightrangeGraph
diff --git a/lightrag/api/graph_viewer_webui/src/lib/constants.ts b/lightrag/api/graph_viewer_webui/src/lib/constants.ts
index 0d2fb6c5..c7b59cbd 100644
--- a/lightrag/api/graph_viewer_webui/src/lib/constants.ts
+++ b/lightrag/api/graph_viewer_webui/src/lib/constants.ts
@@ -19,3 +19,7 @@ export const searchResultLimit = 20
 
 export const minNodeSize = 4
 export const maxNodeSize = 20
+
+export const healthCheckInterval = 15 // seconds
+
+export const defaultQueryLabel = '*'
diff --git a/lightrag/api/graph_viewer_webui/src/lib/utils.ts b/lightrag/api/graph_viewer_webui/src/lib/utils.ts
index 4e883619..f5cc52d2 100644
--- a/lightrag/api/graph_viewer_webui/src/lib/utils.ts
+++ b/lightrag/api/graph_viewer_webui/src/lib/utils.ts
@@ -15,6 +15,10 @@ export function randomColor() {
   return code
 }
 
+export function errorMessage(error: any) {
+  return error instanceof Error ? error.message : `${error}`
+}
+
 type WithSelectors<S> = S extends { getState: () => infer T }
   ? S & { use: { [K in keyof T]: () => T[K] } }
   : never
diff --git a/lightrag/api/graph_viewer_webui/src/stores/graph.ts b/lightrag/api/graph_viewer_webui/src/stores/graph.ts
index 66502af6..b78e9bf8 100644
--- a/lightrag/api/graph_viewer_webui/src/stores/graph.ts
+++ b/lightrag/api/graph_viewer_webui/src/stores/graph.ts
@@ -63,7 +63,6 @@ interface GraphState {
   selectedEdge: string | null
   focusedEdge: string | null
 
-  queryLabel: string | null
   rawGraph: RawGraph | null
   sigmaGraph: DirectedGraph | null
 
@@ -78,7 +77,6 @@ interface GraphState {
 
   setMoveToSelectedNode: (moveToSelectedNode: boolean) => void
 
-  setQueryLabel: (queryLabel: string | null) => void
   setRawGraph: (rawGraph: RawGraph | null) => void
   setSigmaGraph: (sigmaGraph: DirectedGraph | null) => void
 }
@@ -91,7 +89,6 @@ const useGraphStoreBase = create<GraphState>()((set) => ({
 
   moveToSelectedNode: false,
 
-  queryLabel: null,
   rawGraph: null,
   sigmaGraph: null,
 
@@ -113,17 +110,11 @@ const useGraphStoreBase = create<GraphState>()((set) => ({
       focusedNode: null,
       selectedEdge: null,
       focusedEdge: null,
-      queryLabel: null,
       rawGraph: null,
       sigmaGraph: null,
       moveToSelectedNode: false
     }),
 
-  setQueryLabel: (queryLabel: string | null) =>
-    set({
-      queryLabel
-    }),
-
   setRawGraph: (rawGraph: RawGraph | null) =>
     set({
       rawGraph
diff --git a/lightrag/api/graph_viewer_webui/src/stores/settings.ts b/lightrag/api/graph_viewer_webui/src/stores/settings.ts
index e1d6537f..01b9357a 100644
--- a/lightrag/api/graph_viewer_webui/src/stores/settings.ts
+++ b/lightrag/api/graph_viewer_webui/src/stores/settings.ts
@@ -1,6 +1,7 @@
 import { create } from 'zustand'
 import { persist, createJSONStorage } from 'zustand/middleware'
 import { createSelectors } from '@/lib/utils'
+import { defaultQueryLabel } from '@/lib/constants'
 
 type Theme = 'dark' | 'light' | 'system'
 
@@ -10,7 +11,11 @@ interface SettingsState {
   enableEdgeEvents: boolean
   enableHideUnselectedEdges: boolean
   showEdgeLabel: boolean
+
   setTheme: (theme: Theme) => void
+
+  queryLabel: string
+  setQueryLabel: (queryLabel: string) => void
 }
 
 const useSettingsStoreBase = create<SettingsState>()(
@@ -22,16 +27,26 @@ const useSettingsStoreBase = create<SettingsState>()(
       enableHideUnselectedEdges: true,
       showEdgeLabel: false,
 
-      setTheme: (theme: Theme) => set({ theme })
+      queryLabel: defaultQueryLabel,
+
+      setTheme: (theme: Theme) => set({ theme }),
+
+      setQueryLabel: (queryLabel: string) =>
+        set({
+          queryLabel
+        })
     }),
     {
       name: 'settings-storage',
       storage: createJSONStorage(() => localStorage),
-      version: 2,
+      version: 3,
       migrate: (state: any, version: number) => {
         if (version < 2) {
           state.showEdgeLabel = false
         }
+        if (version < 3) {
+          state.queryLabel = defaultQueryLabel
+        }
       }
     }
   )