Update tools/scroll_page.py

tools/scroll_page.py

@@ -1,37 +1,57 @@
 from smolagents.tools import Tool
-from helium import scroll_down, S
+from helium import scroll_down, scroll_up, get_driver
 from selenium.webdriver.common.by import By
 
-def scroll_page(driver, pixels=600, selector=None):
+def scroll_page(driver, selector=None, num_pixels=1200, direction="down"):
     """
-    Scroll the page by a specified number of pixels or to a specific element.
+    Scroll the page to a specific element or by a number of pixels.
     
     Args:
         driver: Selenium WebDriver instance
-        pixels (int): Number of pixels to scroll down (default: 600)
-        selector (str): Optional CSS selector to scroll to a specific element
+        selector (str): CSS selector to scroll to (optional)
+        num_pixels (int): Number of pixels to scroll (default: 1200)
+        direction (str): Scroll direction ('down' or 'up') (default: 'down')
     
     Returns:
-        str: Observation message indicating scroll action
+        str: Result of the scroll action
     """
-    if selector:
-        try:
-            element = S(selector)
-            driver.execute_script("arguments[0].scrollIntoView(true);", element.web_element)
-            return f"Scrolled to element with selector: {selector}"
-        except Exception as e:
-            return f"Failed to scroll to selector {selector}: {str(e)}"
-    else:
-        driver.execute_script(f"window.scrollBy(0, {pixels});")
-        return f"Scrolled down by {pixels} pixels"
+    try:
+        if selector:
+            element = driver.find_element(By.CSS_SELECTOR, selector)
+            driver.execute_script("arguments[0].scrollIntoView(true);", element)
+            return f"Scrolled to element with selector {selector}"
+        else:
+            if direction == "down":
+                scroll_down(num_pixels)
+                return f"Scrolled down {num_pixels} pixels"
+            elif direction == "up":
+                scroll_up(num_pixels)
+                return f"Scrolled up {num_pixels} pixels"
+            else:
+                return f"Invalid direction: {direction}"
+    except Exception as e:
+        return f"Failed to scroll: {str(e)}"
 
 # Register the tool
-tool = Tool(
+scroll_page_tool = Tool(
     name="scroll_page",
-    description="Scrolls the page by a specified number of pixels or to an element with a given CSS selector.",
+    description="Scrolls the page to a specific element or by a number of pixels.",
     inputs={
-        "pixels": {"type": "int", "default": 600, "description": "Number of pixels to scroll down"},
-        "selector": {"type": "str", "default": None, "description": "CSS selector of the element to scroll to"}
+        "selector": {
+            "type": "str",
+            "default": None,
+            "description": "CSS selector to scroll to"
+        },
+        "num_pixels": {
+            "type": "int",
+            "default": 1200,
+            "description": "Number of pixels to scroll"
+        },
+        "direction": {
+            "type": "str",
+            "default": "down",
+            "description": "Scroll direction: 'down' or 'up'"
+        }
     },
     output_type="str",
     function=scroll_page