import { useState, useCallback, useEffect, memo } from 'react'
import { useNavigate } from '@tanstack/react-router'
import ReactFlow, {
Controls,
Background,
BackgroundVariant,
MiniMap,
useNodesState,
useEdgesState,
Panel,
Handle,
Position,
type Node,
type Edge,
type NodeTypes,
} from 'reactflow'
import 'reactflow/dist/style.css'
import dagre from 'dagre'
import { Button } from '@/components/ui/button'
import { Input } from '@/components/ui/input'
import { Badge } from '@/components/ui/badge'
import { ScrollArea } from '@/components/ui/scroll-area'
import {
Dialog,
DialogContent,
DialogHeader,
DialogTitle,
} from '@/components/ui/dialog'
import {
AlertDialog,
AlertDialogAction,
AlertDialogCancel,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
} from '@/components/ui/alert-dialog'
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from '@/components/ui/select'
import {
Search,
RefreshCw,
Info,
Database,
Network,
FileText,
} from 'lucide-react'
import { useToast } from '@/hooks/use-toast'
import { getKnowledgeGraph, getKnowledgeStats, searchKnowledgeNode, type KnowledgeNode, type KnowledgeEdge, type KnowledgeStats } from '@/lib/knowledge-api'
import { cn } from '@/lib/utils'
// 自定义节点组件 - 实体节点
const EntityNode = memo(({ data }: { data: { label: string; content: string } }) => {
return (
)
})
EntityNode.displayName = 'EntityNode'
// 自定义节点组件 - 段落节点
const ParagraphNode = memo(({ data }: { data: { label: string; content: string } }) => {
return (
)
})
ParagraphNode.displayName = 'ParagraphNode'
const nodeTypes: NodeTypes = {
entity: EntityNode,
paragraph: ParagraphNode,
}
// 使用 dagre 进行自动布局
function calculateLayout(nodes: KnowledgeNode[], edges: KnowledgeEdge[]): { nodes: Node[]; edges: Edge[] } {
const dagreGraph = new dagre.graphlib.Graph()
dagreGraph.setDefaultEdgeLabel(() => ({}))
dagreGraph.setGraph({ rankdir: 'TB', ranksep: 100, nodesep: 80 })
const flowNodes: Node[] = []
const flowEdges: Edge[] = []
// 设置节点到 dagre 图
nodes.forEach((node) => {
dagreGraph.setNode(node.id, { width: 150, height: 50 })
})
// 设置边到 dagre 图
edges.forEach((edge) => {
dagreGraph.setEdge(edge.source, edge.target)
})
// 执行布局计算
dagre.layout(dagreGraph)
// 获取布局后的节点位置
nodes.forEach((node) => {
const nodeWithPosition = dagreGraph.node(node.id)
flowNodes.push({
id: node.id,
type: node.type,
position: {
x: nodeWithPosition.x - 75,
y: nodeWithPosition.y - 25,
},
data: {
label: node.content.slice(0, 20) + (node.content.length > 20 ? '...' : ''),
content: node.content,
},
})
})
// 创建边
edges.forEach((edge, index) => {
const flowEdge: Edge = {
id: `edge-${index}`,
source: edge.source,
target: edge.target,
// 节点数超过200时禁用动画提升性能
animated: nodes.length <= 200 && edge.weight > 5,
style: {
strokeWidth: Math.min(edge.weight / 2, 5),
opacity: 0.6,
},
}
// 只在节点数少于100时显示边的标签
if (edge.weight > 10 && nodes.length < 100) {
flowEdge.label = `${edge.weight.toFixed(0)}`
}
flowEdges.push(flowEdge)
})
return { nodes: flowNodes, edges: flowEdges }
}
export function KnowledgeGraphPage() {
const navigate = useNavigate()
const [loading, setLoading] = useState(false)
const [stats, setStats] = useState(null)
const [searchQuery, setSearchQuery] = useState('')
const [nodeType, setNodeType] = useState<'all' | 'entity' | 'paragraph'>('all')
const [nodeLimit, setNodeLimit] = useState(50)
const [customLimit, setCustomLimit] = useState('50')
const [showCustomInput, setShowCustomInput] = useState(false)
const [showInitialConfirm, setShowInitialConfirm] = useState(true)
const [userConfirmedLoad, setUserConfirmedLoad] = useState(false) // 用户是否确认加载
const [showHighNodeWarning, setShowHighNodeWarning] = useState(false)
const [nodes, setNodes, onNodesChange] = useNodesState([])
const [edges, setEdges, onEdgesChange] = useEdgesState([])
const [nodeCount, setNodeCount] = useState(0)
const [selectedNodeData, setSelectedNodeData] = useState(null)
const [selectedEdgeData, setSelectedEdgeData] = useState<{ source: KnowledgeNode; target: KnowledgeNode; edge: KnowledgeEdge } | null>(null)
const { toast } = useToast()
// 缓存 MiniMap 的 nodeColor 函数
const miniMapNodeColor = useCallback((node: Node) => {
if (node.type === 'entity') return '#6366f1'
if (node.type === 'paragraph') return '#10b981'
return '#6b7280'
}, [])
// 加载知识图谱数据
const loadGraph = useCallback(async (skipWarning = false) => {
try {
// 检查是否需要警告用户
if (!skipWarning && nodeLimit > 200) {
setShowHighNodeWarning(true)
return
}
setLoading(true)
const [graphData, statsData] = await Promise.all([
getKnowledgeGraph(nodeLimit, nodeType),
getKnowledgeStats(),
])
setStats(statsData)
if (graphData.nodes.length === 0) {
toast({
title: '提示',
description: '知识库为空,请先导入知识数据',
})
setNodes([])
setEdges([])
return
}
const { nodes: flowNodes, edges: flowEdges } = calculateLayout(graphData.nodes, graphData.edges)
setNodes(flowNodes)
setEdges(flowEdges)
setNodeCount(flowNodes.length)
if (statsData && statsData.total_nodes > nodeLimit) {
toast({
title: '提示',
description: `知识图谱包含 ${statsData.total_nodes} 个节点,当前显示 ${flowNodes.length} 个`,
})
}
toast({
title: '加载成功',
description: `已加载 ${flowNodes.length} 个节点,${flowEdges.length} 条边`,
})
} catch (error) {
console.error('加载知识图谱失败:', error)
toast({
title: '加载失败',
description: error instanceof Error ? error.message : '未知错误',
variant: 'destructive',
})
} finally {
setLoading(false)
}
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [nodeLimit, nodeType, toast]) // setNodes 和 setEdges 是稳定的,不需要包含
// 搜索节点
const handleSearch = useCallback(async () => {
if (!searchQuery.trim()) {
toast({
title: '提示',
description: '请输入搜索关键词',
})
return
}
try {
const results = await searchKnowledgeNode(searchQuery)
if (results.length === 0) {
toast({
title: '未找到',
description: '没有找到匹配的节点',
})
return
}
// 高亮搜索结果
const resultIds = new Set(results.map(r => r.id))
setNodes(nds =>
nds.map(node => ({
...node,
style: {
...node.style,
opacity: resultIds.has(node.id) ? 1 : 0.3,
filter: resultIds.has(node.id) ? 'brightness(1.2)' : 'brightness(0.8)',
},
}))
)
toast({
title: '搜索完成',
description: `找到 ${results.length} 个匹配节点`,
})
} catch (error) {
console.error('搜索失败:', error)
toast({
title: '搜索失败',
description: error instanceof Error ? error.message : '未知错误',
variant: 'destructive',
})
}
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [searchQuery, toast]) // setNodes 是稳定的
// 重置高亮
const handleResetHighlight = useCallback(() => {
setNodes(nds =>
nds.map(node => ({
...node,
style: {
...node.style,
opacity: 1,
filter: 'brightness(1)',
},
}))
)
// eslint-disable-next-line react-hooks/exhaustive-deps
}, []) // setNodes 是稳定的
// 初始确认后加载
const handleInitialConfirm = useCallback(() => {
setShowInitialConfirm(false)
setUserConfirmedLoad(true) // 设置用户确认标记
loadGraph()
}, [loadGraph])
// 高节点数确认后加载
const handleHighNodeConfirm = useCallback(() => {
setShowHighNodeWarning(false) // 立即关闭高节点数警告对话框
// 使用 setTimeout 确保对话框关闭后再开始加载
setTimeout(() => {
loadGraph(true)
}, 0)
}, [loadGraph])
// 节点点击事件
const onNodeClick = useCallback((_: React.MouseEvent, node: Node) => {
setSelectedNodeData({
id: node.id,
type: node.type as 'entity' | 'paragraph',
content: node.data.content,
})
}, [])
// 当节点数量或类型改变时自动刷新
useEffect(() => {
// 跳过初始确认对话框时的加载
if (showInitialConfirm) return
// 只有用户确认后才能自动刷新
if (!userConfirmedLoad) return
// 参数变化时加载,会根据节点数自动判断是否需要警告
loadGraph()
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [nodeLimit, nodeType, showInitialConfirm, userConfirmedLoad]) // 不依赖 loadGraph
// 边点击事件
const onEdgeClick = useCallback((_: React.MouseEvent, edge: Edge) => {
const sourceNode = nodes.find(n => n.id === edge.source)
const targetNode = nodes.find(n => n.id === edge.target)
const edgeData = edges.find(e => e.id === edge.id)
if (sourceNode && targetNode && edgeData) {
setSelectedEdgeData({
source: {
id: sourceNode.id,
type: sourceNode.type as 'entity' | 'paragraph',
content: sourceNode.data.content,
},
target: {
id: targetNode.id,
type: targetNode.type as 'entity' | 'paragraph',
content: targetNode.data.content,
},
edge: {
source: edge.source,
target: edge.target,
weight: parseFloat(edge.label as string || '0'),
},
})
}
}, [nodes, edges])
return (
{/* 顶部工具栏 */}
{stats && (
节点: {stats.total_nodes}
边: {stats.total_edges}
实体: {stats.entity_nodes}
段落: {stats.paragraph_nodes}
)}
{/* 搜索和控制栏 */}
{/* 主内容区域 */}
{loading ? (
) : nodes.length === 0 ? (
) : (
{/* 节点数超过500时禁用MiniMap提升性能 */}
{nodeCount <= 500 && (
)}
{/* 图例 */}
图例
{nodeCount > 200 && (
性能模式
已禁用动画
{nodeCount > 500 &&
已禁用缩略图
}
)}
)}
{/* 节点详情对话框 */}
{/* 边详情对话框 */}
{/* 初始加载确认对话框 */}
加载知识图谱
知识图谱的动态展示会消耗较多系统资源。
确定要加载知识图谱吗?
navigate({ to: '/' })}>
取消 (返回首页)
确认加载
{/* 高节点数警告对话框 */}
⚠️ 节点数量较多
您正在尝试加载 {nodeLimit >= 10000 ? '全部 (最多10000个)' : nodeLimit} 个节点。
节点数量过多可能导致:
- 页面加载时间较长
- 浏览器卡顿或崩溃
- 系统资源占用过高
建议先选择较少的节点数量 (50-200 个)。
{
setShowHighNodeWarning(false)
// 将节点数重置为安全值
if (nodeLimit > 200) {
setNodeLimit(50)
setShowCustomInput(false)
}
}}>
取消
我了解风险,继续加载
)
}