AutoAgent/autoagent/tools/md_obs.py

from browsergym.utils.obs import _process_bid, IGNORED_AXTREE_ROLES, IGNORED_AXTREE_PROPERTIES

def flatten_md_axtree_to_str(
    AX_tree,
    extra_properties: dict = None,
    with_visible: bool = False,
    with_clickable: bool = False,
    with_center_coords: bool = False,
    with_bounding_box_coords: bool = False,
    with_som: bool = False,
    skip_generic: bool = True,
    filter_visible_only: bool = False,
    filter_with_bid_only: bool = False,
    filter_som_only: bool = False,
    coord_decimals: int = 0,
    ignored_roles=IGNORED_AXTREE_ROLES,
    ignored_properties=IGNORED_AXTREE_PROPERTIES,
    remove_redundant_static_text: bool = True,
    hide_bid_if_invisible: bool = False,
    hide_all_children: bool = False,
    hide_all_bids: bool = False,
) -> str:
    """Formats the accessibility tree into a string text"""
    node_id_to_idx = {}
    for idx, node in enumerate(AX_tree["nodes"]):
        node_id_to_idx[node["nodeId"]] = idx
    # print("******************************* AX_tree:", AX_tree)

    def dfs(node_idx: int, depth: int, parent_node_filtered: bool, parent_node_name: str) -> str:
        tree_str = ""
        node = AX_tree["nodes"][node_idx]
        # indent = "\t" * depth
        indent = ""
        skip_node = False  # node will not be printed, with no effect on children nodes
        filter_node = False  # node will not be printed, possibly along with its children nodes
        node_role = node["role"]["value"]
        node_name = ""

        if node_role in ignored_roles:
            skip_node = True
            pass
        elif "name" not in node:
            skip_node = True
            pass
        else:
            node_name = node["name"]["value"]
            if "value" in node and "value" in node["value"]:
                node_value = node["value"]["value"]
            else:
                node_value = None

            # extract bid
            bid = node.get("browsergym_id", None)

            # extract node attributes
            attributes = []
            for property in node.get("properties", []):
                if not "value" in property:
                    continue
                if not "value" in property["value"]:
                    continue

                prop_name = property["name"]
                prop_value = property["value"]["value"]

                if prop_name in ignored_properties:
                    continue
                elif prop_name in ("required", "focused", "atomic"):
                    if prop_value:
                        attributes.append(prop_name)
                else:
                    attributes.append(f"{prop_name}={repr(prop_value)}")

            if skip_generic and node_role == "generic" and not attributes:
                skip_node = True

            if hide_all_children and parent_node_filtered:
                skip_node = True

            if node_role == "StaticText":
                if parent_node_filtered:
                    skip_node = True
                elif remove_redundant_static_text and node_name in parent_node_name:
                    skip_node = True
            else:
                filter_node, extra_attributes_to_print = _process_bid(
                    bid,
                    extra_properties=extra_properties,
                    with_visible=with_visible,
                    with_clickable=with_clickable,
                    with_center_coords=with_center_coords,
                    with_bounding_box_coords=with_bounding_box_coords,
                    with_som=with_som,
                    filter_visible_only=filter_visible_only,
                    filter_with_bid_only=filter_with_bid_only,
                    filter_som_only=filter_som_only,
                    coord_decimals=coord_decimals,
                )

                # if either is True, skip the node
                skip_node = skip_node or filter_node

                # insert extra attributes before regular attributes
                attributes = extra_attributes_to_print + attributes

            # actually print the node string
            if not skip_node:
                if node_role == "generic" and not node_name:
                    node_str = f"{node_role}"
                else:
                    node_str = f"{node_name.strip()}"

                if not (
                    hide_all_bids
                    or bid is None
                    or (
                        hide_bid_if_invisible
                        and extra_properties.get(bid, {}).get("visibility", 0) < 0.5
                    )
                ):
                    node_str = node_str

                if node_value is not None:
                    node_str += f' value={repr(node["value"]["value"])}'

                if attributes:
                    node_str += ", ".join([""] + attributes)

                tree_str += f"{indent}{node_str}"

        for child_node_id in node["childIds"]:
            if child_node_id not in node_id_to_idx or child_node_id == node["nodeId"]:
                continue
            # mark this to save some tokens
            child_depth = depth if skip_node else (depth + 1)
            child_str = dfs(
                node_id_to_idx[child_node_id],
                child_depth,
                parent_node_filtered=filter_node,
                parent_node_name=node_name,
            )
            if child_str:
                if tree_str:
                    tree_str += "\n"
                tree_str += child_str

        return tree_str

    tree_str = dfs(0, 1, False, "")
    return tree_str

# def flatten_md_axtree_to_str(
#     AX_tree,
#     extra_properties: dict = None,
#     with_visible: bool = False,
#     with_clickable: bool = False,
#     with_center_coords: bool = False,
#     with_bounding_box_coords: bool = False,
#     with_som: bool = False,
#     skip_generic: bool = True,
#     filter_visible_only: bool = False,
#     filter_with_bid_only: bool = False,
#     filter_som_only: bool = False,
#     coord_decimals: int = 0,
#     ignored_roles=IGNORED_AXTREE_ROLES,
#     ignored_properties=IGNORED_AXTREE_PROPERTIES,
#     remove_redundant_static_text: bool = True,
#     hide_bid_if_invisible: bool = False,
#     hide_all_children: bool = False,
#     hide_all_bids: bool = False,
# ) -> str:
#     """Formats the accessibility tree into a markdown string"""
#     if not AX_tree or "nodes" not in AX_tree or not AX_tree["nodes"]:
#         return "No content available"
        
#     node_id_to_idx = {}
#     for idx, node in enumerate(AX_tree["nodes"]):
#         node_id_to_idx[node["nodeId"]] = idx

#     def dfs(node_idx: int, depth: int, parent_node_filtered: bool, parent_node_name: str) -> str:
#         tree_str = ""
#         node = AX_tree["nodes"][node_idx]
        
#         # 处理节点内容
#         if "name" in node and "value" in node["name"]:
#             content = node["name"]["value"].strip()
#             if content:
#                 # 移除多余的空白字符
#                 content = ' '.join(content.split())
#                 # 添加到输出
#                 tree_str += content + "\n"
        
#         # 递归处理子节点
#         for child_node_id in node.get("childIds", []):
#             if child_node_id in node_id_to_idx and child_node_id != node["nodeId"]:
#                 child_str = dfs(
#                     node_id_to_idx[child_node_id],
#                     depth + 1,
#                     parent_node_filtered=False,
#                     parent_node_name=node.get("name", {}).get("value", "")
#                 )
#                 if child_str:
#                     tree_str += child_str

#         return tree_str

#     tree_str = dfs(0, 0, False, "")
#     return tree_str.strip()