feat: Add legend components and toggle buttons

lightrag_webui/src/components/graph/Legend.tsx

@@ -0,0 +1,41 @@
+import React from 'react'
+import { useTranslation } from 'react-i18next'
+import { useGraphStore } from '@/stores/graph'
+import { Card } from '@/components/ui/Card'
+import { ScrollArea } from '@/components/ui/ScrollArea'
+
+interface LegendProps {
+  className?: string
+}
+
+const Legend: React.FC<LegendProps> = ({ className }) => {
+  const { t } = useTranslation()
+  const typeColorMap = useGraphStore.use.typeColorMap()
+
+  if (!typeColorMap || typeColorMap.size === 0) {
+    return null
+  }
+
+  return (
+    <Card className={`p-2 max-w-xs ${className}`}>
+      <h3 className="text-sm font-medium mb-2">{t('graphPanel.legend')}</h3>
+      <ScrollArea className="max-h-40">
+        <div className="flex flex-col gap-1">
+          {Array.from(typeColorMap.entries()).map(([type, color]) => (
+            <div key={type} className="flex items-center gap-2">
+              <div
+                className="w-4 h-4 rounded-full"
+                style={{ backgroundColor: color }}
+              />
+              <span className="text-xs truncate" title={type}>
+                {type}
+              </span>
+            </div>
+          ))}
+        </div>
+      </ScrollArea>
+    </Card>
+  )
+}
+
+export default Legend

lightrag_webui/src/components/graph/LegendButton.tsx

@@ -0,0 +1,32 @@
+import { useCallback } from 'react'
+import { BookOpenIcon } from 'lucide-react'
+import Button from '@/components/ui/Button'
+import { controlButtonVariant } from '@/lib/constants'
+import { useSettingsStore } from '@/stores/settings'
+import { useTranslation } from 'react-i18next'
+
+/**
+ * Component that toggles legend visibility.
+ */
+const LegendButton = () => {
+  const { t } = useTranslation()
+  const showLegend = useSettingsStore.use.showLegend()
+  const setShowLegend = useSettingsStore.use.setShowLegend()
+
+  const toggleLegend = useCallback(() => {
+    setShowLegend(!showLegend)
+  }, [showLegend, setShowLegend])
+
+  return (
+    <Button
+      variant={controlButtonVariant}
+      onClick={toggleLegend}
+      tooltip={t('graphPanel.sideBar.legendControl.toggleLegend')}
+      size="icon"
+    >
+      <BookOpenIcon />
+    </Button>
+  )
+}
+
+export default LegendButton

lightrag_webui/src/features/GraphViewer.tsx

@@ -18,6 +18,8 @@ import GraphSearch from '@/components/graph/GraphSearch'
 import GraphLabels from '@/components/graph/GraphLabels'
 import PropertiesView from '@/components/graph/PropertiesView'
 import SettingsDisplay from '@/components/graph/SettingsDisplay'
+import Legend from '@/components/graph/Legend'
+import LegendButton from '@/components/graph/LegendButton'
 
 import { useSettingsStore } from '@/stores/settings'
 import { useGraphStore } from '@/stores/graph'
