From cc87844825310d6e8546a6bfd68655f85d89e68a Mon Sep 17 00:00:00 2001
From: Kerem Yilmaz <kerem@skyvern.com>
Date: Sat, 13 Apr 2024 00:38:01 -0700
Subject: [PATCH] Pass current url to extract-action prompt (#184)

---
 skyvern/forge/agent.py                          | 9 +++++++--
 skyvern/forge/prompts/skyvern/extract-action.j2 | 6 ++++--
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/skyvern/forge/agent.py b/skyvern/forge/agent.py
index 6bf19314..79da5f5a 100644
--- a/skyvern/forge/agent.py
+++ b/skyvern/forge/agent.py
@@ -608,12 +608,17 @@ class ForgeAgent(Agent):
         action_results_str = json.dumps([action_result.model_dump() for action_result in action_results])
         # Generate the extract action prompt
         navigation_goal = task.navigation_goal
+        starting_url = task.url
+        current_url = (
+            await browser_state.page.evaluate("() => document.location.href") if browser_state.page else starting_url
+        )
         extract_action_prompt = prompt_engine.load_prompt(
             "extract-action",
             navigation_goal=navigation_goal,
             navigation_payload_str=json.dumps(task.navigation_payload),
-            url=task.url,
-            elements=scraped_page.element_tree_trimmed,  # scraped_page.element_tree,
+            starting_url=starting_url,
+            current_url=current_url,
+            elements=scraped_page.element_tree_trimmed,
             data_extraction_goal=task.data_extraction_goal,
             action_history=action_results_str,
             error_code_mapping_str=json.dumps(task.error_code_mapping) if task.error_code_mapping else None,
diff --git a/skyvern/forge/prompts/skyvern/extract-action.j2 b/skyvern/forge/prompts/skyvern/extract-action.j2
index 10e226da..2195241b 100644
--- a/skyvern/forge/prompts/skyvern/extract-action.j2
+++ b/skyvern/forge/prompts/skyvern/extract-action.j2
@@ -6,7 +6,7 @@ Each element is tagged with an ID.
 If you see any information in red in the page screenshot, this means a condition wasn't satisfied. prioritize actions with the red information.
 If you see a popup in the page screenshot, prioritize actions on the popup.
 
-{% if "lever" in url %}
+{% if "lever" in starting_url %}
 DO NOT UPDATE ANY LOCATION FIELDS
 {% endif %}
 
@@ -41,11 +41,13 @@ Reply in JSON format with the following keys:
 Consider the action history from the last step and the screenshot together, if actions from the last step don't yield positive impact, try other actions or other action combinations.
 {% endif %}
 
-Clickable elements from `{{ url }}`:
+Clickable elements from `{{ current_url }}`:
 ```
 {{ elements }}
 ```
 
+The URL of the page you're on right now is `{{ current_url }}`.
+
 User goal:
 ```
 {{ navigation_goal }}