Merge pull request #74 from PathOnAI/lats-page-v2

Lats page v2

Author: TataKKKL

visual-tree-search-app/components/LATSVisual.tsx

@@ -95,7 +95,7 @@ const LATSVisual: React.FC<SimpleSearchVisualProps> = ({ messages }) => {
         }
 
         // Handle tree structure updates
-        if ((data.type === 'tree_update_node_expansion' || data.type === 'tree_update_node_children_evaluation') 
+        if ((data.type === 'tree_update_node_expansion' || data.type === 'tree_update_node_children_evaluation' || data.typ === 'tree_update_node_backpropagation') 
             && Array.isArray(data.tree)) {
           // Preserve simulation flags when updating from tree
           if (updatedTreeNodes.some(node => node.isSimulated)) {
@@ -151,11 +151,8 @@ const LATSVisual: React.FC<SimpleSearchVisualProps> = ({ messages }) => {
 
         // 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
-          }));
+          // Remove all simulated nodes from the tree
+          updatedTreeNodes = updatedTreeNodes.filter(node => !node.isSimulated);
 
           newSimulatedNodes = []; // Clear simulated nodes list
           newSimulationStartNodeId = null; // Clear simulation start node

visual-tree-search-backend/app/api/lwats/agents_async/SearchAgents/base_agent.py

@@ -137,7 +137,7 @@ async def remove_simulated_trajectory(self, starting_node, terminal_node: LATSNo
             }
             trajectory_data.append(node_data)
 
-        await self.websocket_simulation_removed(trajectory_data, websocket=None)
+        await self.websocket_simulation_removed(trajectory_data, websocket=websocket)
         pass
 
     def _get_trajectory_data(self, terminal_node: LATSNode):