diff --git a/lightrag_webui/src/components/graph/GraphLabels.tsx b/lightrag_webui/src/components/graph/GraphLabels.tsx index 1e73eb38..474eb410 100644 --- a/lightrag_webui/src/components/graph/GraphLabels.tsx +++ b/lightrag_webui/src/components/graph/GraphLabels.tsx @@ -1,35 +1,15 @@ import { useCallback } from 'react' import { AsyncSelect } from '@/components/ui/AsyncSelect' -import { getGraphLabels } from '@/api/lightrag' import { useSettingsStore } from '@/stores/settings' import { useGraphStore } from '@/stores/graph' import { labelListLimit } from '@/lib/constants' import MiniSearch from 'minisearch' -const lastGraph: any = { - graph: null, - searchEngine: null, - labels: [] -} - const GraphLabels = () => { const label = useSettingsStore.use.queryLabel() - const graph = useGraphStore.use.sigmaGraph() - - const getSearchEngine = useCallback(async () => { - if (lastGraph.graph == graph) { - return { - labels: lastGraph.labels, - searchEngine: lastGraph.searchEngine - } - } - const labels = ['*'].concat(await getGraphLabels()) - - // Ensure query label exists - if (!labels.includes(useSettingsStore.getState().queryLabel)) { - useSettingsStore.getState().setQueryLabel(labels[0]) - } + const graphLabels = useGraphStore.use.graphLabels() + const getSearchEngine = useCallback(() => { // Create search engine const searchEngine = new MiniSearch({ idField: 'id', @@ -44,22 +24,18 @@ const GraphLabels = () => { }) // Add documents - const documents = labels.map((str, index) => ({ id: index, value: str })) + const documents = graphLabels.map((str, index) => ({ id: index, value: str })) searchEngine.addAll(documents) - lastGraph.graph = graph - lastGraph.searchEngine = searchEngine - lastGraph.labels = labels - return { - labels, + labels: graphLabels, searchEngine } - }, [graph]) + }, [graphLabels]) const fetchData = useCallback( async (query?: string): Promise => { - const { labels, searchEngine } = await getSearchEngine() + const { labels, searchEngine } = getSearchEngine() let result: string[] = labels if (query) { diff --git a/lightrag_webui/src/hooks/useLightragGraph.tsx b/lightrag_webui/src/hooks/useLightragGraph.tsx index 1349e857..3951b783 100644 --- a/lightrag_webui/src/hooks/useLightragGraph.tsx +++ b/lightrag_webui/src/hooks/useLightragGraph.tsx @@ -205,6 +205,21 @@ const useLightrangeGraph = () => { state.setSigmaGraph(createSigmaGraph(data)) data?.buildDynamicMap() state.setRawGraph(data) + + // Extract labels from graph data + if (data) { + const labelSet = new Set(['*']) + for (const node of data.nodes) { + if (node.labels && Array.isArray(node.labels)) { + for (const label of node.labels) { + labelSet.add(label) + } + } + } + state.setGraphLabels(Array.from(labelSet).sort()) + } else { + state.setGraphLabels(['*']) + } }) } } else { diff --git a/lightrag_webui/src/stores/graph.ts b/lightrag_webui/src/stores/graph.ts index b7c2120c..082c3161 100644 --- a/lightrag_webui/src/stores/graph.ts +++ b/lightrag_webui/src/stores/graph.ts @@ -65,6 +65,7 @@ interface GraphState { rawGraph: RawGraph | null sigmaGraph: DirectedGraph | null + graphLabels: string[] moveToSelectedNode: boolean @@ -79,6 +80,7 @@ interface GraphState { setRawGraph: (rawGraph: RawGraph | null) => void setSigmaGraph: (sigmaGraph: DirectedGraph | null) => void + setGraphLabels: (labels: string[]) => void } const useGraphStoreBase = create()((set) => ({ @@ -91,6 +93,7 @@ const useGraphStoreBase = create()((set) => ({ rawGraph: null, sigmaGraph: null, + graphLabels: [], setSelectedNode: (nodeId: string | null, moveToSelectedNode?: boolean) => set({ selectedNode: nodeId, moveToSelectedNode }), @@ -112,6 +115,7 @@ const useGraphStoreBase = create()((set) => ({ focusedEdge: null, rawGraph: null, sigmaGraph: null, + graphLabels: [], moveToSelectedNode: false }), @@ -121,6 +125,8 @@ const useGraphStoreBase = create()((set) => ({ }), setSigmaGraph: (sigmaGraph: DirectedGraph | null) => set({ sigmaGraph }), + + setGraphLabels: (labels: string[]) => set({ graphLabels: labels }), setMoveToSelectedNode: (moveToSelectedNode?: boolean) => set({ moveToSelectedNode }) }))