add cleanup function (#631)

skyvern/forge/agent_functions.py

@@ -1,3 +1,5 @@
+from typing import Dict, List
+
 from playwright.async_api import Page
 
 from skyvern.exceptions import StepUnableToExecuteError
@@ -8,6 +10,11 @@ from skyvern.forge.sdk.schemas.tasks import Task, TaskStatus
 from skyvern.webeye.browser_factory import BrowserState
 
 
+def _remove_rect(element: dict) -> None:
+    if "rect" in element:
+        del element["rect"]
+
+
 class AgentFunction:
     async def validate_step_execution(
         self,
@@ -56,3 +63,29 @@ class AgentFunction:
         page: Page,
     ) -> list[AsyncOperation]:
         return []
+
+    async def cleanup_element_tree(
+        self,
+        url: str,
+        element_tree: List[Dict],
+    ) -> List[Dict]:
+        """
+        Remove rect and attribute.unique_id from the elements.
+        The reason we're doing it is to
+        1. reduce unnecessary data so that llm get less distrction
+        TODO later: 2. reduce tokens sent to llm to save money
+        :param elements: List of elements to remove xpaths from.
+        :return: List of elements without xpaths.
+        """
+        queue = []
+        for element in element_tree:
+            queue.append(element)
+        while queue:
+            queue_ele = queue.pop(0)
+            _remove_rect(queue_ele)
+            # TODO: we can come back to test removing the unique_id
+            # from element attributes to make sure this won't increase hallucination
+            # _remove_unique_id(queue_ele)
+            if "children" in queue_ele:
+                queue.extend(queue_ele["children"])
+        return element_tree