Files
lightrag/lightrag_webui/src/components/graph/GraphLabels.tsx
2025-02-25 18:32:53 +08:00

102 lines
2.7 KiB
TypeScript

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])
}
// 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<string[]> => {
const { labels, searchEngine } = await getSearchEngine()
let result: string[] = labels
if (query) {
// Search labels
result = searchEngine.search(query).map((r) => labels[r.id])
}
return result.length <= labelListLimit
? result
: [...result.slice(0, labelListLimit), `And ${result.length - labelListLimit} others`]
},
[getSearchEngine]
)
const setQueryLabel = useCallback((label: string) => {
if (label.startsWith('And ') && label.endsWith(' others')) return
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