fix: describe custom actions in HighLevelActionSet.to_tool_description

browsergym/core/src/browsergym/core/action/highlevel.py

@@ -372,6 +372,7 @@ def __init__(
 
         # parse the actions and build the action space
         self.action_set: dict[str, HighLevelAction] = {}
+        self._action_functions: dict[str, callable] = {}
         self.python_includes = ""
 
         # include playwright imports
@@ -427,6 +428,7 @@ def __init__(
             if func.__name__ in self.action_set:
                 raise ValueError(f"Duplicated action '{func.__name__}'")
 
+            self._action_functions[func.__name__] = func
             self.action_set[func.__name__] = HighLevelAction(
                 # entrypoint=func,
                 signature=signature,
@@ -574,7 +576,9 @@ def to_tool_description(self, api="openai", add_examples=True) -> list[dict]:
         for tool_name, action in self.action_set.items():
             # Parse the signature to extract parameter names and types
             parameters = {"type": "object", "properties": {}, "required": []}
-            signature = inspect.signature(globals()[tool_name])
+            # use the function registered in this action set, so that custom actions
+            # (including ones that reuse a built-in action name) are described correctly
+            signature = inspect.signature(self._action_functions[tool_name])
             for param_name, param in signature.parameters.items():
                 param_type = "string"  # Default to string if type is not specified
                 if param.annotation != inspect.Parameter.empty:

tests/core/test_actions_highlevel.py

@@ -110,6 +110,52 @@ def test_action_parser():
     ]
 
 
+def _custom_click_twice(bid: str, times: int = 2):
+    """
+    Clicks an element several times.
+
+    Examples:
+        _custom_click_twice('a51', 3)
+    """
+    pass
+
+
+def click(bid: str, note: str = ""):
+    """
+    A custom action that reuses the name of a built-in action.
+
+    Examples:
+        click('a51', 'submit')
+    """
+    pass
+
+
+def test_tool_description_custom_actions():
+    action_set = HighLevelActionSet(subsets=["custom"], custom_actions=[_custom_click_twice])
+    tools = {tool["name"]: tool for tool in action_set.to_tool_description()}
+    assert set(tools) == {"noop", "_custom_click_twice"}
+    assert tools["_custom_click_twice"]["parameters"] == {
+        "type": "object",
+        "properties": {
+            "bid": {"type": "string"},
+            "times": {"type": "number", "default": 2},
+        },
+        "required": ["bid"],
+    }
+
+    # a custom action named like a built-in one must be described with its own signature
+    action_set = HighLevelActionSet(subsets=["custom"], custom_actions=[click])
+    tools = {tool["name"]: tool for tool in action_set.to_tool_description(api="anthropic")}
+    assert tools["click"]["input_schema"] == {
+        "type": "object",
+        "properties": {
+            "bid": {"type": "string"},
+            "note": {"type": "string", "default": ""},
+        },
+        "required": ["bid"],
+    }
+
+
 def test_valid_action():
     action_set = HighLevelActionSet()