@cognizant-ai-lab/ui-common 1.3.3 → 1.4.0

package/dist/components/MultiAgentAccelerator/AgentFlow.js

@@ -18,45 +18,30 @@ import AdjustRoundedIcon from "@mui/icons-material/AdjustRounded";
 import ChatBubbleOutlineIcon from "@mui/icons-material/ChatBubbleOutline";
 import HubOutlinedIcon from "@mui/icons-material/HubOutlined";
 import ScatterPlotOutlinedIcon from "@mui/icons-material/ScatterPlotOutlined";
-import { ToggleButton, ToggleButtonGroup, useColorScheme, useTheme } from "@mui/material";
 import Box from "@mui/material/Box";
+import { useTheme } from "@mui/material/styles";
+import ToggleButton from "@mui/material/ToggleButton";
+import ToggleButtonGroup from "@mui/material/ToggleButtonGroup";
 import Tooltip from "@mui/material/Tooltip";
 import Typography from "@mui/material/Typography";
+import { applyNodeChanges, Background, ConnectionMode, ControlButton, Controls, ReactFlow, useReactFlow, useStore, } from "@xyflow/react";
 import { useCallback, useEffect, useMemo, useRef, useState } from "react";
-import { applyNodeChanges, Background, ConnectionMode, ControlButton, Controls, ReactFlow, useReactFlow, useStore, } from "reactflow";
 import { AgentNode, NODE_HEIGHT, NODE_WIDTH } from "./AgentNode.js";
-import { BACKGROUND_COLORS, BACKGROUND_COLORS_DARK_IDX, BASE_RADIUS, DEFAULT_FRONTMAN_X_POS, DEFAULT_FRONTMAN_Y_POS, HEATMAP_COLORS, LEVEL_SPACING, } from "./const.js";
-import { addThoughtBubbleEdge, layoutLinear, layoutRadial, removeThoughtBubbleEdge } from "./GraphLayouts.js";
+import { BASE_RADIUS, DEFAULT_FRONTMAN_X_POS, DEFAULT_FRONTMAN_Y_POS, LEVEL_SPACING } from "./const.js";
+import { addThoughtBubbleEdge, layoutLinear, layoutRadial } from "./GraphLayouts.js";
 import { PlasmaEdge } from "./PlasmaEdge.js";
 import { ThoughtBubbleEdge } from "./ThoughtBubbleEdge.js";
 import { ThoughtBubbleOverlay } from "./ThoughtBubbleOverlay.js";
-import { isDarkMode } from "../../utils/Theme.js";
+import { usePalette } from "../../Theme/Palettes.js";
 import { getZIndex } from "../../utils/zIndexLayers.js";
 // #endregion: Types
 // #region: Constants
 // Timeout for thought bubbles is set to 10 seconds
 const THOUGHT_BUBBLE_TIMEOUT_MS = 10_000;
-// Maximum number of thought bubbles allowed on screen at any time. As more arrive, older ones are removed.
-const MAX_THOUGHT_BUBBLES = 5;
 // #endregion: Constants
