From 40a1a94a311c1c20fcd49d34992419dd91a8224e Mon Sep 17 00:00:00 2001 From: ArnoChen Date: Tue, 25 Feb 2025 18:28:31 +0800 Subject: [PATCH] add graph depth and layout iteration settings --- lightrag/api/routers/graph_routes.py | 4 +- lightrag_webui/src/api/lightrag.ts | 4 +- .../src/components/graph/GraphControl.tsx | 4 +- .../src/components/graph/GraphLabels.tsx | 108 +++++++++-------- .../src/components/graph/GraphSearch.tsx | 10 +- .../src/components/graph/LayoutsControl.tsx | 7 +- .../src/components/graph/Settings.tsx | 109 ++++++++++++++++-- .../src/components/retrieval/ChatMessage.tsx | 6 +- lightrag_webui/src/hooks/useLightragGraph.tsx | 15 ++- lightrag_webui/src/lib/constants.ts | 1 + lightrag_webui/src/stores/settings.ts | 51 +++++--- 11 files changed, 233 insertions(+), 86 deletions(-) diff --git a/lightrag/api/routers/graph_routes.py b/lightrag/api/routers/graph_routes.py index 28a5561a..95a72758 100644 --- a/lightrag/api/routers/graph_routes.py +++ b/lightrag/api/routers/graph_routes.py @@ -20,8 +20,8 @@ def create_graph_routes(rag, api_key: Optional[str] = None): return await rag.get_graph_labels() @router.get("/graphs", dependencies=[Depends(optional_api_key)]) - async def get_knowledge_graph(label: str): + async def get_knowledge_graph(label: str, max_depth: int = 3): """Get knowledge graph for a specific label""" - return await rag.get_knowledge_graph(node_label=label, max_depth=3) + return await rag.get_knowledge_graph(node_label=label, max_depth=max_depth) return router diff --git a/lightrag_webui/src/api/lightrag.ts b/lightrag_webui/src/api/lightrag.ts index 1de79898..c9838335 100644 --- a/lightrag_webui/src/api/lightrag.ts +++ b/lightrag_webui/src/api/lightrag.ts @@ -161,8 +161,8 @@ axiosInstance.interceptors.response.use( ) // API methods -export const queryGraphs = async (label: string): Promise => { - const response = await axiosInstance.get(`/graphs?label=${label}`) +export const queryGraphs = async (label: string, maxDepth: number): Promise => { + const response = await axiosInstance.get(`/graphs?label=${label}&max_depth=${maxDepth}`) return response.data } diff --git a/lightrag_webui/src/components/graph/GraphControl.tsx b/lightrag_webui/src/components/graph/GraphControl.tsx index ecdf121a..d7946240 100644 --- a/lightrag_webui/src/components/graph/GraphControl.tsx +++ b/lightrag_webui/src/components/graph/GraphControl.tsx @@ -26,8 +26,10 @@ const GraphControl = ({ disableHoverEffect }: { disableHoverEffect?: boolean }) const registerEvents = useRegisterEvents() const setSettings = useSetSettings() const loadGraph = useLoadGraph() + + const maxIterations = useSettingsStore.use.graphLayoutMaxIterations() const { assign: assignLayout } = useLayoutForceAtlas2({ - iterations: 20 + iterations: maxIterations }) const { theme } = useTheme() diff --git a/lightrag_webui/src/components/graph/GraphLabels.tsx b/lightrag_webui/src/components/graph/GraphLabels.tsx index b3c325ad..a3849e1f 100644 --- a/lightrag_webui/src/components/graph/GraphLabels.tsx +++ b/lightrag_webui/src/components/graph/GraphLabels.tsx @@ -1,67 +1,81 @@ -import { useCallback, useState } from 'react' +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 [labels, setLabels] = useState<{ - labels: string[] - searchEngine: MiniSearch | null - }>({ - labels: [], - searchEngine: null - }) - const [fetched, setFetched] = useState(false) + 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]) + } + + // Create search engine + const 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) + + lastGraph.graph = graph + lastGraph.searchEngine = searchEngine + lastGraph.labels = labels + + return { + labels, + searchEngine + } + }, [graph]) const fetchData = useCallback( async (query?: string): Promise => { - let _labels = labels.labels - let _searchEngine = labels.searchEngine + const { labels, searchEngine } = await getSearchEngine() - 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 + let result: string[] = labels + if (query) { + // Search labels + result = searchEngine.search(query).map((r) => labels[r.id]) } - // Search labels - return _searchEngine.search(query).map((result) => _labels[result.id]) + return result.length <= labelListLimit + ? result + : [...result.slice(0, labelListLimit), `And ${result.length - labelListLimit} others`] }, - [labels, fetched, setLabels, setFetched] + [getSearchEngine] ) const setQueryLabel = useCallback((label: string) => { + if (label.startsWith('And ') && label.endsWith(' others')) return useSettingsStore.getState().setQueryLabel(label) }, []) diff --git a/lightrag_webui/src/components/graph/GraphSearch.tsx b/lightrag_webui/src/components/graph/GraphSearch.tsx index 16d4c62d..3edc3ede 100644 --- a/lightrag_webui/src/components/graph/GraphSearch.tsx +++ b/lightrag_webui/src/components/graph/GraphSearch.tsx @@ -46,7 +46,7 @@ export const GraphSearchInput = ({ }) => { const graph = useGraphStore.use.sigmaGraph() - const search = useMemo(() => { + const searchEngine = useMemo(() => { if (lastGraph.graph == graph) { return lastGraph.searchEngine } @@ -83,9 +83,9 @@ export const GraphSearchInput = ({ const loadOptions = useCallback( async (query?: string): Promise => { if (onFocus) onFocus(null) - if (!query || !search) return [] - const result: OptionItem[] = search.search(query).map((result) => ({ - id: result.id, + if (!query || !searchEngine) return [] + const result: OptionItem[] = searchEngine.search(query).map((r) => ({ + id: r.id, type: 'nodes' })) @@ -101,7 +101,7 @@ export const GraphSearchInput = ({ } ] }, - [search, onFocus] + [searchEngine, onFocus] ) return ( diff --git a/lightrag_webui/src/components/graph/LayoutsControl.tsx b/lightrag_webui/src/components/graph/LayoutsControl.tsx index d7fcc2b8..1bb29c90 100644 --- a/lightrag_webui/src/components/graph/LayoutsControl.tsx +++ b/lightrag_webui/src/components/graph/LayoutsControl.tsx @@ -13,6 +13,7 @@ import Button from '@/components/ui/Button' import { Popover, PopoverContent, PopoverTrigger } from '@/components/ui/Popover' import { Command, CommandGroup, CommandItem, CommandList } from '@/components/ui/Command' import { controlButtonVariant } from '@/lib/constants' +import { useSettingsStore } from '@/stores/settings' import { GripIcon, PlayIcon, PauseIcon } from 'lucide-react' @@ -75,13 +76,15 @@ const LayoutsControl = () => { const sigma = useSigma() const [layout, setLayout] = useState('Circular') const [opened, setOpened] = useState(false) + + const maxIterations = useSettingsStore.use.graphLayoutMaxIterations() const layoutCircular = useLayoutCircular() const layoutCirclepack = useLayoutCirclepack() const layoutRandom = useLayoutRandom() const layoutNoverlap = useLayoutNoverlap({ settings: { margin: 1 } }) - const layoutForce = useLayoutForce({ maxIterations: 20 }) - const layoutForceAtlas2 = useLayoutForceAtlas2({ iterations: 20 }) + const layoutForce = useLayoutForce({ maxIterations: maxIterations }) + const layoutForceAtlas2 = useLayoutForceAtlas2({ iterations: maxIterations }) const workerNoverlap = useWorkerLayoutNoverlap() const workerForce = useWorkerLayoutForce() const workerForceAtlas2 = useWorkerLayoutForceAtlas2() diff --git a/lightrag_webui/src/components/graph/Settings.tsx b/lightrag_webui/src/components/graph/Settings.tsx index 3a6cc51c..4d2b998d 100644 --- a/lightrag_webui/src/components/graph/Settings.tsx +++ b/lightrag_webui/src/components/graph/Settings.tsx @@ -1,9 +1,10 @@ +import { useState, useCallback, useEffect } from 'react' import { Popover, PopoverContent, PopoverTrigger } from '@/components/ui/Popover' import Checkbox from '@/components/ui/Checkbox' import Button from '@/components/ui/Button' import Separator from '@/components/ui/Separator' import Input from '@/components/ui/Input' -import { useState, useCallback, useEffect } from 'react' + import { controlButtonVariant } from '@/lib/constants' import { useSettingsStore } from '@/stores/settings' import { useBackendState } from '@/stores/state' @@ -35,6 +36,74 @@ const LabeledCheckBox = ({ ) } +/** + * Component that displays a number input with a label. + */ +const LabeledNumberInput = ({ + value, + onEditFinished, + label, + min, + max +}: { + value: number + onEditFinished: (value: number) => void + label: string + min: number + max?: number +}) => { + const [currentValue, setCurrentValue] = useState(value) + + const onValueChange = useCallback( + (e: React.ChangeEvent) => { + const text = e.target.value.trim() + if (text.length === 0) { + setCurrentValue(null) + return + } + const newValue = Number.parseInt(text) + if (!isNaN(newValue) && newValue !== currentValue) { + if (min !== undefined && newValue < min) { + return + } + if (max !== undefined && newValue > max) { + return + } + setCurrentValue(newValue) + } + }, + [currentValue, min, max] + ) + + const onBlur = useCallback(() => { + if (currentValue !== null && value !== currentValue) { + onEditFinished(currentValue) + } + }, [value, currentValue, onEditFinished]) + + return ( +
+ + { + if (e.key === 'Enter') { + onBlur() + } + }} + /> +
+ ) +} + /** * Component that displays a popover with settings options. */ @@ -45,11 +114,12 @@ export default function Settings() { const showPropertyPanel = useSettingsStore.use.showPropertyPanel() const showNodeSearchBar = useSettingsStore.use.showNodeSearchBar() const showNodeLabel = useSettingsStore.use.showNodeLabel() - const enableEdgeEvents = useSettingsStore.use.enableEdgeEvents() const enableNodeDrag = useSettingsStore.use.enableNodeDrag() const enableHideUnselectedEdges = useSettingsStore.use.enableHideUnselectedEdges() const showEdgeLabel = useSettingsStore.use.showEdgeLabel() + const graphQueryMaxDepth = useSettingsStore.use.graphQueryMaxDepth() + const graphLayoutMaxIterations = useSettingsStore.use.graphLayoutMaxIterations() const enableHealthCheck = useSettingsStore.use.enableHealthCheck() const apiKey = useSettingsStore.use.apiKey() @@ -102,6 +172,16 @@ export default function Settings() { [] ) + const setGraphQueryMaxDepth = useCallback((depth: number) => { + if (depth < 1) return + useSettingsStore.setState({ graphQueryMaxDepth: depth }) + }, []) + + const setGraphLayoutMaxIterations = useCallback((iterations: number) => { + if (iterations < 1) return + useSettingsStore.setState({ graphLayoutMaxIterations: iterations }) + }, []) + const setApiKey = useCallback(async () => { useSettingsStore.setState({ apiKey: tempApiKey || null }) await useBackendState.getState().check() @@ -129,6 +209,14 @@ export default function Settings() { onCloseAutoFocus={(e) => e.preventDefault()} >
+ + + + - - + diff --git a/lightrag_webui/src/components/retrieval/ChatMessage.tsx b/lightrag_webui/src/components/retrieval/ChatMessage.tsx index 2da01c6e..ea3dba12 100644 --- a/lightrag_webui/src/components/retrieval/ChatMessage.tsx +++ b/lightrag_webui/src/components/retrieval/ChatMessage.tsx @@ -2,6 +2,7 @@ import { ReactNode, useCallback } from 'react' import { Message } from '@/api/lightrag' import useTheme from '@/hooks/useTheme' import Button from '@/components/ui/Button' +import { cn } from '@/lib/utils' import ReactMarkdown from 'react-markdown' import remarkGfm from 'remark-gfm' @@ -101,7 +102,10 @@ const CodeHighlight = ({ className, children, node, ...props }: CodeHighlightPro {String(children).replace(/\n$/, '')} ) : ( - + {children} ) diff --git a/lightrag_webui/src/hooks/useLightragGraph.tsx b/lightrag_webui/src/hooks/useLightragGraph.tsx index b9c291bc..9393dfaf 100644 --- a/lightrag_webui/src/hooks/useLightragGraph.tsx +++ b/lightrag_webui/src/hooks/useLightragGraph.tsx @@ -50,11 +50,11 @@ export type NodeType = { } export type EdgeType = { label: string } -const fetchGraph = async (label: string) => { +const fetchGraph = async (label: string, maxDepth: number) => { let rawData: any = null try { - rawData = await queryGraphs(label) + rawData = await queryGraphs(label, maxDepth) } catch (e) { useBackendState.getState().setErrorMessage(errorMessage(e), 'Query Graphs Error!') return null @@ -161,12 +161,13 @@ const createSigmaGraph = (rawGraph: RawGraph | null) => { return graph } -const lastQueryLabel = { label: '' } +const lastQueryLabel = { label: '', maxQueryDepth: 0 } const useLightrangeGraph = () => { const queryLabel = useSettingsStore.use.queryLabel() const rawGraph = useGraphStore.use.rawGraph() const sigmaGraph = useGraphStore.use.sigmaGraph() + const maxQueryDepth = useSettingsStore.use.graphQueryMaxDepth() const getNode = useCallback( (nodeId: string) => { @@ -184,11 +185,13 @@ const useLightrangeGraph = () => { useEffect(() => { if (queryLabel) { - if (lastQueryLabel.label !== queryLabel) { + if (lastQueryLabel.label !== queryLabel || lastQueryLabel.maxQueryDepth !== maxQueryDepth) { lastQueryLabel.label = queryLabel + lastQueryLabel.maxQueryDepth = maxQueryDepth + const state = useGraphStore.getState() state.reset() - fetchGraph(queryLabel).then((data) => { + fetchGraph(queryLabel, maxQueryDepth).then((data) => { // console.debug('Query label: ' + queryLabel) state.setSigmaGraph(createSigmaGraph(data)) data?.buildDynamicMap() @@ -200,7 +203,7 @@ const useLightrangeGraph = () => { state.reset() state.setSigmaGraph(new DirectedGraph()) } - }, [queryLabel]) + }, [queryLabel, maxQueryDepth]) const lightrageGraph = useCallback(() => { if (sigmaGraph) { diff --git a/lightrag_webui/src/lib/constants.ts b/lightrag_webui/src/lib/constants.ts index 0a07e1c4..aca6bef6 100644 --- a/lightrag_webui/src/lib/constants.ts +++ b/lightrag_webui/src/lib/constants.ts @@ -16,6 +16,7 @@ export const edgeColorSelected = '#F57F17' export const edgeColorHighlighted = '#B2EBF2' export const searchResultLimit = 20 +export const labelListLimit = 40 export const minNodeSize = 4 export const maxNodeSize = 20 diff --git a/lightrag_webui/src/stores/settings.ts b/lightrag_webui/src/stores/settings.ts index 60ae8f90..d33969cb 100644 --- a/lightrag_webui/src/stores/settings.ts +++ b/lightrag_webui/src/stores/settings.ts @@ -8,9 +8,7 @@ type Theme = 'dark' | 'light' | 'system' type Tab = 'documents' | 'knowledge-graph' | 'retrieval' | 'api' interface SettingsState { - theme: Theme - setTheme: (theme: Theme) => void - + // Graph viewer settings showPropertyPanel: boolean showNodeSearchBar: boolean @@ -21,23 +19,35 @@ interface SettingsState { enableHideUnselectedEdges: boolean enableEdgeEvents: boolean + graphQueryMaxDepth: number + setGraphQueryMaxDepth: (depth: number) => void + + graphLayoutMaxIterations: number + setGraphLayoutMaxIterations: (iterations: number) => void + + // Retrieval settings queryLabel: string setQueryLabel: (queryLabel: string) => void - enableHealthCheck: boolean - setEnableHealthCheck: (enable: boolean) => void - - apiKey: string | null - setApiKey: (key: string | null) => void - - currentTab: Tab - setCurrentTab: (tab: Tab) => void - retrievalHistory: Message[] setRetrievalHistory: (history: Message[]) => void querySettings: Omit updateQuerySettings: (settings: Partial) => void + + // Auth settings + apiKey: string | null + setApiKey: (key: string | null) => void + + // App settings + theme: Theme + setTheme: (theme: Theme) => void + + enableHealthCheck: boolean + setEnableHealthCheck: (enable: boolean) => void + + currentTab: Tab + setCurrentTab: (tab: Tab) => void } const useSettingsStoreBase = create()( @@ -55,7 +65,11 @@ const useSettingsStoreBase = create()( enableHideUnselectedEdges: true, enableEdgeEvents: false, + graphQueryMaxDepth: 3, + graphLayoutMaxIterations: 10, + queryLabel: defaultQueryLabel, + enableHealthCheck: true, apiKey: null, @@ -81,11 +95,18 @@ const useSettingsStoreBase = create()( setTheme: (theme: Theme) => set({ theme }), + setGraphLayoutMaxIterations: (iterations: number) => + set({ + graphLayoutMaxIterations: iterations + }), + setQueryLabel: (queryLabel: string) => set({ queryLabel }), + setGraphQueryMaxDepth: (depth: number) => set({ graphQueryMaxDepth: depth }), + setEnableHealthCheck: (enable: boolean) => set({ enableHealthCheck: enable }), setApiKey: (apiKey: string | null) => set({ apiKey }), @@ -102,7 +123,7 @@ const useSettingsStoreBase = create()( { name: 'settings-storage', storage: createJSONStorage(() => localStorage), - version: 6, + version: 7, migrate: (state: any, version: number) => { if (version < 2) { state.showEdgeLabel = false @@ -137,6 +158,10 @@ const useSettingsStoreBase = create()( } state.retrievalHistory = [] } + if (version < 7) { + state.graphQueryMaxDepth = 3 + state.graphLayoutMaxIterations = 10 + } return state } }