Merge pull request #63 from PathOnAI/refactoring-cleanup

Refactoring cleanup

Author: TataKKKL

.gitignore

@@ -1,3 +1,4 @@
+visual-tree-search-backend/app/api/test_logs/*
 visual-tree-search-backend/log/*
 log/*
 shopping.json

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

@@ -34,10 +34,6 @@
 openai_client = OpenAI()
 
 
-## TODO: remove account reset websocket message
-## browser setup message, ok to leave there in the _reset_browser method
-
-
 class BaseAgent:    
     # no need to pass an initial playwright_manager to the agent class
     def __init__(
@@ -381,6 +377,10 @@ async def websocket_search_complete(self, status, score, path, websocket=None):
                     "path": path,
                     "timestamp": datetime.utcnow().isoformat()
                 })
+        else:
+            print(f"Search complete: {GREEN}{status}{RESET}")
+            print(f"Search score: {GREEN}{score}{RESET}")
+            print(f"Search path: {GREEN}{path}{RESET}")
 
     # shared, not implemented, BFS, DFS and LATS has its own node selection logic
     async def node_selection(self, node, websocket = None):
@@ -485,31 +485,19 @@ def backpropagate(self, node: LATSNode, value: float) -> None:
             node = node.parent
 
     # shared
-    async def simulation(self, node: LATSNode, max_depth: int = 2, num_simulations=1, websocket=None) -> tuple[float, LATSNode]:
+    async def simulation(self, node: LATSNode, websocket=None) -> tuple[float, LATSNode]:
         depth = node.depth
+        num_simulations = self.config.num_simulations
+        max_depth = self.config.max_depth
         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 self.websocket_tree_update(type="tree_update_simulation", tree_data=tree_data, websocket=websocket)
-            # 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)
+        return await self.rollout(node, websocket=websocket)
 
     # refactor simulation, rollout, send_completion_request methods
     # TODO: check, score as reward and then update value of the starting node?
-    async def rollout(self, node: LATSNode, max_depth: int = 2, websocket=None)-> tuple[float, LATSNode]:
+    async def rollout(self, node: LATSNode, websocket=None)-> tuple[float, LATSNode]:
         # Reset browser state
         await self._reset_browser()
         path = self.get_path_to_root(node)
@@ -540,23 +528,14 @@ async def rollout(self, node: LATSNode, max_depth: int = 2, websocket=None)-> tu
                     "action": n.action,
                     "feedback": n.feedback
                 })
-        ## call the prompt agent
         print("current depth: ", len(path) - 1)
         print("max depth: ", self.config.max_depth)
 
-        ## find a better name for this
         trajectory, terminal_node = await self.send_completion_request(self.goal, len(path) - 1, node=n, trajectory=trajectory, websocket=websocket)
         print("print the trajectory")
         print_trajectory(terminal_node)
         print("print the entire tree")
         print_entire_tree(self.root_node)
-        # if websocket:
-        #     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)
@@ -583,12 +562,6 @@ async def send_completion_request(self, plan, depth, node, trajectory=[], websoc
         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",
@@ -684,15 +657,7 @@ async def generate_children(self, node: LATSNode, websocket=None) -> list[dict]:
         path = self.get_path_to_root(node)
 
         # Execute path
-        for n in path[1:]:  # Skip root node
-            # if websocket:
-            #     await websocket.send_json({
-            #         "type": "replaying_action",
-            #         "node_id": id(n),
-            #         "action": n.action,
-            #         "timestamp": datetime.utcnow().isoformat()
-            #     })
-            
+        for n in path[1:]:  # Skip root node       
             success = await playwright_step_execution(
                 n,
                 self.goal,
@@ -702,12 +667,6 @@ async def generate_children(self, node: LATSNode, websocket=None) -> list[dict]:
             )
             if not success:
                 n.is_terminal = True
-                # if websocket:
-                #     await websocket.send_json({
-                #         "type": "replay_failed",
-                #         "node_id": id(n),
-                #         "timestamp": datetime.utcnow().isoformat()
-                #     })
                 return []
 
             if not n.feedback:
@@ -716,26 +675,13 @@ async def generate_children(self, node: LATSNode, websocket=None) -> list[dict]:
                     n.natural_language_description,
                     self.playwright_manager,
                 )
-                # if websocket:
-                #     await websocket.send_json({
-                #         "type": "feedback_generated",
-                #         "node_id": id(n),
-                #         "feedback": n.feedback,
-                #         "timestamp": datetime.utcnow().isoformat()
-                #     })
 
         time.sleep(3)
         page = await self.playwright_manager.get_page()
         page_info = await extract_page_info(page, self.config.fullpage, self.config.log_folder)
 
         messages = [{"role": "user", "content": f"Action is: {n.action}"} for n in path[1:]]
-        
-        # if websocket:
-        #     await websocket.send_json({
-        #         "type": "generating_actions",
-        #         "node_id": id(node),
-        #         "timestamp": datetime.utcnow().isoformat()
-        #     })
+
 
         next_actions = await extract_top_actions(
             [{"natural_language_description": n.natural_language_description, "action": n.action, "feedback": n.feedback} for n in path[1:]],
@@ -779,23 +725,8 @@ async def generate_children(self, node: LATSNode, websocket=None) -> list[dict]:
                         action["element"] = element
                 except Exception as e:
                     action["element"] = None
-                    # if websocket:
-                    #     await websocket.send_json({
-                    #         "type": "element_location_failed",
-                    #         "action": action["action"],
-                    #         "error": str(e),
-                    #         "timestamp": datetime.utcnow().isoformat()
-                    #     })
                 children.append(action)
 
         if not children:
-            node.is_terminal = True
-            # if websocket:
-            #     await websocket.send_json({
-            #         "type": "node_terminal",
-            #         "node_id": id(node),
-            #         "reason": "no_valid_actions",
-            #         "timestamp": datetime.utcnow().isoformat()
-            #     })
-        
+            node.is_terminal = True        
         return children

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

@@ -19,6 +19,7 @@ async def run(self, websocket=None) -> list[LATSNode]:
 
         best_node = await self.lats_search(websocket)
         print_trajectory(best_node)
+        return best_node
 
     async def lats_search(self, websocket=None):
             terminal_nodes = []
@@ -71,11 +72,12 @@ async def lats_search(self, websocket=None):
                 await self.websocket_step_start(step=4, step_name="simulation", websocket=websocket)
                 selected_node = max(node.children, key=lambda child: child.value)
                 await self.websocket_node_selection(selected_node, websocket=websocket, type="node_selected_for_simulation")
-                reward, terminal_node = await self.simulation(selected_node, max_depth=self.config.max_depth, num_simulations=1, websocket=websocket)
+                reward, terminal_node = await self.simulation(selected_node, websocket=websocket)
                 terminal_nodes.append(terminal_node)
                 await self.websocket_simulation_result(reward, terminal_node, websocket=websocket)
 
                 if reward == 1:
+                    await self.websocket_search_complete("success", reward, terminal_node.get_trajectory(), websocket=websocket)
                     return terminal_node
 
                 # Step 5: Backpropagation
@@ -96,10 +98,12 @@ async def lats_search(self, websocket=None):
             ## temp change: if reward is the same, choose the deeper node
             best_child = max(all_nodes_list, key=lambda x: (x.reward, x.depth))
 
-            if best_child.reward == 1:
+            if best_child.value >= 0.75:
                 print("Successful trajectory found")
+                await self.websocket_search_complete("success", best_child.value, best_child.get_trajectory(), websocket=websocket)
             else:
                 print("Unsuccessful trajectory found")
+                await self.websocket_search_complete("partial_success", best_child.value, best_child.get_trajectory(), websocket=websocket)
             await self.playwright_manager.close()
 
             return best_child if best_child is not None else self.root_node

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

@@ -109,27 +109,27 @@ async def bfs(self, websocket=None):
                     print(f"Found satisfactory solution with score {score}")
 
                     # Send completion update if websocket is provided
-                    await self.websocket_search_complete("success", score, current_node.get_trajectory(), websocket=None) 
+                    await self.websocket_search_complete("success", score, current_node.get_trajectory(), websocket=websocket) 
 
-                    return [{"action": node.action} for node in path[1:]]
+                    return current_node
 
         # If we've exhausted all nodes and haven't found a perfect solution,
         # return the best path we found
         if best_path:
             print(f"Returning best path found with score {best_score}")
 
             # Send completion update if websocket is provided
-            await self.websocket_search_complete("partial_success", best_score, best_node.get_trajectory(), websocket=None)
+            await self.websocket_search_complete("partial_success", best_score, best_node.get_trajectory(), websocket=websocket)
 
-            return [{"action": node.action} for node in best_path[1:]]
+            return best_node
 
         # If no path was found at all
         print("No valid path found")
 
         # Send failure update if websocket is provided
-        await self.websocket_search_complete("failure", 0, None, websocket=None)
+        await self.websocket_search_complete("failure", 0, None, websocket=websocket)
 
-        return []
+        return None
 
     # TODO: first evaluate, then expansion
     async def dfs(self, websocket=None) -> List[Dict[str, Any]]:
@@ -209,8 +209,8 @@ async def dfs(self, websocket=None) -> List[Dict[str, Any]]:
                 print(f"Found satisfactory solution with score {score}")
 
                 # Send completion update if websocket is provided
-                await self.websocket_search_complete("success", score, current_node.get_trajectory(), websocket=None)                
-                return [{"action": node.action} for node in path[1:]]
+                await self.websocket_search_complete("success", score, current_node.get_trajectory(), websocket=websocket)                
+                return current_node
 
             # Add non-terminal children to stack in reverse order
             has_unvisited_children = False
@@ -233,15 +233,15 @@ async def dfs(self, websocket=None) -> List[Dict[str, Any]]:
             print(f"Returning best path found with score {best_score}")
 
             # Send completion update if websocket is provided
-            await self.websocket_search_complete("partial_success", best_score, best_node.get_trajectory(), websocket=None)
+            await self.websocket_search_complete("partial_success", best_score, best_node.get_trajectory(), websocket=websocket)
 
-            return [{"action": node.action} for node in best_path[1:]]
+            return best_node
 
         # If no path was found at all
         print("No valid path found")
 
         # Send failure update if websocket is provided
-        await self.websocket_search_complete("failure", 0, None, websocket=None)
+        await self.websocket_search_complete("failure", 0, None, websocket=websocket)
 
-        return []
+        return None
 

visual-tree-search-backend/app/api/lwats/core_async/config.py

@@ -22,7 +22,7 @@ class AgentConfig:
     branching_factor: int = 5
     iterations: int = 1
     max_depth: int = 3
-    num_simulations: int = 100
+    num_simulations: int = 1
     account_reset: bool = True
 
     # Features

visual-tree-search-backend/app/api/test_logs/screenshots/final_state.png

@@ -124,6 +124,9 @@ async def connect_and_test_search(
                 color = COLORS.get(msg_type, COLORS['reset'])
                 print(f"\nWebSocket message - Type: {color}{msg_type}{COLORS['reset']}")
                 print(f"Raw message: {json.dumps(data, indent=2)}")
+
+                if msg_type == "search_complete":
+                    break
 
             except websockets.exceptions.ConnectionClosed:
                 logger.warning("WebSocket connection closed")

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

@@ -116,6 +116,9 @@ async def connect_and_test_search(
                 color = COLORS.get(msg_type, COLORS['reset'])
                 print(f"\nWebSocket message - Type: {color}{msg_type}{COLORS['reset']}")
                 print(f"Raw message: {json.dumps(data, indent=2)}")
+
+                if msg_type == "search_complete":
+                    break
 
             except websockets.exceptions.ConnectionClosed:
                 logger.warning("WebSocket connection closed")