Update tools/interact_element.py

tools/interact_element.py

@@ -2,85 +2,56 @@ from smolagents.tools import Tool
 from helium import write, press, click, Text, Link, S
 from selenium.webdriver.common.keys import Keys
 
-def interact_element(driver, selector=None, text=None, action="click", input_text=None, key=None):
-    """
-    Interact with a web element (click, fill text, press keys).
-    
-    Args:
-        driver: Selenium WebDriver instance
-        selector (str): CSS selector to target the element (optional)
-        text (str): Text to locate the element (optional, for Helium)
-        action (str): Action to perform ('click', 'fill', 'press') (default: 'click')
-        input_text (str): Text to input if action='fill' (optional)
-        key (str): Key to press if action='press' (e.g., 'ENTER') (optional)
-    
-    Returns:
-        str: Result of the interaction
-    """
-    try:
-        if selector:
-            element = driver.find_element(By.CSS_SELECTOR, selector)
-        elif text:
-            if action == "click" and Link(text).exists():
-                element = Link(text).web_element
-            else:
-                element = Text(text).web_element
-        else:
-            return "Must provide selector or text"
+class InteractElementTool(Tool):
+    name = "interact_element"
+    description = "Interacts with a web element (click, fill text, press keys)."
+    inputs = {
+        "selector": {"type": "str", "default": None, "description": "CSS selector to target element"},
+        "text": {"type": "str", "default": None, "description": "Text to locate element"},
+        "action": {"type": "str", "default": "click", "description": "Action: 'click', 'fill', or 'press'"},
+        "input_text": {"type": "str", "default": None, "description": "Text to input for fill action"},
+        "key": {"type": "str", "default": None, "description": "Key to press (e.g., 'ENTER')"}
+    }
+    output_type = "str"
+
+    def __init__(self, driver):
+        self.driver = driver
 
-        if action == "click":
-            click(element)
-            return f"Clicked element with {selector or text}"
-        elif action == "fill":
-            if not input_text:
-                return "input_text required for fill action"
-            write(input_text, into=element)
-            return f"Filled {input_text} into element with {selector or text}"
-        elif action == "press":
-            if not key:
-                return "key required for press action"
-            key_map = {"ENTER": Keys.ENTER, "TAB": Keys.TAB, "RETURN": Keys.RETURN}
-            if key.upper() in key_map:
-                element.send_keys(key_map[key.upper()])
+    def forward(self, **kwargs):
+        selector = kwargs.get("selector")
+        text = kwargs.get("text")
+        action = kwargs.get("action", "click")
+        input_text = kwargs.get("input_text")
+        key = kwargs.get("key")
+        try:
+            if selector:
+                element = self.driver.find_element(By.CSS_SELECTOR, selector)
+            elif text:
+                if action == "click" and Link(text).exists():
+                    element = Link(text).web_element
+                else:
+                    element = Text(text).web_element
             else:
-                press(key)
-            return f"Pressed {key} on element with {selector or text}"
-        else:
-            return f"Unknown action: {action}"
-    except Exception as e:
-        return f"Failed to interact with element {selector or text}: {str(e)}"
+                return "Must provide selector or text"
 
-# Register the tool
-interact_element_tool = Tool(
-    name="interact_element",
-    description="Interacts with a web element (click, fill text, press keys).",
-    inputs={
-        "selector": {
-            "type": "str",
-            "default": None,
-            "description": "CSS selector to target element"
-        },
-        "text": {
-            "type": "str",
-            "default": None,
-            "description": "Text to locate element"
-        },
-        "action": {
-            "type": "str",
-            "default": "click",
-            "description": "Action: 'click', 'fill', or 'press'"
-        },
-        "input_text": {
-            "type": "str",
-            "default": None,
-            "description": "Text to input for fill action"
-        },
-        "key": {
-            "type": "str",
-            "default": None,
-            "description": "Key to press (e.g., 'ENTER')"
-        }
-    },
-    output_type="str",
-    function=interact_element
-) 
+            if action == "click":
+                click(element)
+                return f"Clicked element with {selector or text}"
+            elif action == "fill":
+                if not input_text:
+                    return "input_text required for fill action"
+                write(input_text, into=element)
+                return f"Filled {input_text} into element with {selector or text}"
+            elif action == "press":
+                if not key:
+                    return "key required for press action"
+                key_map = {"ENTER": Keys.ENTER, "TAB": Keys.TAB, "RETURN": Keys.RETURN}
+                if key.upper() in key_map:
+                    element.send_keys(key_map[key.upper()])
+                else:
+                    press(key)
+                return f"Pressed {key} on element with {selector or text}"
+            else:
+                return f"Unknown action: {action}"
+        except Exception as e:
+            return f"Failed to interact with element {selector or text}: {str(e)}"