Merge pull request #95 from PathOnAI/fix_is_terminal

fix is_terminal

Author: TataKKKL

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

@@ -490,6 +490,7 @@ const TreeVisual: React.FC<TreeVisualProps> = ({
             tooltipContent += `<div class="mt-2">${nodeInfo.join(' | ')}</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>`;
@@ -504,6 +505,11 @@ const TreeVisual: React.FC<TreeVisualProps> = ({
           if (typeof d.data.depth === 'number') {
             tooltipContent += `<div>Depth: <span class="font-bold">${d.data.depth}</span></div>`;
           }
+
+          // Add is_terminal info if available
+          if (typeof d.data.is_terminal === 'boolean') {
+            tooltipContent += `<div>Is Terminal: <span class="font-bold">${d.data.is_terminal ? 'Yes' : 'No'}</span></div>`;
+          }
 
           const tooltip = d3.select(tooltipRef.current);
           tooltip.transition()

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

@@ -398,19 +398,10 @@ async def node_expansion(self, node: LATSNode, websocket = None) -> None:
                 goal=node.goal,
                 parent=node
             )
+            if child.depth == self.config.max_depth:
+                child.is_terminal = True
             node.children.append(child)
             await self.websocket_node_created(child, node, websocket=websocket)
-            
-            # Send child creation update if websocket is provided
-            # if websocket:
-            #     await websocket.send_json({
-            #         "type": "node_created",
-            #         "node_id": id(child),
-            #         "parent_id": id(node),
-            #         "action": child.action,
-            #         "description": child.natural_language_description,
-            #         "timestamp": datetime.utcnow().isoformat()
-            #     })
 
 
      # node evaluation
@@ -420,17 +411,17 @@ async def node_children_evaluation(self, node: LATSNode) -> None:
         print(f"{GREEN}-- total {len(node.children)} children to evaluate:{RESET}")
         for i, child in enumerate(node.children):
             print(f"{GREEN}--- evaluating child {i+1}...{RESET}")
-            if child.is_terminal:
+            # if child.is_terminal:
+            #     score = 0
+            # else:
+            trajectory = child.get_trajectory()
+            if len(trajectory) == 0:
                 score = 0
             else:
-                trajectory = child.get_trajectory()
-                if len(trajectory) == 0:
-                    score = 0
-                else:
-                    prompt = create_llm_prompt(trajectory, self.goal)
-                    # , child.observation.image
-                    result = score_trajectory_with_openai(prompt, openai_client, self.config.evaluation_model)
-                    score = result["overall_score"]
+                prompt = create_llm_prompt(trajectory, self.goal)
+                # , child.observation.image
+                result = score_trajectory_with_openai(prompt, openai_client, self.config.evaluation_model)
+                score = result["overall_score"]
             scores.append(score)
 
         for child, score in zip(node.children, scores):
@@ -454,19 +445,19 @@ async def node_evaluation(self, node: LATSNode) -> None:
 
             try:
                 # Score the trajectory
-                if node.is_terminal:
+                # if node.is_terminal:
+                #     score = 0
+                # else:
+                if len(trajectory) == 0:
                     score = 0
                 else:
-                    if len(trajectory) == 0:
-                        score = 0
-                    else:
-                        prompt = create_llm_prompt(trajectory, self.goal)
-                        result = score_trajectory_with_openai(
-                            prompt, 
-                            openai_client, 
-                            model=self.config.evaluation_model
-                        )
-                        score = result["overall_score"]
+                    prompt = create_llm_prompt(trajectory, self.goal)
+                    result = score_trajectory_with_openai(
+                        prompt, 
+                        openai_client, 
+                        model=self.config.evaluation_model
+                    )
+                    score = result["overall_score"]
 
             except Exception as e:
                 error_msg = f"Error scoring node {id(node)}: {str(e)}"

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

@@ -88,7 +88,9 @@ def __init__(
         self.value = 0.0
         self.depth = 0 if parent is None else parent.depth + 1
         self.is_terminal = False
-        # self.reward = 0.0
+        # The goal has been achieved;
+        # The maximum depth has been reached;
+        # A failure condition has been triggered.
         self.exhausted = False  # If all children are terminal
         self.em = 0.0  # Exact match, evaluation metric
         self.observation: Optional[Observation] = None
@@ -106,6 +108,16 @@ def uct(self) -> float:
         return self.value / self.visits + np.sqrt(2 * np.log(self.parent.visits) / self.visits)
 
     def get_best_leaf(self) -> 'LATSNode':
+        """
+        Recursively get the best leaf node from the current node.
+
+        The method searches through unfinished (non-terminal) children,
+        selects the one with the highest UCT score, and continues the 
+        search recursively until a leaf node (with no unfinished children) is reached.
+        
+        Returns:
+            LATSNode: The best leaf node for expansion based on UCT.
+        """
         unfinished_children = [c for c in self.children if not c.is_terminal]
         if not unfinished_children:
             return self

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

@@ -361,19 +361,7 @@ async def mcts_search(self, websocket=None) -> Optional[LATSNode]:
                         print(f"Node {node.action}:")
                         print(f"  Visits: {node.visits}")
                         print(f"  Value: {old_value:.3f} -> {node.value:.3f}")
-                    # add websocket information, just use websocket here
-                    # if websocket:
-                    #     await websocket.send_json({
-                    #         "type": "backpropagation",
-                    #         "node_id": id(node),
-                    #         "node_parent_id": id(node.parent),
-                    #         "node_action": node.action,
-                    #         "node_value": node.value,
-                    #         "node_visits": node.visits,
-                    #         "node_old_value": old_value,
-                    #         "node_description": node.natural_language_description,
-                    #     })
-
+                        
                 tree_data = self._get_tree_data()
                 print_entire_tree(self.root_node)
                 print(tree_data)

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

@@ -47,17 +47,8 @@ async def bfs(self, websocket=None):
             for _ in range(level_size):
                 current_node = queue.popleft()
                 queue_set.remove(current_node)  # Remove from queue tracking
-                
-                # Skip if we've already visited this node
-                if current_node in visited:
-                    continue
-                    
                 visited.add(current_node)
 
-                # Skip terminal nodes
-                if current_node.is_terminal:
-                    continue
-                
                 # Expand current node if it hasn't been expanded yet and hasn't reached max_depth
                 # node expansion for the next level
                 if not current_node.children and current_node.depth < self.config.max_depth:
@@ -78,7 +69,7 @@ async def bfs(self, websocket=None):
 
                 # Add non-terminal children to queue for next level if they haven't reached max_depth
                 for child in current_node.children:
-                    if not child.is_terminal and child not in visited and child not in queue_set and child.depth < self.config.max_depth:
+                    if child not in visited and child not in queue_set and child.depth <= self.config.max_depth:
                         queue.append(child)
                         queue_set.add(child)  # Add to queue tracking
 
@@ -134,76 +125,50 @@ async def bfs(self, websocket=None):
 
         return None
 
-    # TODO: first evaluate, then expansion
-    async def dfs(self, websocket=None) -> List[Dict[str, Any]]:
-        stack = [self.root_node]
+    async def dfs(self, websocket=None):
+        stack = [self.root_node]  # Use a list as a stack
         stack_set = {self.root_node}  # Track nodes in stack
         best_score = float('-inf')
         best_path = None
         best_node = None
         visited = set()  # Track visited nodes to avoid cycles
-        current_path = []  # Track current path for DFS
-        
-        # # Get the live browser URL during initial setup
-        # live_browser_url, session_id = await self._reset_browser(websocket)
-    
 
         while stack:
-            current_node = stack[-1]  # Peek at the top node without removing it
-            
-            # Skip if we've already visited this node
-            if current_node in visited:
-                stack.pop()
-                stack_set.remove(current_node)
-                if current_path:
-                    current_path.pop()  # Remove from current path
-                continue
-                
+            # Get the top node from the stack
+            current_node = stack.pop()
+            stack_set.remove(current_node)  # Remove from stack tracking
             visited.add(current_node)
-            current_path.append(current_node)  # Add to current path
 
-            # Skip terminal nodes
-            if current_node.is_terminal:
-                print(f"Node {id(current_node)} is terminal")
-                stack.pop()
-                stack_set.remove(current_node)
-                current_path.pop()  # Remove from current path
-                continue
-                
             # Expand current node if it hasn't been expanded yet and hasn't reached max_depth
-            # stage 1: node expansion
             if not current_node.children and current_node.depth < self.config.max_depth:
-                    ## during the node expansion process, reset browser for each node
+                # Reset browser for each node expansion
                 live_browser_url, session_id = await self._reset_browser(websocket)
-                # await self.websocket_step_start(step=1, step_name="node_expansion", websocket=websocket)
                 await self.websocket_node_selection(current_node, websocket=websocket)
                 await self.node_expansion(current_node, websocket)
                 tree_data = self._get_tree_data()
+                
                 if websocket:
                     await self.websocket_tree_update(type="tree_update_node_expansion", websocket=websocket, tree_data=tree_data)
                 else:
                     print_entire_tree(self.root_node)
 
-            # Get the path from root to this node
-            path = self.get_path_to_root(current_node)
+            # Node evaluation
             await self.node_evaluation(current_node)
             tree_data = self._get_tree_data()
             if websocket:
                 await self.websocket_tree_update(type="tree_update_node_evaluation", websocket=websocket, tree_data=tree_data)
             else:
                 print("after evaluation")
                 print_entire_tree(self.root_node)
-            path = self.get_path_to_root(current_node)
 
-
+            path = self.get_path_to_root(current_node)
             score = current_node.value
 
             # Update best path if this score is better
             if score > best_score:
                 best_score = score
                 best_path = path
                 best_node = current_node
-
 
             print(f"Node {id(current_node)} score: {score}")
 
@@ -213,23 +178,16 @@ async def dfs(self, websocket=None) -> List[Dict[str, Any]]:
 
                 # Send completion update if websocket is provided
                 await self.websocket_search_complete("success", score, current_node.get_trajectory(), websocket=websocket) 
-                await self.playwright_manager.close()               
+                await self.playwright_manager.close()
+                
                 return current_node
-                        
-            # Add non-terminal children to stack in reverse order
-            has_unvisited_children = False
+            
+            # Add children to stack in reverse order so that the first child is processed first
+            # This maintains the left-to-right exploration order similar to BFS
             for child in reversed(current_node.children):
-                if not child.is_terminal and child not in visited and child not in stack_set:
+                if child not in visited and child not in stack_set and child.depth <= self.config.max_depth:
                     stack.append(child)
                     stack_set.add(child)  # Add to stack tracking
-                    has_unvisited_children = True
-                    break  # Only add one child at a time for DFS
-            
-            # If no unvisited children, remove current node from stack
-            if not has_unvisited_children:
-                stack.pop()
-                stack_set.remove(current_node)
-                current_path.pop()  # Remove from current path
 
         # If we've exhausted all nodes and haven't found a perfect solution,
         # return the best path we found
@@ -249,5 +207,4 @@ async def dfs(self, websocket=None) -> List[Dict[str, Any]]:
         await self.websocket_search_complete("failure", 0, None, websocket=websocket)
         await self.playwright_manager.close()
 
-        return None
-            
+        return None