add 5) Browser Set Up, Account Reset, 6) Tree Update & Trajectory Update, 7) node_selected for LATS

Author: TataKKKL

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

@@ -119,6 +119,12 @@ async def run(self, websocket=None) -> list[LATSNode]:
         print_trajectory(best_node)
 
         if websocket:
+            # trajectory_data = self._get_trajectory_data(best_node)
+            # await websocket.send_json({
+            #     "type": "trajectory_update",
+            #     "trajectory": trajectory_data,
+            #     "timestamp": datetime.utcnow().isoformat()
+            # })
             # TODO: use score instead of reward to determine success
             await websocket.send_json({
                 "type": "search_complete",
@@ -168,6 +174,12 @@ async def lats_search(self, websocket=None) -> LATSNode:
                 })
 
             node = self.select_node(self.root_node)
+            if websocket:
+                await websocket.send_json({
+                    "type": "node_selected",
+                    "node_id": id(node),
+                    "timestamp": datetime.utcnow().isoformat()
+                })
 
             if node is None:
                 print("All paths lead to terminal nodes with reward 0. Ending search.")
@@ -207,6 +219,12 @@ async def lats_search(self, websocket=None) -> LATSNode:
             print(f"{GREEN}Tree:{RESET}")
             better_print(self.root_node)
             print(f"")
+            tree_data = self._get_tree_data()
+            await websocket.send_json({
+                "type": "tree_update",
+                "tree": tree_data,
+                "timestamp": datetime.utcnow().isoformat()
+            })
 
             # Step 3: Evaluation
             print(f"")
@@ -224,6 +242,15 @@ async def lats_search(self, websocket=None) -> LATSNode:
             print(f"{GREEN}Tree:{RESET}")
             better_print(self.root_node)
             print(f"")
+            ## send tree update, since evaluation is added to the tree
+            if websocket:
+                tree_data = self._get_tree_data()
+                await websocket.send_json({
+                    "type": "tree_update",
+                    "tree": tree_data,
+                    "timestamp": datetime.utcnow().isoformat()
+                })
+
 
             # Step 4: Simulation
             print(f"{GREEN}Step 4: simulation{RESET}")
@@ -236,7 +263,7 @@ async def lats_search(self, websocket=None) -> LATSNode:
                     "timestamp": datetime.utcnow().isoformat()
                 })
             ## always = 1
-            reward, terminal_node = await self.simulate(max(node.children, key=lambda child: child.value), max_depth=self.config.max_depth, num_simulations=1)
+            reward, terminal_node = await self.simulate(max(node.children, key=lambda child: child.value), max_depth=self.config.max_depth, num_simulations=1, websocket=websocket)
             terminal_nodes.append(terminal_node)
 
             if reward == 1:
@@ -362,7 +389,8 @@ async def evaluate_node(self, node: LATSNode) -> None:
             child.value = score
             child.reward = score
 