-export const AgentFlow = ({ agentCounts, agentsInNetwork, currentConversations, id, isAwaitingLlm, isStreaming, thoughtBubbleEdges, setThoughtBubbleEdges, }) => {
+export const AgentFlow = ({ agentCounts, agentIconSuggestions, agentsInNetwork, currentConversations, id, isAwaitingLlm, isStreaming, thoughtBubbleEdges, setThoughtBubbleEdges, }) => {
     const theme = useTheme();
     const { fitView } = useReactFlow();
-    // Helper functions to update thought bubble edges state immutably
-    const addThoughtBubbleEdgeHelper = useCallback((conversationId, edge) => {
-        setThoughtBubbleEdges((prev) => {
-            const newMap = new Map(prev);
-            addThoughtBubbleEdge(newMap, conversationId, edge);
-            return newMap;
-        });
-    }, [setThoughtBubbleEdges]);
-    const removeThoughtBubbleEdgeHelper = useCallback((conversationId) => {
-        setThoughtBubbleEdges((prev) => {
-            const newMap = new Map(prev);
-            removeThoughtBubbleEdge(newMap, conversationId);
-            return newMap;
-        });
-    }, [setThoughtBubbleEdges]);
     const handleResize = useCallback(() => {
         fitView(); // Adjusts the view to fit after resizing
     }, [fitView, isAwaitingLlm]);
@@ -64,18 +49,10 @@ export const AgentFlow = ({ agentCounts, agentsInNetwork, currentConversations,
         window.addEventListener("resize", handleResize);
         return () => window.removeEventListener("resize", handleResize);
     }, [handleResize]);
-    // Save this as a mutable ref so child nodes see updates
-    const conversationsRef = useRef(currentConversations);
-    useEffect(() => {
-        conversationsRef.current = currentConversations;
-    }, [currentConversations]);
     const [layout, setLayout] = useState("radial");
     const [coloringOption, setColoringOption] = useState("depth");
     const [enableRadialGuides, setEnableRadialGuides] = useState(true);
     const [showThoughtBubbles, setShowThoughtBubbles] = useState(true);
-    // State for managing active thought bubbles with timers. Use ActiveThoughtBubble which
-    // references the original conversation via `conversationId`.
-    const [activeThoughtBubbles, setActiveThoughtBubbles] = useState([]);
     // Track conversation IDs we've already processed to prevent re-adding after expiry
     const processedConversationIdsRef = useRef(new Set());
     // Track which bubble is currently being hovered
@@ -83,159 +60,96 @@ export const AgentFlow = ({ agentCounts, agentsInNetwork, currentConversations,
     const handleBubbleHoverChange = useCallback((bubbleId) => {
         hoveredBubbleIdRef.current = bubbleId;
     }, []);
+    // Ref for isStreaming, read inside the cleanup interval.
+    const isStreamingRef = useRef(isStreaming);
+    // Keep the ref current after every render.
+    useEffect(() => {
+        isStreamingRef.current = isStreaming;
+    });
     // Clear processed conversation IDs when thought bubble edges are cleared (streaming ends)
     useEffect(() => {
         if (thoughtBubbleEdges.size === 0) {
             processedConversationIdsRef.current.clear();
         }
     }, [thoughtBubbleEdges.size]);
-    // Effect to add and remove thought bubbles - only when we actually have conversations to process
+    // Add new thought bubble edges for incoming conversations.
     useEffect(() => {
-        if (!currentConversations || currentConversations.length === 0) {
-            return; // Skip processing when no conversations
-        }
-        setActiveThoughtBubbles((prevBubbles) => {
-            const newBubbles = [];
-            // Check the text that user sees as one level of duplicate prevention
+        if (!currentConversations || currentConversations.length === 0)
+            return;
+        setThoughtBubbleEdges((prev) => {
             const processedText = new Set();
-            thoughtBubbleEdges.forEach((thoughtBubbleEdge) => {
-                const edgeText = thoughtBubbleEdge.edge.data?.text?.trim();
-                if (edgeText) {
-                    // Add edge text to prevent duplicates
-                    processedText.add(edgeText);
-                }
-            });
-            // Only add bubbles for conversations with unique parsed content
-            // TODO: Even though there are two duplicate checks here, some duplicates still get through
+            for (const entry of prev.values()) {
+                const text = entry.edge.data?.text?.trim();
+                if (text)
+                    processedText.add(text);
+            }
+            let edgesMap = null;
             for (const conv of currentConversations) {
                 const convText = conv.text?.trim();
-                if (
-                // Has text
-                convText &&
-                    // Haven't already displayed this text (duplicate check #1)
+                const agentList = Array.from(conv.agents);
+                if (convText &&
+                    agentList.length >= 2 &&
                     !processedText.has(convText) &&
-                    // Haven't already displayed this conversation ID (duplicate check #2)
                     !processedConversationIdsRef.current.has(conv.id)) {
-                    // Create an AgentConversation object representing the active thought bubble.
-                    // Use the conversation id from the incoming conversation and set startedAt
-                    // to "now" so the bubble timeout is measured from when the bubble appeared.
-                    newBubbles.push({
-                        agents: new Set(conv.agents),
-                        conversationId: conv.id,
-                        startedAt: new Date(),
-                        type: conv.type,
-                        // No text needed. This is for tracking which conversations are active.
-                    });
-                    // Mark this conversation ID as processed. This protects against duplicates if conversations have
-                    // the same ID, although some duplicates still get through.
+                    if (!edgesMap)
+                        edgesMap = new Map(prev);
                     processedConversationIdsRef.current.add(conv.id);
-                    const agentList = Array.from(conv.agents);
-                    // TODO: Check if agentList has at least 2 agents. This allows us to set a source and target.
-                    // However, this can also have more than 2 agents. We may want to think more about that case.
-                    if (agentList.length >= 2) {
-                        // These source and target agents are going to be useful later when we point bubbles to nodes.
-                        const sourceAgent = agentList[0];
-                        const targetAgent = agentList[1];
-                        const edge = {
-                            id: `thought-bubble-${conv.id}`,
-                            source: sourceAgent,
-                            target: targetAgent,
-                            type: "thoughtBubbleEdge",
-                            data: {
-                                text: conv.text,
-                                showAlways: showThoughtBubbles,
-                                conversationId: conv.id,
-                                // Include full agent list so overlays can point to all participants
-                                agents: agentList,
-                                type: conv.type, // Add conversation type for filtering
-                            },
-                            style: { pointerEvents: "none" },
-                        };
-                        addThoughtBubbleEdgeHelper(conv.id, edge);
-                    }
+                    processedText.add(convText);
+                    const edge = {
+                        id: `thought-bubble-${conv.id}`,
+                        source: agentList[0],
+                        target: agentList[1],
+                        type: "thoughtBubbleEdge",
+                        data: {
+                            text: conv.text,
+                            showAlways: showThoughtBubbles,
+                            conversationId: conv.id,
+                            agents: agentList,
+                            type: conv.type,
+                        },
+                        style: { pointerEvents: "none" },
+                    };
+                    addThoughtBubbleEdge(edgesMap, conv.id, edge); // also enforces MAX_GLOBAL_THOUGHT_BUBBLES
                 }
             }
-            if (newBubbles.length === 0) {
-                return prevBubbles; // No changes
-            }
-            const allBubbles = [...prevBubbles, ...newBubbles];
-            // If we're over the limit, first remove expired bubbles, then oldest if still needed
-            if (allBubbles.length > MAX_THOUGHT_BUBBLES) {
-                const now = Date.now();
-                const dropped = [];
-                // First, filter out expired bubbles by checking THOUGHT_BUBBLE_TIMEOUT_MS
-                const activeBubbles = allBubbles.filter((bubble) => {
-                    const age = now - bubble.startedAt.getTime();
-                    const isExpired = age >= THOUGHT_BUBBLE_TIMEOUT_MS;
-                    if (isExpired) {
-                        dropped.push(bubble);
-                    }
-                    return !isExpired;
-                });
-                // If still over the limit after removing expired, remove oldest non-expired
-                let finalBubbles = activeBubbles;
-                if (activeBubbles.length > MAX_THOUGHT_BUBBLES) {
-                    const sorted = activeBubbles.sort((a, b) => a.startedAt.getTime() - b.startedAt.getTime());
-                    finalBubbles = sorted.slice(-MAX_THOUGHT_BUBBLES);
-                    const additionalDropped = sorted.slice(0, -MAX_THOUGHT_BUBBLES);
-                    dropped.push(...additionalDropped);
-                }
-                // Clean up all dropped bubbles from edge cache
-                dropped.forEach((bubble) => {
-                    removeThoughtBubbleEdgeHelper(bubble.conversationId);
-                });
-                return finalBubbles;
-            }
-            return allBubbles;
+            return edgesMap ?? prev;
         });
-    }, [currentConversations, addThoughtBubbleEdgeHelper, removeThoughtBubbleEdgeHelper]);
-    // Independent thought bubble cleanup - run timer only during streaming
+    }, [currentConversations, showThoughtBubbles, setThoughtBubbleEdges]);
+    // Cleanup expired thought bubble edges — created once on mount, reads isStreaming via ref.
     useEffect(() => {
         const cleanupInterval = setInterval(() => {
-            // Only clean up if we're currently streaming
-            if (!isStreaming)
+            if (!isStreamingRef.current)
                 return;
-            setActiveThoughtBubbles((prevBubbles) => {
-                if (prevBubbles.length === 0)
-                    return prevBubbles;
-                const now = Date.now();
-                const filteredBubbles = prevBubbles.filter((bubble) => {
-                    const age = now - bubble.startedAt.getTime();
-                    const shouldKeep = age < THOUGHT_BUBBLE_TIMEOUT_MS;
-                    // Keep bubble if it's being hovered, even if expired
-                    const isHovered = hoveredBubbleIdRef.current === `thought-bubble-${bubble.conversationId}`;
-                    if (isHovered) {
-                        return true;
-                    }
-                    // If bubble is expiring, remove it from the edge cache
-                    if (!shouldKeep) {
-                        removeThoughtBubbleEdgeHelper(bubble.conversationId);
+            const now = Date.now();
+            setThoughtBubbleEdges((prev) => {
+                let changed = false;
+                const edgesMap = new Map(prev);
+                for (const [convId, entry] of prev) {
+                    const isHovered = hoveredBubbleIdRef.current === `thought-bubble-${convId}`;
+                    if (!isHovered && now - entry.timestamp >= THOUGHT_BUBBLE_TIMEOUT_MS) {
+                        edgesMap.delete(convId);
+                        changed = true;
                     }
-                    return shouldKeep;
-                });
-                // Only update if there are changes to prevent unnecessary re-renders
-                if (filteredBubbles.length !== prevBubbles.length) {
-                    return filteredBubbles;
                 }
-                return prevBubbles;
+                return changed ? edgesMap : prev;
             });
-        }, 1000); // Check every second
+        }, 1000);
         return () => clearInterval(cleanupInterval);
-    }, [isStreaming, removeThoughtBubbleEdgeHelper]);
-    const { mode, systemMode } = useColorScheme();
-    const darkMode = isDarkMode(mode, systemMode);
-    // Shadow color for icon. TODO: use MUI theme system instead.
-    const shadowColor = darkMode ? "255, 255, 255" : "0, 0, 0";
+    }, []); // mount/unmount only
+    // Shadow color for icon
+    const shadowColor = theme.palette.mode === "dark" ? theme.palette.common.white : theme.palette.common.black;
     const isHeatmap = coloringOption === "heatmap";
+    const palette = usePalette();
     // Merge agents from active thought bubbles with agentsInNetwork for layout
     // This ensures bubble edges persist even when agents disappear from the network
     const bubbleAgentIds = useMemo(() => {
         const ids = new Set();
-        activeThoughtBubbles.forEach((bubble) => {
-            bubble.agents.forEach((agentId) => ids.add(agentId));
+        thoughtBubbleEdges.forEach(({ edge }) => {
+            const agents = edge.data?.agents ?? [];
+            agents.forEach((agentId) => ids.add(agentId));
         });
         return ids;
-    }, [activeThoughtBubbles]);
+    }, [thoughtBubbleEdges]);
     const mergedAgentsInNetwork = useMemo(() => {
         // Add any missing agents from bubbles as minimal ConnectivityInfo
         const existingIds = new Set(agentsInNetwork.map((a) => a.origin));
@@ -250,17 +164,17 @@ export const AgentFlow = ({ agentCounts, agentsInNetwork, currentConversations,
     // Create the flow layout depending on user preference
     // Memoize layoutResult so it only recalculates when relevant data changes
     const layoutResult = useMemo(() => layout === "linear"
-        ? layoutLinear(isHeatmap ? agentCounts : undefined, mergedAgentsInNetwork, currentConversations, isAwaitingLlm, thoughtBubbleEdges)
-        : layoutRadial(isHeatmap ? agentCounts : undefined, mergedAgentsInNetwork, currentConversations, isAwaitingLlm, thoughtBubbleEdges), [
-        layout,
-        coloringOption,
+        ? layoutLinear(isHeatmap ? agentCounts : undefined, mergedAgentsInNetwork, currentConversations, isAwaitingLlm, thoughtBubbleEdges, agentIconSuggestions)
+        : layoutRadial(isHeatmap ? agentCounts : undefined, mergedAgentsInNetwork, currentConversations, isAwaitingLlm, thoughtBubbleEdges, agentIconSuggestions), [
         agentCounts,
-        mergedAgentsInNetwork,
+        agentIconSuggestions,
+        coloringOption,
         currentConversations,
-        activeThoughtBubbles,
-        thoughtBubbleEdges,
         isAwaitingLlm,
+        layout,
+        mergedAgentsInNetwork,
         showThoughtBubbles,
+        thoughtBubbleEdges,
     ]);
     const [nodes, setNodes] = useState(layoutResult.nodes);
     // Sync up the nodes with the layout result
@@ -268,6 +182,8 @@ export const AgentFlow = ({ agentCounts, agentsInNetwork, currentConversations,
         setNodes(layoutResult.nodes);
     }, [layoutResult.nodes]);
     const edges = layoutResult.edges;
+    // Make sure to extract only thought bubble edges for the overlay.
+    const thoughtBubbleEdgesForOverlay = useMemo(() => edges.filter((e) => e.type === "thoughtBubbleEdge"), [edges]);
     useEffect(() => {
         // Schedule a fitView after the layout is set to ensure the view is adjusted correctly
         setTimeout(() => {
@@ -306,28 +222,23 @@ export const AgentFlow = ({ agentCounts, agentsInNetwork, currentConversations,
             }, children: _jsx("g", { id: `${id}-radial-guides-group`, transform: `translate(${transform[0]}, ${transform[1]}) scale(${transform[2]})`, children: circles }) }));
     };
     // Generate Legend for depth or heatmap colors
-    function getLegend() {
-        const palette = isHeatmap ? HEATMAP_COLORS : BACKGROUND_COLORS;
-        const title = isHeatmap ? "Heat" : "Depth";
-        const length = isHeatmap ? HEATMAP_COLORS.length : Math.min(maxDepth, BACKGROUND_COLORS.length);
+    const getLegend = () => {
+        const length = isHeatmap ? palette.length : Math.min(maxDepth, palette.length);
         return (_jsxs(Box, { id: `${id}-legend`, sx: {
                 position: "absolute",
                 top: "5px",
                 right: "10px",
                 padding: "5px",
                 borderRadius: "5px",
-                boxShadow: `0 0 5px rgba(${shadowColor}, 0.3)`,
+                boxShadow: `0 0 5px color-mix(in srgb, ${shadowColor} 30%, transparent)`,
                 display: "flex",
                 alignItems: "center",
                 zIndex: getZIndex(2, theme),
-            }, children: [_jsx(Typography, { id: `${id}-legend-label`, sx: {
-                        fontSize: "10px",
-                        marginLeft: "0.2rem",
-                    }, children: title }), Array.from({ length }, (_, i) => (_jsx(Box, { id: `${id}-legend-depth-${i}`, style: {
+            }, children: [Array.from({ length }, (_, i) => (_jsx(Box, { id: `${id}-legend-depth-${i}`, style: {
                         alignItems: "center",
                         backgroundColor: palette[i],
                         borderRadius: "50%",
-                        color: i < BACKGROUND_COLORS_DARK_IDX ? "var(--bs-primary)" : "var(--bs-white)",
+                        color: theme.palette.getContrastText(palette[i]),
                         display: "flex",
                         height: "15px",
                         justifyContent: "center",
@@ -346,26 +257,27 @@ export const AgentFlow = ({ agentCounts, agentsInNetwork, currentConversations,
                     }, size: "small", children: [_jsx(ToggleButton, { id: `${id}-depth-toggle`, size: "small", value: "depth", sx: {
                                 fontSize: "0.5rem",
                                 height: "1rem",
-                                backgroundColor: darkMode && coloringOption === "depth" ? "var(--bs-gray-medium)" : undefined,
                             }, children: _jsx(Typography, { id: `${id}-depth-label`, sx: {
                                     fontSize: "10px",
                                 }, children: "Depth" }) }), _jsx(ToggleButton, { id: `${id}-heatmap-toggle`, size: "small", value: "heatmap", sx: {
                                 fontSize: "0.5rem",
                                 height: "1rem",
-                                backgroundColor: darkMode && isHeatmap ? "var(--bs-gray-medium)" : undefined,
                             }, children: _jsx(Typography, { id: `${id}-heatmap-label`, sx: {
                                     fontSize: "10px",
                                 }, children: "Heatmap" }) })] })] }));
-    }
+    };
     // Get the background color for the control buttons based on the layout and dark mode setting
     const getControlButtonBackgroundColor = (isActive) => {
-        return isActive ? (darkMode ? "var(--bs-gray-dark)" : "var(--bs-gray-lighter)") : undefined;
+        if (!isActive) {
+            return undefined;
+        }
+        return theme.palette.mode === "dark" ? theme.palette.grey[800] : theme.palette.grey[200];
     };
     // Only show radial guides if radial layout is selected, radial guides are enabled, and it's not just Frontman
     const shouldShowRadialGuides = enableRadialGuides && layout === "radial" && maxDepth > 1;
     // Generate the control bar for the flow, including layout and radial guides toggles
     const getControls = () => {
-        return (_jsxs(Controls, { id: "react-flow-controls", position: "top-left", style: {
+        return (_jsxs(Controls, { position: "top-left", style: {
                 position: "absolute",
                 top: "0px",
                 left: "0px",
@@ -373,22 +285,31 @@ export const AgentFlow = ({ agentCounts, agentsInNetwork, currentConversations,
                 width: "auto",
             }, showInteractive: true, children: [_jsx(Tooltip, { id: "radial-layout-tooltip", title: "Radial layout", placement: "right", children: _jsx("span", { id: "radial-layout-span", children: _jsx(ControlButton, { id: "radial-layout-button", onClick: () => setLayout("radial"), style: {
                                 backgroundColor: getControlButtonBackgroundColor(layout === "radial"),
-                            }, children: _jsx(HubOutlinedIcon, { id: "radial-layout-icon", sx: { color: darkMode ? "var(--bs-white)" : "var(--bs-dark-mode-dim)" } }) }) }) }), _jsx(Tooltip, { id: "linear-layout-tooltip", title: "Linear layout", placement: "right", children: _jsx("span", { id: "linear-layout-span", children: _jsx(ControlButton, { id: "linear-layout-button", onClick: () => setLayout("linear"), style: {
+                            }, children: _jsx(HubOutlinedIcon, { id: "radial-layout-icon" }) }) }) }), _jsx(Tooltip, { id: "linear-layout-tooltip", title: "Linear layout", placement: "right", children: _jsx("span", { id: "linear-layout-span", children: _jsx(ControlButton, { id: "linear-layout-button", onClick: () => setLayout("linear"), style: {
                                 backgroundColor: getControlButtonBackgroundColor(layout === "linear"),
-                            }, children: _jsx(ScatterPlotOutlinedIcon, { id: "linear-layout-icon", sx: { color: darkMode ? "var(--bs-white)" : "var(--bs-dark-mode-dim)" } }) }) }) }), _jsx(Tooltip, { id: "radial-guides-tooltip", title: `Enable/disable radial guides${layout === "radial" ? "" : " (only available in radial layout)"}`, placement: "right", children: _jsx("span", { id: "radial-guides-span", children: _jsx(ControlButton, { id: "radial-guides-button", onClick: () => setEnableRadialGuides(!enableRadialGuides), style: {
+                            }, children: _jsx(ScatterPlotOutlinedIcon, { id: "linear-layout-icon" }) }) }) }), _jsx(Tooltip, { id: "radial-guides-tooltip", title: `Enable/disable radial guides${layout === "radial" ? "" : " (only available in radial layout)"}`, placement: "right", children: _jsx("span", { id: "radial-guides-span", children: _jsx(ControlButton, { id: "radial-guides-button", onClick: () => setEnableRadialGuides(!enableRadialGuides), style: {
                                 backgroundColor: getControlButtonBackgroundColor(enableRadialGuides),
-                            }, disabled: layout !== "radial", children: _jsx(AdjustRoundedIcon, { id: "radial-guides-icon", sx: { color: darkMode ? "var(--bs-white)" : "var(--bs-dark-mode-dim)" } }) }) }) }), _jsx(Tooltip, { id: "thought-bubble-tooltip", title: `Toggle thought bubbles ${showThoughtBubbles ? "off" : "on"}`, placement: "right", children: _jsx("span", { id: "thought-bubble-span", children: _jsx(ControlButton, { id: "thought-bubble-button", onClick: () => setShowThoughtBubbles(!showThoughtBubbles), style: {
+                            }, disabled: layout !== "radial", children: _jsx(AdjustRoundedIcon, { id: "radial-guides-icon" }) }) }) }), _jsx(Tooltip, { id: "thought-bubble-tooltip", title: `Toggle thought bubbles ${showThoughtBubbles ? "off" : "on"}`, placement: "right", children: _jsx("span", { id: "thought-bubble-span", children: _jsx(ControlButton, { id: "thought-bubble-button", onClick: () => setShowThoughtBubbles(!showThoughtBubbles), style: {
                                 backgroundColor: getControlButtonBackgroundColor(showThoughtBubbles),
-                            }, children: _jsx(ChatBubbleOutlineIcon, { id: "thought-bubble-icon", sx: { color: darkMode ? "var(--bs-white)" : "var(--bs-dark-mode-dim)" } }) }) }) })] }));
+                            }, children: _jsx(ChatBubbleOutlineIcon, { id: "thought-bubble-icon" }) }) }) })] }));
     };
     return (_jsxs(Box, { id: `${id}-outer-box`, sx: {
             height: "100%",
             width: "100%",
+            backgroundColor: theme.palette.background.default,
             "& .react-flow__node": {
-                border: "var(--bs-border-width) var(--bs-border-style) var(--bs-black)",
-                borderRadius: "var(--bs-border-radius-2xl)",
+                border: `1px solid ${theme.palette.divider}`,
+            },
+            "& .react-flow__panel": {
+                backgroundColor: theme.palette.background.paper,
+                border: `1px solid ${theme.palette.divider}`,
+                color: theme.palette.text.primary,
+            },
+            "& .react-flow__controls-button": {
+                backgroundColor: theme.palette.background.paper,
+                borderBottom: `1px solid ${theme.palette.divider}`,
+                color: theme.palette.text.primary,
+                fill: theme.palette.text.primary,
             },
-        }, 
-        // Theme the "React Flow" attribution logo according to dark mode.
-        className: darkMode ? "dark" : undefined, children: [_jsx(ReactFlow, { id: `${id}-react-flow`, nodes: nodes, edges: edges, onNodesChange: onNodesChange, fitView: true, nodeTypes: nodeTypes, edgeTypes: edgeTypes, connectionMode: ConnectionMode.Loose, children: !isAwaitingLlm && (_jsxs(_Fragment, { children: [agentsInNetwork?.length ? getLegend() : null, _jsx(Background, { id: `${id}-background` }), getControls(), shouldShowRadialGuides ? getRadialGuides() : null] })) }), _jsx(ThoughtBubbleOverlay, { nodes: nodes, edges: edges, showThoughtBubbles: showThoughtBubbles, isStreaming: isStreaming, onBubbleHoverChange: handleBubbleHoverChange })] }));
+        }, children: [_jsx(ReactFlow, { id: `${id}-react-flow`, nodes: nodes, edges: edges, onNodesChange: onNodesChange, fitView: true, nodeTypes: nodeTypes, edgeTypes: edgeTypes, connectionMode: ConnectionMode.Loose, children: !isAwaitingLlm && (_jsxs(_Fragment, { children: [agentsInNetwork?.length ? getLegend() : null, _jsx(Background, { id: `${id}-background` }), getControls(), shouldShowRadialGuides ? getRadialGuides() : null] })) }), _jsx(ThoughtBubbleOverlay, { nodes: nodes, edges: thoughtBubbleEdgesForOverlay, showThoughtBubbles: showThoughtBubbles, isStreaming: isStreaming, onBubbleHoverChange: handleBubbleHoverChange })] }));
 };

package/dist/components/MultiAgentAccelerator/AgentNode.d.ts

@@ -1,13 +1,15 @@
+import { NodeProps } from "@xyflow/react";
+import type { Node as RFNode } from "@xyflow/react";
 import { FC } from "react";
-import { NodeProps } from "reactflow";
-import { AgentConversation } from "../../utils/agentConversations.js";
-export interface AgentNodeProps {
+import { AgentConversation } from "./AgentConversations.js";
+export interface AgentNodeProps extends Record<string, unknown> {
     readonly agentCounts?: Map<string, number>;
     readonly agentName: string;
     readonly depth: number;
+    readonly displayAs?: string;
     readonly getConversations: () => AgentConversation[] | null;
     readonly isAwaitingLlm?: boolean;
-    readonly displayAs?: string;
+    readonly agentIconSuggestion?: string;
 }
 export declare const NODE_HEIGHT = 100;
 export declare const NODE_WIDTH = 100;
@@ -15,4 +17,4 @@ export declare const NODE_WIDTH = 100;
  * A node representing an agent in the network for use in react-flow.
  * @param props See AgentNodeProps
  */
-export declare const AgentNode: FC<NodeProps<AgentNodeProps>>;
+export declare const AgentNode: FC<NodeProps<RFNode<AgentNodeProps>>>;