iruno/test_wprm3

1import logging2import os3from typing import Any, List, Tuple4 5from browsergym.core.action.highlevel import HighLevelActionSet6from browsergym.utils.obs import (7    flatten_axtree_to_str,8    flatten_dom_to_str,9    prune_html,10)11from browsergym.experiments import Agent12 13from utils import remove_inline_comments_safe, image_to_jpg_base64_url14 15import openai16 17 18logger = logging.getLogger(__name__)19 20openai.api_key = os.getenv("OPENAI_API_KEY")21 22 23 24class BrowserAgent(Agent):25    def obs_preprocessor(self, obs: dict) -> dict:26        return {27            "chat_messages": obs["chat_messages"],28            "som_screenshot": obs["som_screenshot"],29            "goal_object": obs["goal_object"],30            "last_action": obs["last_action"],31            "last_action_error": obs["last_action_error"],32            "open_pages_urls": obs["open_pages_urls"],33            "open_pages_titles": obs["open_pages_titles"],34            "active_page_index": obs["active_page_index"],35            "axtree_txt": flatten_axtree_to_str(obs["axtree_object"], filter_visible_only=True, extra_properties=obs['extra_element_properties'], filter_som_only=True),36            "pruned_html": prune_html(flatten_dom_to_str(obs["dom_object"])),37        }38 39    def __init__(self, model_name: str = "gpt-4o", use_html: bool = False, use_axtree: bool = True, use_screenshot: bool = False):40        super().__init__()41        logger.info(f"Initializing BrowserAgent with model: {model_name}")42        logger.info(f"Observation space: HTML={use_html}, AXTree={use_axtree}, Screenshot={use_screenshot}")43        44        self.model_name = model_name45        self.use_html = use_html46        self.use_axtree = use_axtree47        self.use_screenshot = use_screenshot48        49        if not (use_html or use_axtree):50            raise ValueError("Either use_html or use_axtree must be set to True.")51        52        self.openai_client = openai.OpenAI()53        54        self.action_set = HighLevelActionSet(55            subsets=["chat", "tab", "nav", "bid", "infeas"],56            strict=False,57            multiaction=False,58            demo_mode="default"59        )60        self.action_history = []61 62    def get_action(self, obs: dict) -> tuple[str, dict]:63        logger.debug("Preparing action request")64        65        system_msgs = [{66            "type": "text",67            "text": """\68# Instructions69 70You are a UI Assistant, your goal is to help the user perform tasks using a web browser. You can71communicate with the user via a chat, to which the user gives you instructions and to which you72can send back messages. You have access to a web browser that both you and the user can see,73and with which only you can interact via specific commands.74 75Review the instructions from the user, the current state of the page and all other information76to find the best possible next action to accomplish your goal. Your answer will be interpreted77and executed by a program, make sure to follow the formatting instructions.78"""79        }]80 81        user_msgs = []82        83        # Add chat messages84        user_msgs.append({85            "type": "text",86            "text": "# Chat Messages\n"87        })88        for msg in obs["chat_messages"]:89            if msg["role"] in ("user", "assistant", "infeasible"):90                user_msgs.append({91                    "type": "text",92                    "text": f"- [{msg['role']}] {msg['message']}\n"93                })94                logger.debug(f"Added chat message: [{msg['role']}] {msg['message']}")95            elif msg["role"] == "user_image":96                user_msgs.append({"type": "image_url", "image_url": msg["message"]})97                logger.debug("Added user image message")98 99        # Add open tabs info100        user_msgs.append({101            "type": "text",102            "text": "# Currently open tabs\n"103        })104        for page_index, (page_url, page_title) in enumerate(105            zip(obs["open_pages_urls"], obs["open_pages_titles"])106        ):107            user_msgs.append({108                "type": "text",109                "text": f"""\110Tab {page_index}{" (active tab)" if page_index == obs["active_page_index"] else ""}111  Title: {page_title}112  URL: {page_url}113"""114            })115            logger.debug(f"Added tab info: {page_title} ({page_url})")116 117        # Add accessibility tree if enabled118        if self.use_axtree:119            user_msgs.append({120                "type": "text",121                "text": f"""\122# Current page Accessibility Tree123 124{obs["axtree_txt"]}125 126"""127            })128            logger.debug("Added accessibility tree")129 130        # Add HTML if enabled131        if self.use_html:132            user_msgs.append({133                "type": "text",134                "text": f"""\135# Current page DOM136 137{obs["pruned_html"]}138 139"""140            })141            logger.debug("Added HTML DOM")142 143        # Add screenshot if enabled144        if self.use_screenshot:145            user_msgs.append({146                "type": "text",147                "text": "# Current page Screenshot\n"148            })149            user_msgs.append({150                "type": "image_url",151                "image_url": {152                    "url": image_to_jpg_base64_url(obs["som_screenshot"]),153                    "detail": "auto"154                }155            })156            logger.debug("Added screenshot")157 158        # Add action space description159        user_msgs.append({160            "type": "text",161            "text": f"""\162# Action Space163 164{self.action_set.describe(with_long_description=False, with_examples=True)}165 166Here are examples of actions with chain-of-thought reasoning:167 168I now need to click on the Submit button to send the form. I will use the click action on the button, which has bid 12.169```click("12")```170 171I found the information requested by the user, I will send it to the chat.172```send_msg_to_user("The price for a 15\\" laptop is 1499 USD.")```173 174"""175        })176 177        # Add action history and errors178        if self.action_history:179            user_msgs.append({180                "type": "text",181                "text": "# History of past actions\n"182            })183            for action in self.action_history:184                user_msgs.append({185                    "type": "text",186                    "text": f"\n{action}\n"187                })188                logger.debug(f"Added past action: {action}")189 190            if obs["last_action_error"]:191                user_msgs.append({192                    "type": "text",193                    "text": f"""\194# Error message from last action195 196{obs["last_action_error"]}197 198"""199                })200                logger.warning(f"Last action error: {obs['last_action_error']}")201 202        # Ask for next action203        user_msgs.append({204            "type": "text",205            "text": """\206# Next action207 208You will now think step by step and produce your next best action. Reflect on your past actions, any resulting error message, and the current state of the page before deciding on your next action.209Note: You might use 'goto' action if you're in a blank page.210"""211        })212 213        # Log the full prompt for debugging214        prompt_text_strings = []215        for message in system_msgs + user_msgs:216            match message["type"]:217                case "text":218                    prompt_text_strings.append(message["text"])219                case "image_url":220                    image_url = message["image_url"]221                    if isinstance(message["image_url"], dict):222                        image_url = image_url["url"]223                    if image_url.startswith("data:image"):224                        prompt_text_strings.append(225                            "image_url: " + image_url[:30] + "... (truncated)"226                        )227                    else:228                        prompt_text_strings.append("image_url: " + image_url)229                case _:230                    raise ValueError(231                        f"Unknown message type {repr(message['type'])} in the task goal."232                    )233        full_prompt_txt = "\n".join(prompt_text_strings)234        logger.debug(full_prompt_txt)235 236        # Query OpenAI model237        logger.info("Sending request to OpenAI")238        response = self.openai_client.chat.completions.create(239            model=self.model_name,240            messages=[241                {"role": "system", "content": system_msgs},242                {"role": "user", "content": user_msgs}243            ],244            n=20,245            temperature=0.8246        )247        parses = []248        for i, choice in enumerate(response.choices):249            response = choice.message.content250            try:251                parses.append({252                    'response': response,253                    'thought': response.split('```')[0].strip(),254                    'action': remove_inline_comments_safe(response.split('```')[1].strip('`').strip().strip('`').strip()),255                })256            except Exception as e:257                logger.error(f"Error parsing action: {e}")258                logger.error(f"Response: {response}")259                logger.error(f"Choice: {choice}")260                logger.error(f"Index: {i}")261                logger.error(f"Response: {response}")262                263        candidates = self.get_top_k_actions(parses)264        logger.info(f"Received action from OpenAI: {[cand['action'] for cand in candidates]}")265        return candidates, {}266 267    def get_top_k_actions(self, parses, k=3):268        count_dict = {}269        action_to_parsed = {}270        for parsed in parses:271            action = parsed["action"]272            if action in count_dict:273                count_dict[action] += 1274            else:275                count_dict[action] = 1276                action_to_parsed[action] = parsed.copy()277        278        # Get the top_k most frequent actions279        sorted_actions = sorted(count_dict.items(), key=lambda x: x[1], reverse=True)280        top_k_actions = [action_to_parsed[action] for action, _ in sorted_actions[:k]]281        282        return top_k_actions