-    async def simulate(self, node: LATSNode, max_depth: int = 2, num_simulations=1) -> tuple[float, LATSNode]:
+    ## TODO: make number of simulations configurable
+    async def simulate(self, node: LATSNode, max_depth: int = 2, num_simulations=1, websocket=None) -> tuple[float, LATSNode]:
         """
         Perform a rollout simulation from a node.
         
@@ -378,13 +406,39 @@ async def simulate(self, node: LATSNode, max_depth: int = 2, num_simulations=1)
         print_trajectory(node)
         print("print the entire tree")
         print_entire_tree(self.root_node)
-        return await self.rollout(node, max_depth=max_depth)
+        if websocket:
+            tree_data = self._get_tree_data()
+            await websocket.send_json({
+                "type": "tree_update",
+                "tree": tree_data,
+                "timestamp": datetime.utcnow().isoformat()
+            })
+            trajectory_data = self._get_trajectory_data(node)
+            await websocket.send_json({
+                "type": "trajectory_update",
+                "trajectory": trajectory_data,
+                "timestamp": datetime.utcnow().isoformat()
+            })
+        return await self.rollout(node, max_depth=max_depth, websocket=websocket)
 
-    async def send_completion_request(self, plan, depth, node, trajectory=[]):
+    async def send_completion_request(self, plan, depth, node, trajectory=[], websocket=None):
         print("print the trajectory")
         print_trajectory(node)
         print("print the entire tree")
         print_entire_tree(self.root_node)
+        if websocket:
+            # tree_data = self._get_tree_data()
+            # await websocket.send_json({
+            #     "type": "tree_update",
+            #     "tree": tree_data,
+            #     "timestamp": datetime.utcnow().isoformat()
+            # })
+            trajectory_data = self._get_trajectory_data(node)
+            await websocket.send_json({
+                "type": "trajectory_update",
+                "trajectory": trajectory_data,
+                "timestamp": datetime.utcnow().isoformat()
+            })
 
         if depth >= self.config.max_depth:
             return trajectory, node
@@ -447,20 +501,20 @@ async def send_completion_request(self, plan, depth, node, trajectory=[]):
                 if goal_finished:
                     return trajectory, new_node
 
-                return await self.send_completion_request(plan, depth + 1, new_node, trajectory)
+                return await self.send_completion_request(plan, depth + 1, new_node, trajectory, websocket)
 
             except Exception as e:
                 print(f"Attempt {attempt + 1} failed with error: {e}")
                 if attempt + 1 == retry_count:
                     print("Max retries reached. Skipping this step and retrying the whole request.")
                     # Retry the entire request from the same state
-                    return await self.send_completion_request(plan, depth, node, trajectory)
+                    return await self.send_completion_request(plan, depth, node, trajectory, websocket)
 
         # If all retries and retries of retries fail, return the current trajectory and node
         return trajectory, node
 
 
-    async def rollout(self, node: LATSNode, max_depth: int = 2)-> tuple[float, LATSNode]:
+    async def rollout(self, node: LATSNode, max_depth: int = 2, websocket=None)-> tuple[float, LATSNode]:
         # Reset browser state
         await self._reset_browser()
         path = self.get_path_to_root(node)
@@ -494,11 +548,24 @@ async def rollout(self, node: LATSNode, max_depth: int = 2)-> tuple[float, LATSN
         ## call the prompt agent
         print("current depth: ", len(path) - 1)
         print("max depth: ", self.config.max_depth)
-        trajectory, node = await self.send_completion_request(self.goal, len(path) - 1, node=n, trajectory=trajectory)
+        trajectory, node = await self.send_completion_request(self.goal, len(path) - 1, node=n, trajectory=trajectory, websocket=websocket)
         print("print the trajectory")
         print_trajectory(node)
         print("print the entire tree")
         print_entire_tree(self.root_node)
+        if websocket:
+            # tree_data = self._get_tree_data()
+            # await websocket.send_json({
+            #     "type": "tree_update",
+            #     "tree": tree_data,
+            #     "timestamp": datetime.utcnow().isoformat()
+            # })
+            trajectory_data = self._get_trajectory_data(node)
+            await websocket.send_json({
+                "type": "trajectory_update",
+                "trajectory": trajectory_data,
+                "timestamp": datetime.utcnow().isoformat()
+            })
 
         page = await self.playwright_manager.get_page()
         page_info = await extract_page_info(page, self.config.fullpage, self.config.log_folder)
@@ -801,3 +868,40 @@ def _get_tree_data(self):
             tree_data.append(node_data)
 
         return tree_data
+    
+    def _get_trajectory_data(self, terminal_node: LATSNode):
+        """Get trajectory data in a format suitable for visualization
+        
+        Args:
+            terminal_node: The leaf node to start the trajectory from
+            
+        Returns:
+            list: List of node data dictionaries representing the trajectory
+        """
+        trajectory_data = []
+        path = []
+        
+        # Collect path from terminal to root
+        current = terminal_node
+        while current is not None:
+            path.append(current)
+            current = current.parent
+            
+        # Process nodes in order from root to terminal
+        for level, node in enumerate(reversed(path)):
+            node_data = {
+                "id": id(node),
+                "level": level,
+                "action": node.action if node.action else "ROOT",
+                "description": node.natural_language_description,
+                "visits": node.visits,
+                "value": float(f"{node.value:.3f}") if hasattr(node, 'value') else None,
+                "reward": float(f"{node.reward:.3f}") if hasattr(node, 'reward') else None,
+                "is_terminal": node.is_terminal,
+                "feedback": node.feedback if hasattr(node, 'feedback') else None,
+                "is_root": not hasattr(node, 'parent') or node.parent is None,
+                "is_terminal_node": node == terminal_node
+            }
+            trajectory_data.append(node_data)
+            
+        return trajectory_data 

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

@@ -1145,9 +1145,51 @@ def _get_tree_data(self):
                 "action": node.action if node.action else "ROOT",
                 "description": node.natural_language_description,
                 "depth": node.depth,
-                "is_terminal": node.is_terminal
+                "is_terminal": node.is_terminal,
+                "value": node.value,
+                "visits": node.visits,
+                "reward": node.reward
             }
             tree_data.append(node_data)
 
         return tree_data
+    
+    def _get_trajectory_data(self, terminal_node: LATSNode):
+        """Get trajectory data in a format suitable for visualization
+        
+        Args:
+            terminal_node: The leaf node to start the trajectory from
+            
+        Returns:
+            list: List of node data dictionaries representing the trajectory
+        """
+        trajectory_data = []
+        path = []
+        
+        # Collect path from terminal to root
+        current = terminal_node
+        while current is not None:
+            path.append(current)
+            current = current.parent
+            
+        # Process nodes in order from root to terminal
+        for level, node in enumerate(reversed(path)):
+            node_data = {
+                "id": id(node),
+                "level": level,
+                "action": node.action if node.action else "ROOT",
+                "description": node.natural_language_description,
+                "visits": node.visits,
+                "value": float(f"{node.value:.3f}") if hasattr(node, 'value') else None,
+                "reward": float(f"{node.reward:.3f}") if hasattr(node, 'reward') else None,
+                "is_terminal": node.is_terminal,
+                "feedback": node.feedback if hasattr(node, 'feedback') else None,
+                "is_root": not hasattr(node, 'parent') or node.parent is None,
+                "is_terminal_node": node == terminal_node
+            }
+            trajectory_data.append(node_data)
+            
+        return trajectory_data 
+
+    
 

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

@@ -0,0 +1,36 @@
+    def _get_trajectory_data(self, terminal_node: LATSNode):
+        """Get trajectory data in a format suitable for visualization
+        
+        Args:
+            terminal_node: The leaf node to start the trajectory from
+            
+        Returns:
+            list: List of node data dictionaries representing the trajectory
+        """
+        trajectory_data = []
+        path = []
+        
+        # Collect path from terminal to root
+        current = terminal_node
+        while current is not None:
+            path.append(current)
+            current = current.parent
+            
+        # Process nodes in order from root to terminal
+        for level, node in enumerate(reversed(path)):
+            node_data = {
+                "id": id(node),
+                "level": level,
+                "action": node.action if node.action else "ROOT",
+                "description": node.natural_language_description,
+                "visits": node.visits,
+                "value": float(f"{node.value:.3f}") if hasattr(node, 'value') else None,
+                "reward": float(f"{node.reward:.3f}") if hasattr(node, 'reward') else None,
+                "is_terminal": node.is_terminal,
+                "feedback": node.feedback if hasattr(node, 'feedback') else None,
+                "is_root": not hasattr(node, 'parent') or node.parent is None,
+                "is_terminal_node": node == terminal_node
+            }
+            trajectory_data.append(node_data)
+            
+        return trajectory_data 

visual-tree-search-backend/test/test-tree-search-ws-lats.py

@@ -85,8 +85,11 @@ async def connect_and_test_search(
                     if status == "scored":
                         logger.info(f"Node score: {data.get('score')}")
 
+                elif msg_type == "trajectory_update":
+                    logger.info(f"Trajectory update received with {data.get('trajectory')}")
+                
                 elif msg_type == "tree_update":
-                    logger.info(f"Tree update received with {len(data.get('nodes', []))} nodes")
+                    logger.info(f"Tree update received with {data.get('tree')}")
 
                 elif msg_type == "best_path_update":
                     logger.info(f"Best path update: score={data.get('score')}, path length={len(data.get('path', []))}")

visual-tree-search-backend/test/test-tree-search-ws-simple.py

@@ -80,7 +80,7 @@ async def connect_and_test_search(
                         logger.info(f"Node score: {data.get('score')}")
 
                 elif msg_type == "tree_update":
-                    logger.info(f"Tree update received with {len(data.get('nodes', []))} nodes")
+                    logger.info(f"Tree update received with {data.get('tree')}")
 
                 elif msg_type == "best_path_update":
                     logger.info(f"Best path update: score={data.get('score')}, path={data.get('path')}")