optimize svg logic (#701)

skyvern/forge/agent.py

@@ -532,6 +532,7 @@ class ForgeAgent:
                 task,
                 step,
                 browser_state,
+                organization,
             )
             detailed_agent_step_output.scraped_page = scraped_page
             detailed_agent_step_output.extract_action_prompt = extract_action_prompt
@@ -890,6 +891,7 @@ class ForgeAgent:
         step: Step,
         browser_state: BrowserState,
         scrape_type: ScrapeType,
+        organization: Organization | None = None,
     ) -> ScrapedPage | None:
         if scrape_type == ScrapeType.NORMAL:
             pass
@@ -912,7 +914,7 @@ class ForgeAgent:
         return await scrape_website(
             browser_state,
             task.url,
-            app.AGENT_FUNCTION.cleanup_element_tree,
+            app.AGENT_FUNCTION.cleanup_element_tree_factory(task=task, step=step, organization=organization),
             scrape_exclude=app.scrape_exclude,
         )
 
@@ -921,6 +923,7 @@ class ForgeAgent:
         task: Task,
         step: Step,
         browser_state: BrowserState,
+        organization: Organization | None = None,
     ) -> tuple[ScrapedPage, str]:
         # start the async tasks while running scrape_website
         self.async_operation_pool.run_operation(task.task_id, AgentPhase.scrape)
@@ -934,7 +937,11 @@ class ForgeAgent:
         for idx, scrape_type in enumerate(SCRAPE_TYPE_ORDER):
             try:
                 scraped_page = await self._scrape_with_type(
-                    task=task, step=step, browser_state=browser_state, scrape_type=scrape_type
+                    task=task,
+                    step=step,
+                    browser_state=browser_state,
+                    scrape_type=scrape_type,
+                    organization=organization,
                 )
                 break
             except FailedToTakeScreenshot as e: