Add Node Expansion and Pruning Features

This commit is contained in:
yangdx
2025-03-14 15:58:50 +08:00
parent eb11a3db27
commit de016025ba
7 changed files with 854 additions and 152 deletions

View File

@@ -96,6 +96,14 @@ interface GraphState {
// Methods to set global flags
setGraphDataFetchAttempted: (attempted: boolean) => void
setLabelsFetchAttempted: (attempted: boolean) => void
// Event trigger methods for node operations
triggerNodeExpand: (nodeId: string | null) => void
triggerNodePrune: (nodeId: string | null) => void
// Node operation state
nodeToExpand: string | null
nodeToPrune: string | null
}
const useGraphStoreBase = create<GraphState>()((set, get) => ({
@@ -192,7 +200,267 @@ const useGraphStoreBase = create<GraphState>()((set, get) => ({
// Methods to set global flags
setGraphDataFetchAttempted: (attempted: boolean) => set({ graphDataFetchAttempted: attempted }),
setLabelsFetchAttempted: (attempted: boolean) => set({ labelsFetchAttempted: attempted })
setLabelsFetchAttempted: (attempted: boolean) => set({ labelsFetchAttempted: attempted }),
// Node operation state
nodeToExpand: null,
nodeToPrune: null,
// Event trigger methods for node operations
triggerNodeExpand: (nodeId: string | null) => set({ nodeToExpand: nodeId }),
triggerNodePrune: (nodeId: string | null) => set({ nodeToPrune: nodeId }),
// Legacy node expansion and pruning methods - will be removed after refactoring
expandNode: async (nodeId: string) => {
const state = get();
if (!state.sigmaGraph || !state.rawGraph || !nodeId) {
console.error('Cannot expand node: graph or node not available');
return;
}
try {
// Set fetching state
state.setIsFetching(true);
// Import queryGraphs dynamically to avoid circular dependency
const { queryGraphs } = await import('@/api/lightrag');
// Get the node to expand
const nodeToExpand = state.rawGraph.getNode(nodeId);
if (!nodeToExpand) {
console.error('Node not found:', nodeId);
state.setIsFetching(false);
return;
}
// Get the label of the node to expand
const label = nodeToExpand.labels[0];
if (!label) {
console.error('Node has no label:', nodeId);
state.setIsFetching(false);
return;
}
// Fetch the extended subgraph with depth 2
const extendedGraph = await queryGraphs(label, 2, 0);
if (!extendedGraph || !extendedGraph.nodes || !extendedGraph.edges) {
console.error('Failed to fetch extended graph');
state.setIsFetching(false);
return;
}
// Process nodes to add required properties for RawNodeType
const processedNodes: RawNodeType[] = [];
for (const node of extendedGraph.nodes) {
// Generate random color
const randomColorValue = () => Math.floor(Math.random() * 256);
const color = `rgb(${randomColorValue()}, ${randomColorValue()}, ${randomColorValue()})`;
// Create a properly typed RawNodeType
processedNodes.push({
id: node.id,
labels: node.labels,
properties: node.properties,
size: 10, // Default size
x: Math.random(), // Random position
y: Math.random(), // Random position
color: color, // Random color
degree: 0 // Initial degree
});
}
// Process edges to add required properties for RawEdgeType
const processedEdges: RawEdgeType[] = [];
for (const edge of extendedGraph.edges) {
// Create a properly typed RawEdgeType
processedEdges.push({
id: edge.id,
source: edge.source,
target: edge.target,
type: edge.type,
properties: edge.properties,
dynamicId: '' // Will be set when adding to sigma graph
});
}
// Store current node positions
const nodePositions: Record<string, {x: number, y: number}> = {};
state.sigmaGraph.forEachNode((node) => {
nodePositions[node] = {
x: state.sigmaGraph!.getNodeAttribute(node, 'x'),
y: state.sigmaGraph!.getNodeAttribute(node, 'y')
};
});
// Get existing node IDs
const existingNodeIds = new Set(state.sigmaGraph.nodes());
// Create a map from id to processed node for quick lookup
const processedNodeMap = new Map<string, RawNodeType>();
for (const node of processedNodes) {
processedNodeMap.set(node.id, node);
}
// Create a map from id to processed edge for quick lookup
const processedEdgeMap = new Map<string, RawEdgeType>();
for (const edge of processedEdges) {
processedEdgeMap.set(edge.id, edge);
}
// Add new nodes from the processed nodes
for (const newNode of processedNodes) {
// Skip if node already exists
if (existingNodeIds.has(newNode.id)) {
continue;
}
// Check if this node is connected to the selected node
const isConnected = processedEdges.some(
edge => (edge.source === nodeId && edge.target === newNode.id) ||
(edge.target === nodeId && edge.source === newNode.id)
);
if (isConnected) {
// Add the new node to the graph
state.sigmaGraph.addNode(newNode.id, {
label: newNode.labels.join(', '),
color: newNode.color,
x: nodePositions[nodeId].x + (Math.random() - 0.5) * 0.5,
y: nodePositions[nodeId].y + (Math.random() - 0.5) * 0.5,
size: newNode.size,
borderColor: '#000',
borderSize: 0.2
});
// Add the node to the raw graph
if (!state.rawGraph.getNode(newNode.id)) {
// Add to nodes array
state.rawGraph.nodes.push(newNode);
// Update nodeIdMap
state.rawGraph.nodeIdMap[newNode.id] = state.rawGraph.nodes.length - 1;
}
}
}
// Add new edges
for (const newEdge of processedEdges) {
// Only add edges where both source and target exist in the graph
if (state.sigmaGraph.hasNode(newEdge.source) && state.sigmaGraph.hasNode(newEdge.target)) {
// Skip if edge already exists
if (state.sigmaGraph.hasEdge(newEdge.source, newEdge.target)) {
continue;
}
// Add the edge to the sigma graph
newEdge.dynamicId = state.sigmaGraph.addDirectedEdge(newEdge.source, newEdge.target, {
label: newEdge.type || undefined
});
// Add the edge to the raw graph
if (!state.rawGraph.getEdge(newEdge.id, false)) {
// Add to edges array
state.rawGraph.edges.push(newEdge);
// Update edgeIdMap
state.rawGraph.edgeIdMap[newEdge.id] = state.rawGraph.edges.length - 1;
// Update dynamic edge map
state.rawGraph.edgeDynamicIdMap[newEdge.dynamicId] = state.rawGraph.edges.length - 1;
}
}
}
// Update the dynamic edge map
state.rawGraph.buildDynamicMap();
// Restore positions for existing nodes
Object.entries(nodePositions).forEach(([nodeId, position]) => {
if (state.sigmaGraph!.hasNode(nodeId)) {
state.sigmaGraph!.setNodeAttribute(nodeId, 'x', position.x);
state.sigmaGraph!.setNodeAttribute(nodeId, 'y', position.y);
}
});
// Refresh the layout
state.refreshLayout();
} catch (error) {
console.error('Error expanding node:', error);
} finally {
// Reset fetching state
state.setIsFetching(false);
}
},
pruneNode: (nodeId: string) => {
const state = get();
if (!state.sigmaGraph || !state.rawGraph || !nodeId) {
console.error('Cannot prune node: graph or node not available');
return;
}
try {
// Check if the node exists
if (!state.sigmaGraph.hasNode(nodeId)) {
console.error('Node not found:', nodeId);
return;
}
// If the node is selected or focused, clear selection
if (state.selectedNode === nodeId || state.focusedNode === nodeId) {
state.clearSelection();
}
// Remove the node from the sigma graph (this will also remove connected edges)
state.sigmaGraph.dropNode(nodeId);
// Remove the node from the raw graph
const nodeIndex = state.rawGraph.nodeIdMap[nodeId];
if (nodeIndex !== undefined) {
// Find all edges connected to this node
const edgesToRemove = state.rawGraph.edges.filter(
edge => edge.source === nodeId || edge.target === nodeId
);
// Remove edges from raw graph
for (const edge of edgesToRemove) {
const edgeIndex = state.rawGraph.edgeIdMap[edge.id];
if (edgeIndex !== undefined) {
// Remove from edges array
state.rawGraph.edges.splice(edgeIndex, 1);
// Update edgeIdMap for all edges after this one
for (const [id, idx] of Object.entries(state.rawGraph.edgeIdMap)) {
if (idx > edgeIndex) {
state.rawGraph.edgeIdMap[id] = idx - 1;
}
}
// Remove from edgeIdMap
delete state.rawGraph.edgeIdMap[edge.id];
// Remove from edgeDynamicIdMap
delete state.rawGraph.edgeDynamicIdMap[edge.dynamicId];
}
}
// Remove node from nodes array
state.rawGraph.nodes.splice(nodeIndex, 1);
// Update nodeIdMap for all nodes after this one
for (const [id, idx] of Object.entries(state.rawGraph.nodeIdMap)) {
if (idx > nodeIndex) {
state.rawGraph.nodeIdMap[id] = idx - 1;
}
}
// Remove from nodeIdMap
delete state.rawGraph.nodeIdMap[nodeId];
// Rebuild the dynamic edge map
state.rawGraph.buildDynamicMap();
}
} catch (error) {
console.error('Error pruning node:', error);
}
}
}))
const useGraphStore = createSelectors(useGraphStoreBase)