@@ -116,6 +118,7 @@ const GraphViewer = () => {
   const showPropertyPanel = useSettingsStore.use.showPropertyPanel()
   const showNodeSearchBar = useSettingsStore.use.showNodeSearchBar()
   const enableNodeDrag = useSettingsStore.use.enableNodeDrag()
+  const showLegend = useSettingsStore.use.showLegend()
 
   // Initialize sigma settings once on component mount
   // All dynamic settings will be updated in GraphControl using useSetSettings
@@ -195,6 +198,7 @@ const GraphViewer = () => {
           <LayoutsControl />
           <ZoomControl />
           <FullScreenControl />
+          <LegendButton />
           <Settings />
           {/* <ThemeToggle /> */}
         </div>
@@ -205,6 +209,12 @@ const GraphViewer = () => {
           </div>
         )}
 
+        {showLegend && (
+          <div className="absolute bottom-10 left-15">
+            <Legend className="bg-background/60 backdrop-blur-lg" />
+          </div>
+        )}
+
         {/* <div className="absolute bottom-2 right-2 flex flex-col rounded-xl border-2">
           <MiniMap width="100px" height="100px" />
         </div> */}

lightrag_webui/src/hooks/useLightragGraph.tsx

@@ -12,22 +12,31 @@ import { useSettingsStore } from '@/stores/settings'
 import seedrandom from 'seedrandom'
 
 // Helper function to generate a color based on type
-const getNodeColorByType = (nodeType: string | undefined, typeColorMap: React.RefObject<Map<string, string>>): string => {
+const getNodeColorByType = (nodeType: string | undefined): string => {
   const defaultColor = '#CCCCCC'; // Default color for nodes without a type or undefined type
   if (!nodeType) {
     return defaultColor;
   }
-  if (!typeColorMap.current.has(nodeType)) {
+
+  const typeColorMap = useGraphStore.getState().typeColorMap;
+
+  if (!typeColorMap.has(nodeType)) {
     // Generate a color based on the type string itself for consistency
     // Seed the global random number generator based on the node type
     seedrandom(nodeType, { global: true });
     // Call randomColor without arguments; it will use the globally seeded Math.random()
     const newColor = randomColor();
-    typeColorMap.current.set(nodeType, newColor);
+
+    const newMap = new Map(typeColorMap);
+    newMap.set(nodeType, newColor);
+    useGraphStore.setState({ typeColorMap: newMap });
+
+    return newColor;
   }
+
   // Restore the default random seed if necessary, though usually not required for this use case
   // seedrandom(Date.now().toString(), { global: true });
-  return typeColorMap.current.get(nodeType) || defaultColor; // Add fallback just in case
+  return typeColorMap.get(nodeType) || defaultColor; // Add fallback just in case
 };
 
 
@@ -240,8 +249,6 @@ const useLightrangeGraph = () => {
   const nodeToExpand = useGraphStore.use.nodeToExpand()
   const nodeToPrune = useGraphStore.use.nodeToPrune()
 
-    // Ref to store the mapping from node type to color
-    const typeColorMap = useRef<Map<string, string>>(new Map());
 
   // Use ref to track if data has been loaded and initial load
   const dataLoadedRef = useRef(false)
@@ -333,7 +340,7 @@ const useLightrangeGraph = () => {
           data.nodes.forEach(node => {
             // Use entity_type instead of type
             const nodeEntityType = node.properties?.entity_type as string | undefined;
-            node.color = getNodeColorByType(nodeEntityType, typeColorMap);
+            node.color = getNodeColorByType(nodeEntityType);
           });
         }
 
@@ -446,9 +453,9 @@ const useLightrangeGraph = () => {
         // Process nodes to add required properties for RawNodeType
         const processedNodes: RawNodeType[] = [];
         for (const node of extendedGraph.nodes) {
-          // Get color based on entity_type using the helper function and the shared map
+          // Get color based on entity_type using the helper function
-          const nodeEntityType = node.properties?.entity_type as string | undefined; // Use entity_type
+          const nodeEntityType = node.properties?.entity_type as string | undefined;
-          const color = getNodeColorByType(nodeEntityType, typeColorMap);
+          const color = getNodeColorByType(nodeEntityType);
 
           // Create a properly typed RawNodeType
           processedNodes.push({

lightrag_webui/src/locales/en.json

@@ -127,6 +127,7 @@
     }
   },
   "graphPanel": {
+    "legend": "Legend",
     "sideBar": {
       "settings": {
         "settings": "Settings",
@@ -171,6 +172,9 @@
       "fullScreenControl": {
         "fullScreen": "Full Screen",
         "windowed": "Windowed"
+      },
+      "legendControl": {
+        "toggleLegend": "Toggle Legend"
       }
     },
     "statusIndicator": {

lightrag_webui/src/locales/zh.json

@@ -127,6 +127,7 @@
     }
   },
   "graphPanel": {
+    "legend": "图例",
     "sideBar": {
       "settings": {
         "settings": "设置",
@@ -171,6 +172,9 @@
       "fullScreenControl": {
         "fullScreen": "全屏",
         "windowed": "窗口"
+      },
+      "legendControl": {
+        "toggleLegend": "切换图例显示"
       }
     },
     "statusIndicator": {

lightrag_webui/src/stores/graph.ts

@@ -77,6 +77,8 @@ interface GraphState {
   graphIsEmpty: boolean
   lastSuccessfulQueryLabel: string
 
+  typeColorMap: Map<string, string>
+
   // Global flags to track data fetching attempts
   graphDataFetchAttempted: boolean
   labelsFetchAttempted: boolean
@@ -136,6 +138,8 @@ const useGraphStoreBase = create<GraphState>()((set) => ({
   sigmaInstance: null,
   allDatabaseLabels: ['*'],
 
+  typeColorMap: new Map<string, string>(),
+
   searchEngine: null,
 
   setGraphIsEmpty: (isEmpty: boolean) => set({ graphIsEmpty: isEmpty }),
@@ -166,7 +170,6 @@ const useGraphStoreBase = create<GraphState>()((set) => ({
       searchEngine: null,
       moveToSelectedNode: false,
       graphIsEmpty: false
-      // Do not reset lastSuccessfulQueryLabel here as it's used to track query history
     });
   },
 
@@ -199,6 +202,8 @@ const useGraphStoreBase = create<GraphState>()((set) => ({
 
   setSigmaInstance: (instance: any) => set({ sigmaInstance: instance }),
 
+  setTypeColorMap: (typeColorMap: Map<string, string>) => set({ typeColorMap }),
+
   setSearchEngine: (engine: MiniSearch | null) => set({ searchEngine: engine }),
   resetSearchEngine: () => set({ searchEngine: null }),
 

lightrag_webui/src/stores/settings.ts

@@ -16,6 +16,8 @@ interface SettingsState {
   // Graph viewer settings
   showPropertyPanel: boolean
   showNodeSearchBar: boolean
+  showLegend: boolean
+  setShowLegend: (show: boolean) => void
 
   showNodeLabel: boolean
   enableNodeDrag: boolean
@@ -68,6 +70,7 @@ const useSettingsStoreBase = create<SettingsState>()(
       language: 'en',
       showPropertyPanel: true,
       showNodeSearchBar: true,
+      showLegend: false,
 
       showNodeLabel: true,
       enableNodeDrag: true,
@@ -145,7 +148,8 @@ const useSettingsStoreBase = create<SettingsState>()(
           querySettings: { ...state.querySettings, ...settings }
         })),
 
-      setShowFileName: (show: boolean) => set({ showFileName: show })
+      setShowFileName: (show: boolean) => set({ showFileName: show }),
+      setShowLegend: (show: boolean) => set({ showLegend: show })
     }),
     {
       name: 'settings-storage',