diff --git a/packages/memory-graph/src/__tests__/cluster-level-graph.test.ts b/packages/memory-graph/src/__tests__/cluster-level-graph.test.ts new file mode 100644 index 00000000..5e32bfe2 --- /dev/null +++ b/packages/memory-graph/src/__tests__/cluster-level-graph.test.ts @@ -0,0 +1,133 @@ +import { describe, expect, it } from "vitest" +import { buildClusterLevelGraph } from "../hooks/cluster-level-graph" +import type { GraphEdge, GraphNode, MemoryNodeData } from "../types" + +function makeDocumentNode(index: number): GraphNode { + const id = `doc-${index}` + return { + id, + type: "document", + x: index * 120, + y: 0, + data: { + id, + title: id, + summary: null, + type: "text", + createdAt: "2024-01-01", + updatedAt: "2024-01-01", + memories: [], + }, + size: 50, + borderColor: "#58C7E8", + clusterColor: "#58C7E8", + isHovered: false, + isDragging: false, + } +} + +function makeMemoryNode(documentId: string, index: number): GraphNode { + const id = `${documentId}-mem-${index}` + const data: MemoryNodeData = { + id, + memory: id, + content: id, + documentId, + isStatic: false, + isLatest: true, + isForgotten: false, + forgetAfter: null, + forgetReason: null, + version: 1, + parentMemoryId: null, + spaceId: "space", + createdAt: "2024-01-01", + updatedAt: "2024-01-01", + } + + return { + id, + type: "memory", + x: index * 80, + y: 120, + data, + size: 36, + borderColor: "#58C7E8", + clusterColor: "#58C7E8", + isHovered: false, + isDragging: false, + } +} + +function makeGraph(documentCount: number, memoriesPerDocument: number) { + const nodes: GraphNode[] = [] + const edges: GraphEdge[] = [] + + for (let docIndex = 0; docIndex < documentCount; docIndex++) { + const doc = makeDocumentNode(docIndex) + nodes.push(doc) + + for (let memIndex = 0; memIndex < memoriesPerDocument; memIndex++) { + const memory = makeMemoryNode(doc.id, memIndex) + nodes.push(memory) + edges.push({ + id: `dm-${doc.id}-${memory.id}`, + source: doc.id, + target: memory.id, + edgeType: "derives", + visualProps: { opacity: 0.4, thickness: 1.2 }, + }) + } + } + + return { nodes, edges } +} + +describe("buildClusterLevelGraph", () => { + it("collapses dense raw graphs into aggregate cluster nodes", () => { + const graph = makeGraph(500, 2) + const clustered = buildClusterLevelGraph({ + ...graph, + enabled: true, + expandedClusterId: null, + revealRawGraph: false, + canvasWidth: 1200, + canvasHeight: 800, + }) + + expect(clustered.isClustered).toBe(true) + expect(clustered.nodes.length).toBeLessThan(graph.nodes.length) + expect(clustered.nodes.every((node) => node.type === "cluster")).toBe(true) + expect(clustered.nodes[0]?.data).toMatchObject({ + documentCount: 50, + memoryCount: 100, + nodeCount: 150, + }) + }) + + it("expands a selected cluster while keeping other groups collapsed", () => { + const graph = makeGraph(120, 2) + const collapsed = buildClusterLevelGraph({ + ...graph, + enabled: true, + expandedClusterId: null, + revealRawGraph: false, + canvasWidth: 1200, + canvasHeight: 800, + }) + const clusterId = collapsed.nodes[0]?.id ?? null + + const expanded = buildClusterLevelGraph({ + ...graph, + enabled: true, + expandedClusterId: clusterId, + revealRawGraph: false, + canvasWidth: 1200, + canvasHeight: 800, + }) + + expect(expanded.nodes.some((node) => node.type === "document")).toBe(true) + expect(expanded.nodes.some((node) => node.type === "memory")).toBe(true) + expect(expanded.nodes.some((node) => node.type === "cluster")).toBe(true) + }) +}) diff --git a/packages/memory-graph/src/canvas/renderer.ts b/packages/memory-graph/src/canvas/renderer.ts index cc396bd8..f25f93e2 100644 --- a/packages/memory-graph/src/canvas/renderer.ts +++ b/packages/memory-graph/src/canvas/renderer.ts @@ -1,4 +1,5 @@ import type { + ClusterNodeData, DocumentNodeData, GraphEdge, GraphNode, @@ -28,6 +29,11 @@ function nodeMatchesDocumentHighlights( ): boolean { if (highlightIds.size === 0) return false if (node.type === "document") return highlightIds.has(node.id) + if (node.type === "cluster") { + return (node.data as ClusterNodeData).sampleDocumentIds.some((id) => + highlightIds.has(id), + ) + } return highlightIds.has((node.data as MemoryNodeData).documentId) } @@ -438,7 +444,7 @@ function drawNodes( highlightFocus && !isSelected && !isHovered && !isHighlighted if (screenSize < 8 && !isSelected && !isHovered && !isHighlighted) { - if (node.type === "document") { + if (node.type === "document" || node.type === "cluster") { docDots.push({ x: screen.x, y: screen.y, s: Math.max(3, screenSize) }) } else { const md = node.data as MemoryNodeData @@ -465,7 +471,19 @@ function drawNodes( } ctx.globalAlpha = alpha - if (node.type === "document") { + if (node.type === "cluster") { + drawClusterNode( + ctx, + screen.x, + screen.y, + screenSize, + node, + isSelected, + isHovered, + isHighlighted, + colors, + ) + } else if (node.type === "document") { drawDocumentNode( ctx, screen.x, @@ -631,6 +649,82 @@ function drawNodes( ctx.globalAlpha = 1 } +function drawClusterNode( + ctx: CanvasRenderingContext2D, + sx: number, + sy: number, + size: number, + node: GraphNode, + isSelected: boolean, + isHovered: boolean, + isHighlighted: boolean, + colors: GraphThemeColors, +): void { + const data = node.data as ClusterNodeData + const radius = size * 0.5 + const clusterColor = node.clusterColor ?? colors.accent + + if (isSelected || isHovered) { + ctx.save() + ctx.shadowColor = isSelected ? colors.accent : clusterColor + ctx.shadowBlur = isSelected ? 24 : 16 + ctx.shadowOffsetX = 0 + ctx.shadowOffsetY = 0 + } + + const gradient = ctx.createRadialGradient( + sx - radius * 0.25, + sy - radius * 0.3, + radius * 0.1, + sx, + sy, + radius, + ) + gradient.addColorStop( + 0, + mixHexColors(colors.memFillHover, clusterColor, 0.36), + ) + gradient.addColorStop(0.62, mixHexColors(colors.docFill, clusterColor, 0.24)) + gradient.addColorStop(1, mixHexColors(colors.bg, clusterColor, 0.58)) + ctx.fillStyle = gradient + ctx.beginPath() + ctx.arc(sx, sy, radius, 0, Math.PI * 2) + ctx.fill() + + ctx.strokeStyle = isSelected || isHighlighted ? colors.accent : clusterColor + ctx.lineWidth = isSelected || isHighlighted ? 3 : isHovered ? 2.5 : 1.8 + ctx.stroke() + + ctx.globalAlpha *= 0.45 + ctx.strokeStyle = clusterColor + ctx.lineWidth = Math.max(1, size * 0.018) + for (const scale of [0.68, 0.44]) { + ctx.beginPath() + ctx.arc(sx, sy, radius * scale, 0, Math.PI * 2) + ctx.stroke() + } + ctx.globalAlpha /= 0.45 + + if (isSelected || isHovered) { + ctx.restore() + } + + if (size >= 58) { + const count = data.documentCount.toLocaleString() + const label = data.memoryCount.toLocaleString() + ctx.save() + ctx.textAlign = "center" + ctx.textBaseline = "middle" + ctx.fillStyle = colors.textPrimary + ctx.font = `600 ${Math.max(12, Math.min(18, size * 0.18))}px system-ui, sans-serif` + ctx.fillText(count, sx, sy - Math.max(5, size * 0.07)) + ctx.fillStyle = colors.textSecondary + ctx.font = `500 ${Math.max(9, Math.min(12, size * 0.11))}px system-ui, sans-serif` + ctx.fillText(`${label} mem`, sx, sy + Math.max(9, size * 0.1)) + ctx.restore() + } +} + function drawDocumentNode( ctx: CanvasRenderingContext2D, sx: number, @@ -850,7 +944,7 @@ function drawGlow( sx: number, sy: number, size: number, - nodeType: "document" | "memory", + nodeType: "document" | "memory" | "cluster", colors: GraphThemeColors, isHoverOnly = false, ): void { @@ -866,6 +960,9 @@ function drawGlow( const half = glowSize * 0.5 const r = 8 * (glowSize / 50) roundRect(ctx, sx - half, sy - half, glowSize, glowSize, r) + } else if (nodeType === "cluster") { + ctx.beginPath() + ctx.arc(sx, sy, size * 0.5 * scale, 0, Math.PI * 2) } else { drawHexagon(ctx, sx, sy, size * 0.5 * scale) } diff --git a/packages/memory-graph/src/canvas/simulation.ts b/packages/memory-graph/src/canvas/simulation.ts index 4182f5e2..1348f80e 100644 --- a/packages/memory-graph/src/canvas/simulation.ts +++ b/packages/memory-graph/src/canvas/simulation.ts @@ -50,11 +50,12 @@ export class ForceSimulation { "collide", d3 .forceCollide() - .radius((d) => - d.type === "document" - ? FORCE_CONFIG.collisionRadius.document - : FORCE_CONFIG.collisionRadius.memory, - ) + .radius((d) => { + if (d.type === "document") + return FORCE_CONFIG.collisionRadius.document + if (d.type === "cluster") return d.size * 0.62 + return FORCE_CONFIG.collisionRadius.memory + }) .strength(FORCE_CONFIG.collisionStrength), ) diff --git a/packages/memory-graph/src/components/legend.tsx b/packages/memory-graph/src/components/legend.tsx index d007ac56..8e252e42 100644 --- a/packages/memory-graph/src/components/legend.tsx +++ b/packages/memory-graph/src/components/legend.tsx @@ -1,5 +1,10 @@ import { memo, useState } from "react" -import type { GraphEdge, GraphNode, GraphThemeColors } from "../types" +import type { + ClusterNodeData, + GraphEdge, + GraphNode, + GraphThemeColors, +} from "../types" interface LegendProps { nodes?: GraphNode[] @@ -294,19 +299,32 @@ export const Legend = memo(function Legend({ const [isExpanded, setIsExpanded] = useState(false) const [connectionsExpanded, setConnectionsExpanded] = useState(true) - const memoryCount = nodes.filter((n) => n.type === "memory").length - const documentCount = nodes.filter((n) => n.type === "document").length + const clusterNodes = nodes.filter((n) => n.type === "cluster") + const clusterMemoryCount = clusterNodes.reduce( + (total, node) => total + (node.data as ClusterNodeData).memoryCount, + 0, + ) + const clusterDocumentCount = clusterNodes.reduce( + (total, node) => total + (node.data as ClusterNodeData).documentCount, + 0, + ) + const memoryCount = + nodes.filter((n) => n.type === "memory").length + clusterMemoryCount + const documentCount = + nodes.filter((n) => n.type === "document").length + clusterDocumentCount const connectionCount = edges.length const derivesCount = countEdgesByType(edges, "derives") const updatesCount = countEdgesByType(edges, "updates") const extendsCount = countEdgesByType(edges, "extends") const activeUpdateCount = getActiveUpdateCount(edges, hoveredNode) - const clusterCount = new Set( - nodes - .filter((node) => node.type === "memory") - .map((node) => node.clusterKey) - .filter(Boolean), - ).size + const clusterCount = + clusterNodes.length || + new Set( + nodes + .filter((node) => node.type === "memory") + .map((node) => node.clusterKey) + .filter(Boolean), + ).size const updateNodeCount = nodes.filter(isUpdateMemoryNode).length const outerStyle: React.CSSProperties = { diff --git a/packages/memory-graph/src/components/memory-graph.tsx b/packages/memory-graph/src/components/memory-graph.tsx index 0ec65ed2..69ddd17c 100644 --- a/packages/memory-graph/src/components/memory-graph.tsx +++ b/packages/memory-graph/src/components/memory-graph.tsx @@ -2,6 +2,7 @@ import { useCallback, useEffect, useMemo, useRef, useState } from "react" import { ForceSimulation } from "../canvas/simulation" import { VersionChainIndex } from "../canvas/version-chain" import type { ViewportState } from "../canvas/viewport" +import { buildClusterLevelGraph } from "../hooks/cluster-level-graph" import { useGraphData } from "../hooks/use-graph-data" import { useGraphTheme } from "../hooks/use-graph-theme" import type { @@ -15,6 +16,9 @@ import { LoadingIndicator } from "./loading-indicator" import { NavigationControls } from "./navigation-controls" import { NodeHoverPopover } from "./node-hover-popover" +const CLUSTER_MODE_NODE_THRESHOLD = 3000 +const RAW_GRAPH_REVEAL_ZOOM = 0.62 + export function MemoryGraph({ documents = [], isLoading: externalIsLoading = false, @@ -55,6 +59,9 @@ export function MemoryGraph({ // React state only for things that affect DOM const [hoveredNode, setHoveredNode] = useState(null) const [selectedNode, setSelectedNode] = useState(null) + const [expandedClusterId, setExpandedClusterId] = useState( + null, + ) const [zoomDisplay, setZoomDisplay] = useState(50) // Monotonic counter that increments on any viewport change (pan or zoom) // Used as a dependency proxy to recalculate popover positions @@ -97,13 +104,38 @@ export function MemoryGraph({ const hasContainerSize = containerSize.width > 0 && containerSize.height > 0 - const { nodes, edges } = useGraphData( + const { nodes: rawNodes, edges: rawEdges } = useGraphData( hasContainerSize ? limitedDocuments : [], null, containerSize.width, containerSize.height, colors, ) + const revealRawGraph = + rawNodes.length <= CLUSTER_MODE_NODE_THRESHOLD || + zoomDisplay / 100 >= RAW_GRAPH_REVEAL_ZOOM + const clusterGraph = useMemo( + () => + buildClusterLevelGraph({ + nodes: rawNodes, + edges: rawEdges, + enabled: rawNodes.length > CLUSTER_MODE_NODE_THRESHOLD, + expandedClusterId, + revealRawGraph, + canvasWidth: containerSize.width, + canvasHeight: containerSize.height, + }), + [ + rawNodes, + rawEdges, + expandedClusterId, + revealRawGraph, + containerSize.width, + containerSize.height, + ], + ) + const nodes = clusterGraph.nodes + const edges = clusterGraph.edges const isCompactViewport = containerSize.width > 0 && containerSize.width < 640 const graphFitHeight = isCompactViewport ? Math.max(containerSize.height - 170, 240) @@ -231,6 +263,10 @@ export function MemoryGraph({ if (nodes.length === 0) hasAutoFittedRef.current = false }, [nodes.length]) + useEffect(() => { + if (!clusterGraph.isClustered) setExpandedClusterId(null) + }, [clusterGraph.isClustered]) + useEffect(() => { if (isCompactViewport) { hasAutoFittedRef.current = false @@ -278,9 +314,32 @@ export function MemoryGraph({ [], ) - const handleNodeClick = useCallback((id: string | null) => { - setSelectedNode((prev) => (id === null ? null : prev === id ? null : id)) - }, []) + const handleNodeClick = useCallback( + (id: string | null) => { + if (id === null) { + setSelectedNode(null) + setExpandedClusterId(null) + return + } + + const node = nodes.find((n) => n.id === id) + if (node?.type === "cluster") { + setSelectedNode(null) + setHoveredNode(null) + setExpandedClusterId((prev) => (prev === id ? null : id)) + viewportRef.current?.centerOn( + node.x, + node.y, + containerSize.width, + graphFitHeight, + ) + return + } + + setSelectedNode((prev) => (prev === id ? null : id)) + }, + [nodes, containerSize.width, graphFitHeight], + ) const handleNodeDragStart = useCallback((_id: string) => { // Drag is handled imperatively by InputHandler @@ -497,7 +556,12 @@ export function MemoryGraph({ const node = nodes.find((n) => n.id === selectedNode) if (!node) return - if (node.type === "document") { + if (node.type === "cluster") { + const clusters = nodes.filter((n) => n.type === "cluster") + const idx = clusters.findIndex((n) => n.id === selectedNode) + const next = clusters[(idx + 1) % clusters.length] + if (next) selectAndCenter(next.id) + } else if (node.type === "document") { const docs = nodes.filter((n) => n.type === "document") const idx = docs.findIndex((n) => n.id === selectedNode) const next = docs[(idx + 1) % docs.length] @@ -522,7 +586,12 @@ export function MemoryGraph({ const node = nodes.find((n) => n.id === selectedNode) if (!node) return - if (node.type === "document") { + if (node.type === "cluster") { + const clusters = nodes.filter((n) => n.type === "cluster") + const idx = clusters.findIndex((n) => n.id === selectedNode) + const prev = clusters[(idx - 1 + clusters.length) % clusters.length] + if (prev) selectAndCenter(prev.id) + } else if (node.type === "document") { const docs = nodes.filter((n) => n.type === "document") const idx = docs.findIndex((n) => n.id === selectedNode) const prev = docs[(idx - 1 + docs.length) % docs.length] diff --git a/packages/memory-graph/src/components/node-hover-popover.tsx b/packages/memory-graph/src/components/node-hover-popover.tsx index 97ec06e8..53a1645f 100644 --- a/packages/memory-graph/src/components/node-hover-popover.tsx +++ b/packages/memory-graph/src/components/node-hover-popover.tsx @@ -1,6 +1,7 @@ import { memo, useCallback, useEffect, useMemo, useRef, useState } from "react" import type { ChainEntry } from "../canvas/version-chain" import type { + ClusterNodeData, DocumentNodeData, GraphNode, GraphThemeColors, @@ -323,6 +324,7 @@ export const NodeHoverPopover = memo( const TOTAL_W = CARD_W + 12 + SHORTCUTS_W const isMemory = node.type === "memory" + const isCluster = node.type === "cluster" const data = node.data const hasChain = Boolean(versionChain && versionChain.length > 1) @@ -399,18 +401,27 @@ export const NodeHoverPopover = memo( }, [screenX, screenY, nodeRadius, containerBounds, TOTAL_W, TOTAL_H]) const content = useMemo(() => { + if (isCluster) { + const cd = data as ClusterNodeData + return cd.summary || cd.title + } if (isMemory) { const md = data as MemoryNodeData return md.memory || md.content || "" } const dd = data as DocumentNodeData return dd.summary || dd.title || "" - }, [isMemory, data]) + }, [isCluster, isMemory, data]) - const docData = !isMemory ? (data as DocumentNodeData) : null + const docData = node.type === "document" ? (data as DocumentNodeData) : null + const clusterData = isCluster ? (data as ClusterNodeData) : null // For document nodes, node.id IS the document ID - const documentId = isMemory ? (data as MemoryNodeData).documentId : node.id + const documentId = isCluster + ? null + : isMemory + ? (data as MemoryNodeData).documentId + : node.id const overlayStyle: React.CSSProperties = { pointerEvents: "none", @@ -578,7 +589,26 @@ export const NodeHoverPopover = memo( )}
- {memoryMeta ? ( + {clusterData ? ( + <> + + {clusterData.documentCount} documents + + + {clusterData.memoryCount} memories + + + ) : memoryMeta ? ( (
- {isMemory ? ( + {isCluster ? ( + + ) : isMemory ? ( ) : ( @@ -645,7 +677,7 @@ export const NodeHoverPopover = memo( onClick={onNavigateUp} /> )} - {(isMemory ? hasChain : true) && ( + {(isMemory ? hasChain : !isCluster) && ( (
diff --git a/packages/memory-graph/src/hooks/cluster-level-graph.ts b/packages/memory-graph/src/hooks/cluster-level-graph.ts new file mode 100644 index 00000000..9e0b2e42 --- /dev/null +++ b/packages/memory-graph/src/hooks/cluster-level-graph.ts @@ -0,0 +1,382 @@ +import type { + ClusterNodeData, + GraphEdge, + GraphNode, + MemoryNodeData, +} from "../types" + +const CLUSTER_TARGET_RAW_NODES = 150 +const CLUSTER_MIN_SIZE = 68 +const CLUSTER_MAX_SIZE = 142 +const GOLDEN_ANGLE = Math.PI * (3 - Math.sqrt(5)) + +export interface ClusterLevelGraph { + nodes: GraphNode[] + edges: GraphEdge[] + isClustered: boolean + clusterCount: number + rawNodeToClusterId: Map +} + +interface ClusterGroup { + id: string + key: string + color: string + rawNodes: GraphNode[] + documentCount: number + memoryCount: number + sampleDocumentIds: string[] + sampleMemoryIds: string[] +} + +export function buildClusterLevelGraph({ + nodes, + edges, + enabled, + expandedClusterId, + revealRawGraph, + canvasWidth, + canvasHeight, +}: { + nodes: GraphNode[] + edges: GraphEdge[] + enabled: boolean + expandedClusterId: string | null + revealRawGraph: boolean + canvasWidth: number + canvasHeight: number +}): ClusterLevelGraph { + if (!enabled || revealRawGraph || nodes.length === 0) { + return { + nodes, + edges, + isClustered: false, + clusterCount: 0, + rawNodeToClusterId: new Map(), + } + } + + const groups = buildClusterGroups(nodes) + if (groups.length === 0) { + return { + nodes, + edges, + isClustered: false, + clusterCount: 0, + rawNodeToClusterId: new Map(), + } + } + + const rawNodeToClusterId = new Map() + for (const group of groups) { + for (const node of group.rawNodes) { + rawNodeToClusterId.set(node.id, group.id) + } + } + + const visibleNodes: GraphNode[] = [] + const visibleIds = new Set() + const clusterNodeById = new Map() + + for (let index = 0; index < groups.length; index++) { + const group = groups[index] + if (!group) continue + + const clusterNode = createClusterNode( + group, + index, + groups.length, + canvasWidth, + canvasHeight, + ) + + if (expandedClusterId === group.id) { + for (const rawNode of cloneExpandedNodes(group.rawNodes, clusterNode)) { + visibleNodes.push(rawNode) + visibleIds.add(rawNode.id) + } + continue + } + + clusterNodeById.set(group.id, clusterNode) + visibleNodes.push(clusterNode) + visibleIds.add(clusterNode.id) + } + + const visibleEdges = buildVisibleEdges( + edges, + visibleIds, + rawNodeToClusterId, + clusterNodeById, + ) + + return { + nodes: visibleNodes, + edges: visibleEdges, + isClustered: true, + clusterCount: groups.length, + rawNodeToClusterId, + } +} + +function cloneExpandedNodes( + nodes: GraphNode[], + clusterNode: GraphNode, +): GraphNode[] { + if (nodes.length === 0) return [] + + let centerX = 0 + let centerY = 0 + for (const node of nodes) { + centerX += node.x + centerY += node.y + } + centerX /= nodes.length + centerY /= nodes.length + + let maxDistance = 1 + for (const node of nodes) { + const dx = node.x - centerX + const dy = node.y - centerY + maxDistance = Math.max(maxDistance, Math.sqrt(dx * dx + dy * dy)) + } + + const scale = Math.min(0.72, 620 / maxDistance) + + return nodes.map((node) => ({ + ...node, + x: clusterNode.x + (node.x - centerX) * scale, + y: clusterNode.y + (node.y - centerY) * scale, + })) +} + +function buildClusterGroups(nodes: GraphNode[]): ClusterGroup[] { + const memoriesByDocumentId = new Map() + const documentNodes: GraphNode[] = [] + const looseNodes: GraphNode[] = [] + + for (const node of nodes) { + if (node.type === "document") { + documentNodes.push(node) + continue + } + + if (node.type === "memory") { + const documentId = (node.data as MemoryNodeData).documentId + const bucket = memoriesByDocumentId.get(documentId) + if (bucket) { + bucket.push(node) + } else { + memoriesByDocumentId.set(documentId, [node]) + } + continue + } + + looseNodes.push(node) + } + + if (documentNodes.length === 0) return chunkLooseNodes(looseNodes) + + const groups: ClusterGroup[] = [] + let current: GraphNode[] = [] + let currentDocuments = 0 + let currentMemories = 0 + + const flush = () => { + if (current.length === 0) return + groups.push( + createGroup(groups.length, current, currentDocuments, currentMemories), + ) + current = [] + currentDocuments = 0 + currentMemories = 0 + } + + for (const doc of documentNodes) { + const memories = memoriesByDocumentId.get(doc.id) ?? [] + const docGroupSize = 1 + memories.length + if ( + current.length > 0 && + current.length + docGroupSize > CLUSTER_TARGET_RAW_NODES + ) { + flush() + } + + current.push(doc, ...memories) + currentDocuments++ + currentMemories += memories.length + + if (current.length >= CLUSTER_TARGET_RAW_NODES) flush() + } + flush() + + if (looseNodes.length > 0) { + for (const looseGroup of chunkLooseNodes(looseNodes, groups.length)) { + groups.push(looseGroup) + } + } + + return groups +} + +function chunkLooseNodes(nodes: GraphNode[], offset = 0): ClusterGroup[] { + const groups: ClusterGroup[] = [] + for (let i = 0; i < nodes.length; i += CLUSTER_TARGET_RAW_NODES) { + const chunk = nodes.slice(i, i + CLUSTER_TARGET_RAW_NODES) + groups.push( + createGroup( + offset + groups.length, + chunk, + chunk.filter((node) => node.type === "document").length, + chunk.filter((node) => node.type === "memory").length, + ), + ) + } + return groups +} + +function createGroup( + index: number, + rawNodes: GraphNode[], + documentCount: number, + memoryCount: number, +): ClusterGroup { + const firstNode = rawNodes[0] + const key = `supercluster:${index}:${firstNode?.clusterKey ?? firstNode?.id ?? "empty"}` + const sampleDocumentIds = rawNodes + .filter((node) => node.type === "document") + .slice(0, 5) + .map((node) => node.id) + const sampleMemoryIds = rawNodes + .filter((node) => node.type === "memory") + .slice(0, 5) + .map((node) => node.id) + const color = + firstNode?.clusterColor ?? firstNode?.borderColor ?? "rgba(88, 199, 232, 1)" + + return { + id: `cluster:${index}`, + key, + color, + rawNodes, + documentCount, + memoryCount, + sampleDocumentIds, + sampleMemoryIds, + } +} + +function createClusterNode( + group: ClusterGroup, + index: number, + total: number, + canvasWidth: number, + canvasHeight: number, +): GraphNode { + const centerX = canvasWidth / 2 + const centerY = canvasHeight / 2 + const spread = Math.max(420, Math.sqrt(total) * 210) + const angle = index * GOLDEN_ANGLE + const radius = spread * Math.sqrt((index + 1) / Math.max(1, total)) + const nodeCount = group.rawNodes.length + const size = clamp( + CLUSTER_MIN_SIZE + Math.sqrt(nodeCount) * 3.8, + CLUSTER_MIN_SIZE, + CLUSTER_MAX_SIZE, + ) + const data: ClusterNodeData = { + id: group.id, + title: `${group.documentCount} documents`, + summary: `${group.memoryCount} memories grouped for fast graph browsing`, + clusterKey: group.key, + documentCount: group.documentCount, + memoryCount: group.memoryCount, + nodeCount, + sampleDocumentIds: group.sampleDocumentIds, + sampleMemoryIds: group.sampleMemoryIds, + } + + return { + id: group.id, + type: "cluster", + x: centerX + Math.cos(angle) * radius, + y: centerY + Math.sin(angle) * radius, + data, + size, + borderColor: group.color, + clusterKey: group.key, + clusterColor: group.color, + isHovered: false, + isDragging: false, + } +} + +function buildVisibleEdges( + edges: GraphEdge[], + visibleIds: Set, + rawNodeToClusterId: Map, + clusterNodeById: Map, +): GraphEdge[] { + const result = new Map< + string, + { edge: GraphEdge; count: number; opacity: number; thickness: number } + >() + + for (const edge of edges) { + const sourceId = getEndpointId(edge.source) + const targetId = getEndpointId(edge.target) + const visibleSource = visibleIds.has(sourceId) + const visibleTarget = visibleIds.has(targetId) + const sourceClusterId = rawNodeToClusterId.get(sourceId) + const targetClusterId = rawNodeToClusterId.get(targetId) + const displaySource = visibleSource ? sourceId : sourceClusterId + const displayTarget = visibleTarget ? targetId : targetClusterId + + if (!displaySource || !displayTarget || displaySource === displayTarget) + continue + + const source = clusterNodeById.get(displaySource) ?? displaySource + const target = clusterNodeById.get(displayTarget) ?? displayTarget + const key = `${displaySource}->${displayTarget}:${edge.edgeType}` + const existing = result.get(key) + + if (existing) { + existing.count++ + existing.opacity = Math.max(existing.opacity, edge.visualProps.opacity) + existing.thickness = Math.max( + existing.thickness, + edge.visualProps.thickness, + ) + continue + } + + result.set(key, { + count: 1, + opacity: edge.visualProps.opacity, + thickness: edge.visualProps.thickness, + edge: { + id: `cluster-edge:${key}`, + source, + target, + edgeType: edge.edgeType, + visualProps: { ...edge.visualProps }, + }, + }) + } + + return [...result.values()].map(({ edge, count, opacity, thickness }) => ({ + ...edge, + visualProps: { + opacity: clamp(opacity + Math.log2(count) * 0.025, 0.12, 0.62), + thickness: clamp(thickness + Math.log2(count) * 0.2, 0.8, 4), + }, + })) +} + +function getEndpointId(endpoint: string | GraphNode): string { + return typeof endpoint === "string" ? endpoint : endpoint.id +} + +function clamp(value: number, min: number, max: number): number { + return Math.max(min, Math.min(max, value)) +} diff --git a/packages/memory-graph/src/types.ts b/packages/memory-graph/src/types.ts index b446a023..4f48f968 100644 --- a/packages/memory-graph/src/types.ts +++ b/packages/memory-graph/src/types.ts @@ -75,10 +75,10 @@ export interface MemoryNodeData { export interface GraphNode { id: string - type: "document" | "memory" + type: "document" | "memory" | "cluster" x: number y: number - data: DocumentNodeData | MemoryNodeData + data: DocumentNodeData | MemoryNodeData | ClusterNodeData size: number borderColor: string clusterKey?: string | null @@ -92,6 +92,18 @@ export interface GraphNode { fy?: number | null } +export interface ClusterNodeData { + id: string + title: string + summary: string + clusterKey: string + documentCount: number + memoryCount: number + nodeCount: number + sampleDocumentIds: string[] + sampleMemoryIds: string[] +} + export interface GraphEdge { id: string source: string | GraphNode