v1 lats visualization done

TataKKKL · TataKKKL · commit 4b353a9a4d3b · 2025-04-15T18:53:55.000+08:00
diff --git a/visual-tree-search-app/components/LATSVisual.tsx b/visual-tree-search-app/components/LATSVisual.tsx
@@ -13,6 +13,7 @@ interface TreeNode {
   visits?: number;
   feedback?: string;
   reward?: number;
+  isSimulated?: boolean; // Flag to track newly simulated nodes
 }
 
 interface Message {
@@ -31,8 +32,10 @@ const LATSVisual: React.FC<SimpleSearchVisualProps> = ({ messages }) => {
   const tooltipRef = useRef<HTMLDivElement | null>(null);
   const { theme } = useTheme();
   const [selectedNodeId, setSelectedNodeId] = useState<number | null>(null);
+  const [simulationStartNodeId, setSimulationStartNodeId] = useState<number | null>(null); // Track simulation starting node (existing node)
   const [treeNodes, setTreeNodes] = useState<TreeNode[]>([]);
   const [containerWidth, setContainerWidth] = useState<number>(0);
+  const [simulatedNodes, setSimulatedNodes] = useState<number[]>([]); // Keep track of new simulated node IDs
 
   // Set up resize observer to make the visualization responsive
   useEffect(() => {
@@ -71,27 +74,91 @@ const LATSVisual: React.FC<SimpleSearchVisualProps> = ({ messages }) => {
 
     let updatedTreeNodes: TreeNode[] = [...treeNodes];
     let newSelectedNodeId = selectedNodeId;
+    let newSimulationStartNodeId = simulationStartNodeId;
+    let newSimulatedNodes = [...simulatedNodes];
     let hasChanges = false;
 
     messages.forEach(msg => {
       try {
         const data = JSON.parse(msg.content);
         
-        // Handle node selection updates
+        // Handle regular node selection (during tree expansion/evaluation)
         if (data.type === 'node_selected' && data.node_id !== undefined) {
           newSelectedNodeId = data.node_id;
           hasChanges = true;
         }
         
+        // Handle simulation start node selection (existing node highlighted as simulation start)
+        if (data.type === 'node_selected_for_simulation' && data.node_id !== undefined) {
+          newSimulationStartNodeId = data.node_id;
+          hasChanges = true;
+        }
+        
         // Handle tree structure updates
-        if (data.type === 'tree_update_node_expansion' && Array.isArray(data.tree)) {
-          updatedTreeNodes = data.tree;
+        if ((data.type === 'tree_update_node_expansion' || data.type === 'tree_update_node_children_evaluation') 
+            && Array.isArray(data.tree)) {
+          // Preserve simulation flags when updating from tree
+          if (updatedTreeNodes.some(node => node.isSimulated)) {
+            // Find all nodes with isSimulated flag
+            const simulatedNodeMap = new Map();
+            updatedTreeNodes.forEach(node => {
+              if (node.isSimulated) {
+                simulatedNodeMap.set(node.id, true);
+              }
+            });
+            
+            // Apply the flag to the updated tree
+            updatedTreeNodes = data.tree.map((node: TreeNode) => ({
+              ...node,
+              isSimulated: simulatedNodeMap.has(node.id) ? true : false
+            }));
+          } else {
+            updatedTreeNodes = data.tree;
+          }
           hasChanges = true;
         }
+
+        // Handle simulated node creation
+        if (data.type === 'node_simulated' && data.node_id !== undefined && data.parent_id !== undefined) {
+          // Check if the node already exists in the tree
+          const nodeExists = updatedTreeNodes.some(node => node.id === data.node_id);
+          
+          if (!nodeExists) {
+            // Add the new simulated node to the tree
+            updatedTreeNodes.push({
+              id: data.node_id,
+              parent_id: data.parent_id,
+              action: data.action,
+              description: data.description,
+              isSimulated: true, // Mark as simulated
+            });
+            
+            // Add to our list of simulated nodes
+            newSimulatedNodes.push(data.node_id);
+            hasChanges = true;
+          } else {
+            // If node already exists, update it to mark as simulated
+            updatedTreeNodes = updatedTreeNodes.map(node => 
+              node.id === data.node_id ? { ...node, isSimulated: true } : node
+            );
+            
+            if (!newSimulatedNodes.includes(data.node_id)) {
+              newSimulatedNodes.push(data.node_id);
+              hasChanges = true;
+            }
+          }
+        }
         
-        // Handle node evaluation updates
-        if (data.type === 'tree_update_node_evaluation' && Array.isArray(data.tree)) {
-          updatedTreeNodes = data.tree;
+        // Handle simulation removal
+        if (data.type === 'removed_simulation') {
+          // Remove simulation flags instead of removing nodes
+          updatedTreeNodes = updatedTreeNodes.map(node => ({
+            ...node, 
+            isSimulated: false // Remove simulation flag
+          }));
+          
+          newSimulatedNodes = []; // Clear simulated nodes list
+          newSimulationStartNodeId = null; // Clear simulation start node
           hasChanges = true;
         }
       } catch {
@@ -102,8 +169,10 @@ const LATSVisual: React.FC<SimpleSearchVisualProps> = ({ messages }) => {
     if (hasChanges) {
       setTreeNodes(updatedTreeNodes);
       setSelectedNodeId(newSelectedNodeId);
+      setSimulationStartNodeId(newSimulationStartNodeId);
+      setSimulatedNodes(newSimulatedNodes);
     }
-  }, [messages, treeNodes, selectedNodeId]);
+  }, [messages]);
 
   // Render the tree visualization
   useEffect(() => {
@@ -136,6 +205,7 @@ const LATSVisual: React.FC<SimpleSearchVisualProps> = ({ messages }) => {
         .style("z-index", "1000")
         .style("max-width", "400px")
         .style("box-shadow", "0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06)")
+        .style("line-height", "1.5")
         .node() as HTMLDivElement;
     };
 
@@ -187,9 +257,41 @@ const LATSVisual: React.FC<SimpleSearchVisualProps> = ({ messages }) => {
         return `M${sourceY},${sourceX}C${(sourceY + targetY) / 2},${sourceX} ${(sourceY + targetY) / 2},${targetX} ${targetY},${targetX}`;
       })
       .attr("fill", "none")
-      .attr("stroke", theme === 'dark' ? "#9CA3AF" : "#6B7280")
-      .attr("stroke-width", 1.5)
-      .attr("stroke-opacity", 0.7);
+      .attr("stroke", d => {
+        // Link to a simulated node gets an orange color
+        if (d.target.data.isSimulated) {
+          return theme === 'dark' ? "#F97316" : "#FB923C"; // Orange for simulated paths
+        }
+        
+        // Link from simulation start node
+        if (d.source.data.id === simulationStartNodeId) {
+          return theme === 'dark' ? "#10B981" : "#34D399"; // Green for simulation start path
+        }
+        
+        // Default link color
+        return theme === 'dark' ? "#9CA3AF" : "#6B7280";
+      })
+      .attr("stroke-width", d => {
+        // Thicker link for simulation paths
+        if (d.target.data.isSimulated || d.source.data.id === simulationStartNodeId) {
+          return 2.5;
+        }
+        return 1.5;
+      })
+      .attr("stroke-opacity", d => {
+        // More visible for simulation paths
+        if (d.target.data.isSimulated || d.source.data.id === simulationStartNodeId) {
+          return 0.9;
+        }
+        return 0.7;
+      })
+      .attr("stroke-dasharray", d => {
+        // Dashed line for simulation paths
+        if (d.target.data.isSimulated) {
+          return "5,3";
+        }
+        return null;
+      });
 
     // Create node groups
     const nodes = g.selectAll(".node")
@@ -203,6 +305,16 @@ const LATSVisual: React.FC<SimpleSearchVisualProps> = ({ messages }) => {
     nodes.append("circle")
       .attr("r", 12)
       .attr("fill", d => {
+        // Simulated node (orange)
+        if (d.data.isSimulated) {
+          return theme === 'dark' ? "#F97316" : "#FDBA74"; // Orange for simulated nodes
+        }
+        
+        // Simulation start node (green)
+        if (d.data.id === simulationStartNodeId) {
+          return theme === 'dark' ? "#10B981" : "#34D399"; // Green for simulation start node
+        }
+        
         // Selected node (blue)
         if (d.data.id === selectedNodeId) {
           return theme === 'dark' ? "#3B82F6" : "#60A5FA";
@@ -216,12 +328,29 @@ const LATSVisual: React.FC<SimpleSearchVisualProps> = ({ messages }) => {
         // Action node (default)
         return theme === 'dark' ? "#4B5563" : "#E5E7EB";
       })
-      .attr("stroke", d => d.data.id === selectedNodeId
-        ? theme === 'dark' ? "#93C5FD" : "#2563EB"
-        : theme === 'dark' ? "#374151" : "#D1D5DB")
-      .attr("stroke-width", d => d.data.id === selectedNodeId ? 3 : 2);
+      .attr("stroke", d => {
+        if (d.data.isSimulated) {
+          return theme === 'dark' ? "#EA580C" : "#F97316"; // Darker orange stroke for simulated nodes
+        }
+        
+        if (d.data.id === simulationStartNodeId) {
+          return theme === 'dark' ? "#059669" : "#10B981"; // Darker green stroke for simulation start
+        }
+        
+        if (d.data.id === selectedNodeId) {
+          return theme === 'dark' ? "#93C5FD" : "#2563EB";
+        }
+        
+        return theme === 'dark' ? "#374151" : "#D1D5DB";
+      })
+      .attr("stroke-width", d => {
+        if (d.data.isSimulated || d.data.id === simulationStartNodeId || d.data.id === selectedNodeId) {
+          return 3;
+        }
+        return 2;
+      });
 
-    // Add node labels directly on the node circles
+    // Add node labels with tooltips
     nodes.append("text")
       .attr("dy", ".35em")
       .attr("x", d => d.children ? -18 : 18)
@@ -230,20 +359,28 @@ const LATSVisual: React.FC<SimpleSearchVisualProps> = ({ messages }) => {
         // For root node
         if (d.data.parent_id === null) return "ROOT";
         
-        // Extract action name from action string
+        // Show full action string
         if (d.data.action) {
-          const actionMatch = d.data.action.match(/^([a-zA-Z_]+)\(/);
-          return actionMatch ? actionMatch[1] : "action";
+          return d.data.action;
         }
         
         return d.data.id.toString().slice(-4);
       })
-      .attr("font-size", "14px")
+      .attr("font-size", "15px")
       .attr("font-weight", "500")
       .attr("fill", d => {
+        if (d.data.isSimulated) {
+          return theme === 'dark' ? "#FDBA74" : "#C2410C"; // Orange for simulated node
+        }
+        
+        if (d.data.id === simulationStartNodeId) {
+          return theme === 'dark' ? "#A7F3D0" : "#047857"; // Green for simulation start
+        }
+        
         if (d.data.id === selectedNodeId) {
-          return theme === 'dark' ? "#93C5FD" : "#2563EB";
+          return theme === 'dark' ? "#93C5FD" : "#1D4ED8"; // Blue for selected node
         }
+        
         return theme === 'dark' ? "#FFFFFF" : "#111827";
       });
 
@@ -265,28 +402,66 @@ const LATSVisual: React.FC<SimpleSearchVisualProps> = ({ messages }) => {
     nodes
       .on("mouseover", function(event, d) {
         if (tooltipRef.current) {
-          let content = `<div><strong>Node ID:</strong> ${d.data.id}</div>`;
-          if (d.data.action) content += `<div><strong>Action:</strong> ${d.data.action}</div>`;
-          if (d.data.description) content += `<div><strong>Description:</strong> ${d.data.description}</div>`;
-          if (typeof d.data.value === 'number') content += `<div><strong>Value:</strong> ${d.data.value.toFixed(2)}</div>`;
-          if (typeof d.data.reward === 'number') content += `<div><strong>Reward:</strong> ${d.data.reward.toFixed(2)}</div>`;
-          if (typeof d.data.visits === 'number') content += `<div><strong>Visits:</strong> ${d.data.visits}</div>`;
-          if (d.data.feedback) content += `<div><strong>Feedback:</strong> ${d.data.feedback}</div>`;
+          let tooltipContent = '';
+          
+          // Add description if available
+          if (d.data.description) {
+            tooltipContent += `<p>${d.data.description}</p>`;
+          }
+          
+          // Add node status information
+          const nodeInfo = [];
+          
+          if (d.data.id === simulationStartNodeId) {
+            nodeInfo.push(`<span class="font-semibold text-green-600 dark:text-green-400">Simulation Starting Node</span>`);
+          }
+          
+          if (d.data.isSimulated) {
+            nodeInfo.push(`<span class="font-semibold text-orange-600 dark:text-orange-400">Simulated Node</span>`);
+          }
+          
+          if (d.data.id === selectedNodeId) {
+            nodeInfo.push(`<span class="font-semibold text-blue-600 dark:text-blue-400">Selected Node</span>`);
+          }
+          
+          if (nodeInfo.length > 0) {
+            tooltipContent += `<div class="mt-2">${nodeInfo.join(' | ')}</div>`;
+          }
+          
+          // Add reward info if available
+          if (typeof d.data.reward === 'number') {
+            tooltipContent += `<div class="mt-1">Reward: <span class="font-bold">${d.data.reward.toFixed(2)}</span></div>`;
+          }
+          
+          // Add value info if available
+          if (typeof d.data.value === 'number') {
+            tooltipContent += `<div>Value: <span class="font-bold">${d.data.value.toFixed(2)}</span></div>`;
+          }
+          
+          // Add visits info if available
+          if (typeof d.data.visits === 'number') {
+            tooltipContent += `<div>Visits: <span class="font-bold">${d.data.visits}</span></div>`;
+          }
+          
+          // Add depth info if available
+          if (typeof d.data.depth === 'number') {
+            tooltipContent += `<div>Depth: <span class="font-bold">${d.data.depth}</span></div>`;
+          }
           
           const tooltip = d3.select(tooltipRef.current);
           tooltip.transition()
             .duration(200)
             .style("opacity", .9);
-          tooltip.html(content)
+          tooltip.html(tooltipContent)
             .style("left", (event.pageX + 15) + "px")
-            .style("top", (event.pageY - 30) + "px");
+            .style("top", (event.pageY - 60) + "px");
         }
       })
       .on("mousemove", function(event) {
         if (tooltipRef.current) {
           d3.select(tooltipRef.current)
             .style("left", (event.pageX + 15) + "px")
-            .style("top", (event.pageY - 30) + "px");
+            .style("top", (event.pageY - 28) + "px");
         }
       })
       .on("mouseout", function() {
@@ -307,7 +482,7 @@ const LATSVisual: React.FC<SimpleSearchVisualProps> = ({ messages }) => {
 
     svg.call(zoom);
 
-  }, [treeNodes, selectedNodeId, theme, containerWidth]);
+  }, [treeNodes, selectedNodeId, simulationStartNodeId, simulatedNodes, theme, containerWidth]);
 
   return (
     <div className="w-[30%] bg-white dark:bg-slate-800 rounded-r-lg overflow-hidden">
@@ -318,6 +493,32 @@ const LATSVisual: React.FC<SimpleSearchVisualProps> = ({ messages }) => {
           </svg>
           Tree Visualization
         </h2>
+        
+        {/* Simulation indicator */}
+        {simulationStartNodeId && (
+          <div className="mt-2 flex items-center">
+            <span className="inline-flex items-center px-2.5 py-0.5 rounded-full text-xs font-medium bg-green-100 text-green-800 dark:bg-green-900 dark:text-green-200">
+              <span className="w-2 h-2 mr-1 rounded-full bg-green-500"></span>
+              Simulation Mode
+            </span>
+          </div>
+        )}
+        
+        {/* Legend */}
+        <div className="mt-2 flex flex-wrap gap-2 text-xs">
+          <div className="flex items-center">
+            <span className="w-3 h-3 rounded-full inline-block mr-1 bg-blue-500 dark:bg-blue-600"></span>
+            <span className="text-gray-700 dark:text-gray-300">Selected</span>
+          </div>
+          <div className="flex items-center">
+            <span className="w-3 h-3 rounded-full inline-block mr-1 bg-green-500 dark:bg-green-600"></span>
+            <span className="text-gray-700 dark:text-gray-300">Sim Start</span>
+          </div>
+          <div className="flex items-center">
+            <span className="w-3 h-3 rounded-full inline-block mr-1 bg-orange-500 dark:bg-orange-600"></span>
+            <span className="text-gray-700 dark:text-gray-300">Simulated</span>
+          </div>
+        </div>
       </div>
       <div 
         ref={containerRef} 
