                                "text": """\
SCREENSHOT:
""",
                            },
                            {
                                "type": "image_url",
                                "image_url": {"url": image_data_to_uri(screenshot_data)},
                            },
                            {
                                "type": "text",
                                "text": f"""\
ACTION: {action}
""",
                            },
                        ]
                    )
                else:
                    raise ValueError("Unexpected format for few-shot example.")

    @cost_tracker_decorator
    def get_action(self, obs: Any) -> tuple[str, dict]:
        """
        Replica of VisualWebArena agent
        https://github.com/web-arena-x/visualwebarena/blob/89f5af29305c3d1e9f97ce4421462060a70c9a03/agent/prompts/prompt_constructor.py#L211
        https://github.com/web-arena-x/visualwebarena/blob/89f5af29305c3d1e9f97ce4421462060a70c9a03/agent/prompts/prompt_constructor.py#L272

        Args:
            obs (Any): Observation from the environment

        Returns:
            tuple[str, dict]: Action and AgentInfo
        """
        user_messages = []

        # 1. add few-shot examples (if any)
        user_messages.extend(self.few_shot_messages)

        # 2. add the current observation to the user prompt
        active_tab = obs["active_page_index"][0]
        open_tab_titles = obs["open_pages_titles"]
        cur_tabs_txt = " | ".join(
            f"Tab {i}{' (current)' if i == active_tab else ''}: {title}"
            for i, title in enumerate(open_tab_titles)
        )
        cur_axtree_txt = obs["axtree_txt"]
        cur_url = obs["url"]
        user_messages.append(
            {
                "type": "text",
                "text": f"""\
OBSERVATION:

{cur_tabs_txt}

{cur_axtree_txt}

URL: {cur_url}

PREVIOUS ACTION: {self.action_history[-1]}
""",
            }
        )

        # if desired, add current page's screenshot
        if self.observation_type in ("axtree_som", "axtree_screenshot"):
            cur_screenshot = obs["screenshot"]
            # if desired, overlay set-of-marks on the screenshot
            if self.observation_type == "axtree_som":
                cur_screenshot = overlay_som(cur_screenshot, obs["extra_element_properties"])
            user_messages.extend(
                [
                    {
                        "type": "text",
                        "text": """\
SCREENSHOT:
""",
                    },
                    {"type": "image_url", "image_url": {"url": image_data_to_uri(cur_screenshot)}},
                ]
            )

        # 3. add the objective (goal) to the user prompt
        user_messages.append(
            {
                "type": "text",
                "text": f"""\
OBJECTIVE:
""",
            }
        )
        user_messages.extend(obs["goal_object"])

        messages = [
            # intro prompt
            make_system_message(content=self.intro_messages),
            # few-shot examples + observation + goal
            make_user_message(content=user_messages),
        ]

        # finally, query the chat model
        answer: dict = retry(self.chat_model, messages, n_retry=3, parser=parser)

        action = answer.get("action", None)
        thought = answer.get("think", None)

        self.action_history.append(action)

        return (
            action,
            AgentInfo(
                think=thought,
                chat_messages=messages,
            ),
        )

# A WebArena agent is a VisualWebArena agent with only axtree observation
WebArenaAgent = partial(
    VisualWebArenaAgentArgs,
    agent_name="WebArenaAgent",
    observation_type="axtree",
)

WA_AGENT_4O_MINI = WebArenaAgent(
    temperature=0.1,
    chat_model_args=CHAT_MODEL_ARGS_DICT["openai/gpt-4o-mini-2024-07-18"],
)

WA_AGENT_4O = WebArenaAgent(
    temperature=0.1,
    chat_model_args=CHAT_MODEL_ARGS_DICT["openai/gpt-4o-mini-2024-07-18"],
)

WA_AGENT_SONNET = WebArenaAgent(


# Best-attempt reproduction of original prompts the VisualWebArena agent

TEMPLATES = {}

# https://github.com/web-arena-x/visualwebarena/blob/89f5af29305c3d1e9f97ce4421462060a70c9a03/agent/prompts/raw/p_cot_id_actree_3s.py#L1
TEMPLATES["axtree"] = {
    "intro": """\
You are an autonomous intelligent agent tasked with navigating a web browser. You will be given web-based tasks. These tasks will be accomplished through the use of specific actions you can issue.

Here's the information you'll have:
The user's objective: This is the task you're trying to complete.
The current web page's accessibility tree: This is a simplified representation of the webpage, providing key information.
The current web page's URL: This is the page you're currently navigating.
The open tabs: These are the tabs you have open.
The previous action: This is the action you just performed. It may be helpful to track your progress.

{action_space_description}

To be successful, it is very important to follow the following rules:
1. You should only issue an action that is valid given the current observation
2. You should only issue one action at a time.
3. You should follow the examples to reason step by step and then issue the next action.
4. Generate the action in the correct format. Start with a "In summary, the next action I will perform is" phrase, followed by action inside ``````. For example, "In summary, the next action I will perform is ```click("1234")```".""",
    "examples": [
        (
            """\
OBSERVATION:
[1744] link 'HP CB782A#ABA 640 Inkjet Fax Machine (Renewed)'
[1749] StaticText '$279.49'
[1757] button 'Add to Cart'
[1760] button 'Add to Wish List'
[1761] button 'Add to Compare'
URL: http://onestopmarket.com/office-products/office-electronics.html
OBJECTIVE: What is the price of HP Inkjet Fax Machine?
PREVIOUS ACTION: None""",
            """\
Let's think step-by-step. This page list the information of HP Inkjet Fax Machine, which is the product identified in the objective. Its price is $279.49. I think I have achieved the objective. I will issue a user message with the answer. In summary, the next action I will perform is ```send_msg_to_user("$279.49")```""",
        ),
        (
            """\
OBSERVATION:
[204] heading '/f/food'
[593] heading '[homemade] Obligatory Halloween Pumpkin Loaf!'
	[942] link '[homemade] Obligatory Halloween Pumpkin Loaf!'
[945] StaticText 'Submitted by '
[30] link 'kneechalice' expanded: False
[1484] StaticText 't3_yid9lu'
[949] time 'October 31, 2022 at 10:10:03 AM EDT'
	[1488] StaticText '1 year ago'
[1489] link '45 comments'
[605] heading '[I ate] Maple Pecan Croissant'
	[963] link '[I ate] Maple Pecan Croissant'
[966] StaticText 'Submitted by '
[37] link 'AccordingtoJP' expanded: False
[1494] StaticText 't3_y3hrpn'
[970] time 'October 13, 2022 at 10:41:09 PM EDT'
	[1498] StaticText '1 year ago'
[1499] link '204 comments'
URL: http://reddit.com
OBJECTIVE: Tell me what the top comment on the croissant post says.
PREVIOUS ACTION: None""",
            """\
Let's think step-by-step. This page has a post titled '[I ate] Maple Pecan Croissant', which is the post mentioned in the objective. In order to find the top comment, I will navigate into the comments section of the post. In summary, the next action I will perform is ```click("1499")```""",
        ),
        (
            """\
OBSERVATION:
[42] link 'My account'
[43] link 'Logout'
[44] link 'Publish Ad'
[25] heading 'What are you looking for today?'
[143] StaticText 'Keyword'
[81] textbox 'e.g., a blue used car' required: False
[146] StaticText 'Category'
[28] heading 'Latest Listings'
[86] link 'Atlas Powered Audio System w/ Tripod'
	[176] img 'Atlas Powered Audio System w/ Tripod'
[511] StaticText '150.00 $'
[88] link 'Neptune Gaming Console'
	[178] img 'Neptune Gaming Console'
[515] StaticText '350.00 $'
URL: http://classifieds.com
OBJECTIVE: Help me find the cheapest dark colored guitar.
PREVIOUS ACTION: None""",
            """\
Let's think step-by-step. The objective is to find the cheapest dark colored guitar on the site. The site has a search box whose ID is [81]. I can search for guitars by entering "guitar". I can submit this by pressing the Enter afterwards. In summary, the next action I will perform is ```fill("81", "guitar")```""",
        ),
    ],
}

# https://github.com/web-arena-x/visualwebarena/blob/89f5af29305c3d1e9f97ce4421462060a70c9a03/agent/prompts/raw/p_multimodal_cot_id_actree_3s.py#L1
TEMPLATES["axtree_screenshot"] = {
    "intro": """\
You are an autonomous intelligent agent tasked with navigating a web browser. You will be given web-based tasks. These tasks will be accomplished through the use of specific actions you can issue.

Here's the information you'll have:
The user's objective: This is the task you're trying to complete.
The current web page's accessibility tree: This is a simplified representation of the webpage, providing key information.
The current web page's URL: This is the page you're currently navigating.
The open tabs: These are the tabs you have open.
The previous action: This is the action you just performed. It may be helpful to track your progress.

{action_space_description}

To be successful, it is very important to follow the following rules:
1. You should only issue an action that is valid given the current observation
2. You should only issue one action at a time.
3. You should follow the examples to reason step by step and then issue the next action.
4. Generate the action in the correct format. Start with a "In summary, the next action I will perform is" phrase, followed by action inside ``````. For example, "In summary, the next action I will perform is ```click("1234")```".""",
    "examples": [
        (
            """\
OBSERVATION:
[1744] link 'HP CB782A#ABA 640 Inkjet Fax Machine (Renewed)'
[1749] StaticText '$279.49'
[1757] button 'Add to Cart'
[1760] button 'Add to Wish List'
[1761] button 'Add to Compare'
URL: http://onestopmarket.com/office-products/office-electronics.html
OBJECTIVE: What is the price of HP Inkjet Fax Machine?
PREVIOUS ACTION: None""",
            """\
Let's think step-by-step. This page list the information of HP Inkjet Fax Machine, which is the product identified in the objective. Its price is $279.49. I think I have achieved the objective. I will issue a user message with the answer. In summary, the next action I will perform is ```send_msg_to_user("$279.49")```""",
            "multimodal_example1.png",
        ),
        (
            """\
OBSERVATION:
[204] heading '/f/food'
[593] heading '[homemade] Obligatory Halloween Pumpkin Loaf!'
	[942] link '[homemade] Obligatory Halloween Pumpkin Loaf!'
[945] StaticText 'Submitted by '
[30] link 'kneechalice' expanded: False
[1484] StaticText 't3_yid9lu'
[949] time 'October 31, 2022 at 10:10:03 AM EDT'
	[1488] StaticText '1 year ago'
[1489] link '45 comments'
[605] heading '[I ate] Maple Pecan Croissant'
	[963] link '[I ate] Maple Pecan Croissant'
[966] StaticText 'Submitted by '
[37] link 'AccordingtoJP' expanded: False
[1494] StaticText 't3_y3hrpn'
[970] time 'October 13, 2022 at 10:41:09 PM EDT'
	[1498] StaticText '1 year ago'
[1499] link '204 comments'
URL: http://reddit.com
OBJECTIVE: Tell me what the top comment on the croissant post says.
PREVIOUS ACTION: None""",
            """\
Let's think step-by-step. This page has a post titled '[I ate] Maple Pecan Croissant', which is the post mentioned in the objective. In order to find the top comment, I will navigate into the comments section of the post. In summary, the next action I will perform is ```click("1499")```""",
            "multimodal_example2.png",
        ),
        (
            """\
OBSERVATION:
[42] link 'My account'
[43] link 'Logout'
[44] link 'Publish Ad'
[25] heading 'What are you looking for today?'
[143] StaticText 'Keyword'
[81] textbox 'e.g., a blue used car' required: False
[146] StaticText 'Category'
[28] heading 'Latest Listings'
[86] link 'Atlas Powered Audio System w/ Tripod'
	[176] img 'Atlas Powered Audio System w/ Tripod'
[511] StaticText '150.00 $'
[88] link 'Neptune Gaming Console'
	[178] img 'Neptune Gaming Console'
[515] StaticText '350.00 $'
URL: http://classifieds.com
OBJECTIVE: Help me find the cheapest dark colored guitar.
PREVIOUS ACTION: None""",
            """\
Let's think step-by-step. The objective is to find the cheapest dark colored guitar on the site. The site has a search box whose ID is [81]. I can search for guitars by entering "guitar". I can submit this by pressing the Enter afterwards. In summary, the next action I will perform is ```fill("81", "guitar")```""",
            "multimodal_example3.png",
        ),
    ],
}

# https://github.com/web-arena-x/visualwebarena/blob/89f5af29305c3d1e9f97ce4421462060a70c9a03/agent/prompts/raw/p_som_cot_id_actree_3s.py#L1
TEMPLATES["axtree_som"] = prompt = {
    "intro": """\
You are an autonomous intelligent agent tasked with navigating a web browser. You will be given web-based tasks. These tasks will be accomplished through the use of specific actions you can issue.

Here's the information you'll have:
The user's objective: This is the task you're trying to complete.
The current web page screenshot: This is a screenshot of the webpage, with each interactable element assigned a unique numerical id. Each bounding box and its respective id shares the same color.
The observation, which lists the IDs of all interactable elements on the current web page with their text content if any, in the format [id] [tagType] [text content]. tagType is the type of the element, such as button, link, or textbox. text content is the text content of the element. For example, [1234] [button] ['Add to Cart'] means that there is a button with id 1234 and text content 'Add to Cart' on the current web page. [] [StaticText] [text] means that the element is of some text that is not interactable.
The current web page's URL: This is the page you're currently navigating.
The open tabs: These are the tabs you have open.
The previous action: This is the action you just performed. It may be helpful to track your progress.

{action_space_description}

To be successful, it is very important to follow the following rules:
1. You should only issue an action that is valid given the current observation
2. You should only issue one action at a time.
3. You should follow the examples to reason step by step and then issue the next action.
4. Generate the action in the correct format. Start with a "In summary, the next action I will perform is" phrase, followed by action inside ``````. For example, "In summary, the next action I will perform is ```click("1234")```".""",
    "examples": [
        (
            """\
OBSERVATION:
[31] [IMG] [Image, description: hp fx-7010dn fax machine, url: http://ec2-3-13-232-171.us-east-2.compute.amazonaws.com:7770/media/catalog/product/cache/89ff578b9cd87e0600daac45c9e1ea98/B/0/B08GKZ3ZKD.0.jpg]
[32] [A] [HP CB782A#ABA 640 Inkjet Fax Machine (Renewed)]
[] [StaticText] [$279.49]
[33] [BUTTON] [Add to Cart]
[34] [A] [Add to Wish List]
[35] [A] [Add to Compare]
URL: http://onestopmarket.com/office-products/office-electronics.html
OBJECTIVE: What is the price of HP Inkjet Fax Machine?
PREVIOUS ACTION: None""",
            """\
Let's think step-by-step. This page list the information of HP Inkjet Fax Machine, which is the product identified in the objective. Its price is $279.49. I think I have achieved the objective. I will issue a user message with the answer. In summary, the next action I will perform is ```send_msg_to_user("$279.49")```""",
            "som_example1.png",
        ),
        (
            """\
OBSERVATION:
[] [StaticText] [/f/food]
[] [StaticText] [[homemade] Obligatory Halloween Pumpkin Loaf!	Submitted by	kneechalice	t3_yid9lu	1 year ago]
[9] [IMG] []
[] [StaticText] [Submitted by	kneechalice	t3_yid9lu	1 year ago]
[10] [A] [kneechalice]
[11] [A] [45 comments]
[] [StaticText] [[I ate] Maple Pecan Croissant	Submitted by	AccordingtoJP	t3_y3hrpn	1 year ago]
[14] [IMG] []
[] [StaticText] [Submitted by	AccordingtoJP	t3_y3hrpn	1 year ago]
[15] [A] [AccordingtoJP]
[16] [A] [204 comments]
URL: http://reddit.com
OBJECTIVE: Tell me what the top comment on the croissant post says.
PREVIOUS ACTION: None""",
            """\
Let's think step-by-step. This page has a post titled '[I ate] Maple Pecan Croissant', which is the post mentioned in the objective. In order to find the top comment, I will navigate into the comments section of the post. In summary, the next action I will perform is ```click("11")```""",
            "som_example2.png",
        ),
        (
            """\
OBSERVATION:
[] [StaticText] [What are you looking for today?]
[5] [INPUT] []
[6] [SELECT] [Select a category]
[7] [BUTTON] [Search]
[] [StaticText] [Latest Listings]
[] [StaticText] [Atlas Powered Audio System w/ Tripod	150.00 $	Music instruments	Borough of Red Lion  (Pennsylvania)	2023/11/16]
[8] [IMG] [Atlas Powered Audio System w/ Tripod]
[9] [A] [Atlas Powered Audio System w/ Tripod]
[] [StaticText] [150.00 $]
[] [StaticText] [Neptune Gaming Console	350.00 $	Video gaming	Pennwyn  (Pennsylvania)	2023/11/16]
[10] [IMG] [Neptune Gaming Console]
[11] [A] [Neptune Gaming Console]
[] [StaticText] [350.00 $]
URL: http://classifieds.com
OBJECTIVE: Help me find the cheapest dark colored guitar.
PREVIOUS ACTION: None""",
            """\
Let's think step-by-step. The objective is to find the cheapest dark colored guitar on the site. The site has a search box whose ID is [5]. I can search for guitars by entering "guitar". I can submit this by pressing the Enter afterwards. In summary, the next action I will perform is ```fill("5", "guitar")```""",
            "som_example3.png",
        ),
    ],
}


import logging
from dataclasses import asdict, dataclass
from typing import TYPE_CHECKING, Any

import bgym

from agentlab.agents.agent_args import AgentArgs
from agentlab.experiments.loop import ExpArgs
from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT
from agentlab.llm.llm_utils import (
    Discussion,
    HumanMessage,
    ParseError,
    SystemMessage,
    extract_code_blocks,
    retry,
)
from agentlab.llm.tracking import cost_tracker_decorator

if TYPE_CHECKING:
    from agentlab.llm.chat_api import BaseModelArgs

@dataclass
class MostBasicAgentArgs(AgentArgs):
    agent_name: str = "BasicAgent"
    temperature: float = 0.1
    use_chain_of_thought: bool = False
    chat_model_args: "BaseModelArgs" = None

    def make_agent(self) -> bgym.Agent:
        return MostBasicAgent(
            temperature=self.temperature,
            use_chain_of_thought=self.use_chain_of_thought,
            chat_model_args=self.chat_model_args,
        )

    def set_reproducibility_mode(self):
        self.temperature = 0

    def prepare(self):
        return self.chat_model_args.prepare_server()

    def close(self):
        return self.chat_model_args.close_server()

class MostBasicAgent(bgym.Agent):
    def __init__(
        self, temperature: float, use_chain_of_thought: bool, chat_model_args: "BaseModelArgs"
    ):
        self.temperature = temperature
        self.use_chain_of_thought = use_chain_of_thought
        self.chat = chat_model_args.make_model()
        self.chat_model_args = chat_model_args

        self.action_set = bgym.HighLevelActionSet(["bid"], multiaction=False)

    @cost_tracker_decorator
    def get_action(self, obs: Any) -> tuple[str, dict]:
        messages = Discussion(SystemMessage("You are a web assistant."))
        messages.append(
            HumanMessage(
                f"""
You are helping a user to accomplish the following goal on a website:

{obs["goal"]}

To do so, you can interact with the environment using the following actions:

{self.action_set.describe(with_long_description=False)}

The inputs to those functions are the bids given in the html.

Here is the current state of the website, in the form of an html:

{obs["pruned_html"]}

The action you provide must be in between triple ticks and leverage the 'bid=' information provided in the html.
Here is an example of how to use the bid action:

```
click('a314')
```

Please provide a single action at a time and wait for the next observation. Provide only a single action per step.
Focus on the bid that are given in the html, and use them to perform the actions.
"""
            )
        )
        if self.use_chain_of_thought:
            messages.add_text(
                f"""
Provide a chain of thoughts reasoning to decompose the task into smaller steps. And execute only the next step.
"""
            )

        def parser(response: str) -> tuple[dict, bool, str]:
            blocks = extract_code_blocks(response)
            if len(blocks) == 0:
                raise ParseError("No code block found in the response")
            action = blocks[0][1]
            thought = response
            return {"action": action, "think": thought}

        ans_dict = retry(self.chat, messages, n_retry=3, parser=parser)

        action = ans_dict.get("action", None)
        thought = ans_dict.get("think", None)

        return (
            action,
            bgym.AgentInfo(
                think=thought,
                chat_messages=messages,
                # put any stats that you care about as long as it is a number or a dict of numbers
                stats={"prompt_length": len(messages), "response_length": len(thought)},
                markdown_page="Add any txt information here, including base 64 images, to display in xray",
                extra_info={"chat_model_args": asdict(self.chat_model_args)},
            ),
        )

# example for a single task
env_args = bgym.EnvArgs(
    task_name="miniwob.click-button",
    task_seed=0,
    max_steps=10,
    headless=True,
)

chat_model_args = CHAT_MODEL_ARGS_DICT["openai/gpt-4o-mini-2024-07-18"]

# example for 2 experiments testing chain of thoughts on a miniwob task
exp_args = [
    ExpArgs(
        agent_args=MostBasicAgentArgs(
            temperature=0.1,
            use_chain_of_thought=True,
            chat_model_args=chat_model_args,
        ),
        env_args=env_args,
        logging_level=logging.INFO,
    ),
    ExpArgs(
        agent_args=MostBasicAgentArgs(
            temperature=0.1,
            use_chain_of_thought=False,
            chat_model_args=chat_model_args,
        ),
        env_args=env_args,
        logging_level=logging.INFO,
    ),
]

AGENT_4o_MINI = MostBasicAgentArgs(
    temperature=0.3,
    use_chain_of_thought=True,
    chat_model_args=chat_model_args,
)

def experiment_config():
    return exp_args


MostBasicAgentArgs

import bgym

from agentlab.agents.agent_args import AgentArgs
from agentlab.experiments.loop import ExpArgs
from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT
from agentlab.llm.llm_utils import (
    Discussion,
    HumanMessage,
    ParseError,
    SystemMessage,
    extract_code_blocks,
    retry,
)
from agentlab.llm.tracking import cost_tracker_decorator

if TYPE_CHECKING:
    from agentlab.llm.chat_api import BaseModelArgs

@dataclass
class MostBasicAgentArgs(AgentArgs):
    agent_name: str = "BasicAgent"
    temperature: float = 0.1
    use_chain_of_thought: bool = False
    chat_model_args: "BaseModelArgs" = None

    def make_agent(self) -> bgym.Agent:
        return MostBasicAgent(
            temperature=self.temperature,
            use_chain_of_thought=self.use_chain_of_thought,
            chat_model_args=self.chat_model_args,
        )

    def set_reproducibility_mode(self):
        self.temperature = 0

    def prepare(self):
        return self.chat_model_args.prepare_server()

    def close(self):
        return self.chat_model_args.close_server()

class MostBasicAgent(bgym.Agent):
    def __init__(
        self, temperature: float, use_chain_of_thought: bool, chat_model_args: "BaseModelArgs"
    ):
        self.temperature = temperature
        self.use_chain_of_thought = use_chain_of_thought
        self.chat = chat_model_args.make_model()
        self.chat_model_args = chat_model_args

        self.action_set = bgym.HighLevelActionSet(["bid"], multiaction=False)

    @cost_tracker_decorator
    def get_action(self, obs: Any) -> tuple[str, dict]:
        messages = Discussion(SystemMessage("You are a web assistant."))
        messages.append(
            HumanMessage(
                f"""
You are helping a user to accomplish the following goal on a website:


MostBasicAgent

    use_chain_of_thought: bool = False
    chat_model_args: "BaseModelArgs" = None

    def make_agent(self) -> bgym.Agent:
        return MostBasicAgent(
            temperature=self.temperature,
            use_chain_of_thought=self.use_chain_of_thought,
            chat_model_args=self.chat_model_args,
        )

    def set_reproducibility_mode(self):
        self.temperature = 0

    def prepare(self):
        return self.chat_model_args.prepare_server()

    def close(self):
        return self.chat_model_args.close_server()

class MostBasicAgent(bgym.Agent):
    def __init__(
        self, temperature: float, use_chain_of_thought: bool, chat_model_args: "BaseModelArgs"
    ):
        self.temperature = temperature
        self.use_chain_of_thought = use_chain_of_thought
        self.chat = chat_model_args.make_model()
        self.chat_model_args = chat_model_args

        self.action_set = bgym.HighLevelActionSet(["bid"], multiaction=False)

    @cost_tracker_decorator
    def get_action(self, obs: Any) -> tuple[str, dict]:
        messages = Discussion(SystemMessage("You are a web assistant."))
        messages.append(
            HumanMessage(
                f"""
You are helping a user to accomplish the following goal on a website:

{obs["goal"]}

To do so, you can interact with the environment using the following actions:

{self.action_set.describe(with_long_description=False)}

The inputs to those functions are the bids given in the html.

Here is the current state of the website, in the form of an html:

{obs["pruned_html"]}

The action you provide must be in between triple ticks and leverage the 'bid=' information provided in the html.
Here is an example of how to use the bid action:

```
click('a314')
```

Please provide a single action at a time and wait for the next observation. Provide only a single action per step.
Focus on the bid that are given in the html, and use them to perform the actions.
"""
            )
        )
        if self.use_chain_of_thought:
            messages.add_text(
                f"""
Provide a chain of thoughts reasoning to decompose the task into smaller steps. And execute only the next step.
"""
            )

        def parser(response: str) -> tuple[dict, bool, str]:
            blocks = extract_code_blocks(response)
            if len(blocks) == 0:
                raise ParseError("No code block found in the response")
            action = blocks[0][1]
            thought = response
            return {"action": action, "think": thought}

        ans_dict = retry(self.chat, messages, n_retry=3, parser=parser)

        action = ans_dict.get("action", None)
        thought = ans_dict.get("think", None)

        return (
            action,
            bgym.AgentInfo(
                think=thought,
                chat_messages=messages,
                # put any stats that you care about as long as it is a number or a dict of numbers
                stats={"prompt_length": len(messages), "response_length": len(thought)},
                markdown_page="Add any txt information here, including base 64 images, to display in xray",
                extra_info={"chat_model_args": asdict(self.chat_model_args)},
            ),
        )

# example for a single task
env_args = bgym.EnvArgs(
    task_name="miniwob.click-button",
    task_seed=0,
    max_steps=10,
    headless=True,
)

chat_model_args = CHAT_MODEL_ARGS_DICT["openai/gpt-4o-mini-2024-07-18"]

# example for 2 experiments testing chain of thoughts on a miniwob task
exp_args = [
    ExpArgs(
        agent_args=MostBasicAgentArgs(
            temperature=0.1,
            use_chain_of_thought=True,
            chat_model_args=chat_model_args,
        ),


        logging_level=logging.INFO,
    ),
    ExpArgs(
        agent_args=MostBasicAgentArgs(
            temperature=0.1,
            use_chain_of_thought=False,
            chat_model_args=chat_model_args,
        ),
        env_args=env_args,
        logging_level=logging.INFO,
    ),
]

AGENT_4o_MINI = MostBasicAgentArgs(
    temperature=0.3,
    use_chain_of_thought=True,
    chat_model_args=chat_model_args,
)

def experiment_config():
    return exp_args


    Discussion,
    HumanMessage,
    ParseError,
    SystemMessage,
    extract_code_blocks,
    retry,
)
from agentlab.llm.tracking import cost_tracker_decorator

if TYPE_CHECKING:
    from agentlab.llm.chat_api import BaseModelArgs

@dataclass
class MostBasicAgentArgs(AgentArgs):
    agent_name: str = "BasicAgent"
    temperature: float = 0.1
    use_chain_of_thought: bool = False
    chat_model_args: "BaseModelArgs" = None

    def make_agent(self) -> bgym.Agent:
        return MostBasicAgent(
            temperature=self.temperature,
            use_chain_of_thought=self.use_chain_of_thought,
            chat_model_args=self.chat_model_args,
        )

    def set_reproducibility_mode(self):
        self.temperature = 0

    def prepare(self):
        return self.chat_model_args.prepare_server()

    def close(self):
        return self.chat_model_args.close_server()

class MostBasicAgent(bgym.Agent):
    def __init__(
        self, temperature: float, use_chain_of_thought: bool, chat_model_args: "BaseModelArgs"
    ):
        self.temperature = temperature
        self.use_chain_of_thought = use_chain_of_thought
        self.chat = chat_model_args.make_model()
        self.chat_model_args = chat_model_args


from agentlab.llm.tracking import cost_tracker_decorator

if TYPE_CHECKING:
    from agentlab.llm.chat_api import BaseModelArgs

@dataclass
class MostBasicAgentArgs(AgentArgs):
    agent_name: str = "BasicAgent"
    temperature: float = 0.1
    use_chain_of_thought: bool = False
    chat_model_args: "BaseModelArgs" = None

    def make_agent(self) -> bgym.Agent:
        return MostBasicAgent(
            temperature=self.temperature,
            use_chain_of_thought=self.use_chain_of_thought,
            chat_model_args=self.chat_model_args,
        )

    def set_reproducibility_mode(self):
        self.temperature = 0

    def prepare(self):
        return self.chat_model_args.prepare_server()

    def close(self):
        return self.chat_model_args.close_server()

class MostBasicAgent(bgym.Agent):
    def __init__(
        self, temperature: float, use_chain_of_thought: bool, chat_model_args: "BaseModelArgs"
    ):
        self.temperature = temperature
        self.use_chain_of_thought = use_chain_of_thought
        self.chat = chat_model_args.make_model()
        self.chat_model_args = chat_model_args

        self.action_set = bgym.HighLevelActionSet(["bid"], multiaction=False)

    @cost_tracker_decorator


prepare

    from agentlab.llm.chat_api import BaseModelArgs

@dataclass
class MostBasicAgentArgs(AgentArgs):
    agent_name: str = "BasicAgent"
    temperature: float = 0.1
    use_chain_of_thought: bool = False
    chat_model_args: "BaseModelArgs" = None

    def make_agent(self) -> bgym.Agent:
        return MostBasicAgent(
            temperature=self.temperature,
            use_chain_of_thought=self.use_chain_of_thought,
            chat_model_args=self.chat_model_args,
        )

    def set_reproducibility_mode(self):
        self.temperature = 0

    def prepare(self):
        return self.chat_model_args.prepare_server()

    def close(self):
        return self.chat_model_args.close_server()

class MostBasicAgent(bgym.Agent):
    def __init__(
        self, temperature: float, use_chain_of_thought: bool, chat_model_args: "BaseModelArgs"
    ):
        self.temperature = temperature
        self.use_chain_of_thought = use_chain_of_thought
        self.chat = chat_model_args.make_model()
        self.chat_model_args = chat_model_args

        self.action_set = bgym.HighLevelActionSet(["bid"], multiaction=False)

    @cost_tracker_decorator
    def get_action(self, obs: Any) -> tuple[str, dict]:
        messages = Discussion(SystemMessage("You are a web assistant."))
        messages.append(


close

@dataclass
class MostBasicAgentArgs(AgentArgs):
    agent_name: str = "BasicAgent"
    temperature: float = 0.1
    use_chain_of_thought: bool = False
    chat_model_args: "BaseModelArgs" = None

    def make_agent(self) -> bgym.Agent:
        return MostBasicAgent(
            temperature=self.temperature,
            use_chain_of_thought=self.use_chain_of_thought,
            chat_model_args=self.chat_model_args,
        )

    def set_reproducibility_mode(self):
        self.temperature = 0

    def prepare(self):
        return self.chat_model_args.prepare_server()

    def close(self):
        return self.chat_model_args.close_server()

class MostBasicAgent(bgym.Agent):
    def __init__(
        self, temperature: float, use_chain_of_thought: bool, chat_model_args: "BaseModelArgs"
    ):
        self.temperature = temperature
        self.use_chain_of_thought = use_chain_of_thought
        self.chat = chat_model_args.make_model()
        self.chat_model_args = chat_model_args

        self.action_set = bgym.HighLevelActionSet(["bid"], multiaction=False)

    @cost_tracker_decorator
    def get_action(self, obs: Any) -> tuple[str, dict]:
        messages = Discussion(SystemMessage("You are a web assistant."))
        messages.append(
            HumanMessage(
                f"""
You are helping a user to accomplish the following goal on a website:


    chat_model_args: "BaseModelArgs" = None

    def make_agent(self) -> bgym.Agent:
        return MostBasicAgent(
            temperature=self.temperature,
            use_chain_of_thought=self.use_chain_of_thought,
            chat_model_args=self.chat_model_args,
        )

    def set_reproducibility_mode(self):
        self.temperature = 0

    def prepare(self):
        return self.chat_model_args.prepare_server()

    def close(self):
        return self.chat_model_args.close_server()

class MostBasicAgent(bgym.Agent):
    def __init__(
        self, temperature: float, use_chain_of_thought: bool, chat_model_args: "BaseModelArgs"
    ):
        self.temperature = temperature
        self.use_chain_of_thought = use_chain_of_thought
        self.chat = chat_model_args.make_model()
        self.chat_model_args = chat_model_args

        self.action_set = bgym.HighLevelActionSet(["bid"], multiaction=False)

    @cost_tracker_decorator
    def get_action(self, obs: Any) -> tuple[str, dict]:
        messages = Discussion(SystemMessage("You are a web assistant."))
        messages.append(
            HumanMessage(
                f"""
You are helping a user to accomplish the following goal on a website:

{obs["goal"]}

To do so, you can interact with the environment using the following actions:

{self.action_set.describe(with_long_description=False)}

The inputs to those functions are the bids given in the html.

Here is the current state of the website, in the form of an html:

{obs["pruned_html"]}


    def prepare(self):
        return self.chat_model_args.prepare_server()

    def close(self):
        return self.chat_model_args.close_server()

class MostBasicAgent(bgym.Agent):
    def __init__(
        self, temperature: float, use_chain_of_thought: bool, chat_model_args: "BaseModelArgs"
    ):
        self.temperature = temperature
        self.use_chain_of_thought = use_chain_of_thought
        self.chat = chat_model_args.make_model()
        self.chat_model_args = chat_model_args

        self.action_set = bgym.HighLevelActionSet(["bid"], multiaction=False)

    @cost_tracker_decorator
    def get_action(self, obs: Any) -> tuple[str, dict]:
        messages = Discussion(SystemMessage("You are a web assistant."))
        messages.append(
            HumanMessage(
                f"""
You are helping a user to accomplish the following goal on a website:

{obs["goal"]}

To do so, you can interact with the environment using the following actions:

{self.action_set.describe(with_long_description=False)}

The inputs to those functions are the bids given in the html.

Here is the current state of the website, in the form of an html:

{obs["pruned_html"]}

The action you provide must be in between triple ticks and leverage the 'bid=' information provided in the html.
Here is an example of how to use the bid action:

```
click('a314')
```

Please provide a single action at a time and wait for the next observation. Provide only a single action per step.
Focus on the bid that are given in the html, and use them to perform the actions.
"""
            )
        )
        if self.use_chain_of_thought:
            messages.add_text(
                f"""
Provide a chain of thoughts reasoning to decompose the task into smaller steps. And execute only the next step.
"""
            )

        def parser(response: str) -> tuple[dict, bool, str]:
            blocks = extract_code_blocks(response)
            if len(blocks) == 0:
                raise ParseError("No code block found in the response")
            action = blocks[0][1]
            thought = response
            return {"action": action, "think": thought}

        ans_dict = retry(self.chat, messages, n_retry=3, parser=parser)

        action = ans_dict.get("action", None)
        thought = ans_dict.get("think", None)

        return (
            action,
            bgym.AgentInfo(
                think=thought,
                chat_messages=messages,
                # put any stats that you care about as long as it is a number or a dict of numbers
                stats={"prompt_length": len(messages), "response_length": len(thought)},
                markdown_page="Add any txt information here, including base 64 images, to display in xray",
                extra_info={"chat_model_args": asdict(self.chat_model_args)},
            ),
        )

# example for a single task
env_args = bgym.EnvArgs(
    task_name="miniwob.click-button",
    task_seed=0,
    max_steps=10,
    headless=True,
)

chat_model_args = CHAT_MODEL_ARGS_DICT["openai/gpt-4o-mini-2024-07-18"]

# example for 2 experiments testing chain of thoughts on a miniwob task
exp_args = [
    ExpArgs(
        agent_args=MostBasicAgentArgs(
            temperature=0.1,
            use_chain_of_thought=True,
            chat_model_args=chat_model_args,
        ),


parser

The action you provide must be in between triple ticks and leverage the 'bid=' information provided in the html.
Here is an example of how to use the bid action:

```
click('a314')
```

Please provide a single action at a time and wait for the next observation. Provide only a single action per step.
Focus on the bid that are given in the html, and use them to perform the actions.
"""
            )
        )
        if self.use_chain_of_thought:
            messages.add_text(
                f"""
Provide a chain of thoughts reasoning to decompose the task into smaller steps. And execute only the next step.
"""
            )

        def parser(response: str) -> tuple[dict, bool, str]:
            blocks = extract_code_blocks(response)
            if len(blocks) == 0:
                raise ParseError("No code block found in the response")
            action = blocks[0][1]
            thought = response
            return {"action": action, "think": thought}

        ans_dict = retry(self.chat, messages, n_retry=3, parser=parser)

        action = ans_dict.get("action", None)
        thought = ans_dict.get("think", None)

        return (
            action,
            bgym.AgentInfo(
                think=thought,
                chat_messages=messages,
                # put any stats that you care about as long as it is a number or a dict of numbers
                stats={"prompt_length": len(messages), "response_length": len(thought)},
                markdown_page="Add any txt information here, including base 64 images, to display in xray",
                extra_info={"chat_model_args": asdict(self.chat_model_args)},
            ),
        )

# example for a single task


import bgym
from bgym import HighLevelActionSetArgs

import agentlab.agents.dynamic_prompting as dp
from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT

from .visual_agent import VisualAgentArgs
from .visual_agent_prompts import PromptFlags

# the other flags are ignored for this agent.
DEFAULT_OBS_FLAGS = dp.ObsFlags(
    use_tabs=True,  # will be overridden by the benchmark when set_benchmark is called after initalizing the agent
    use_error_logs=True,
    use_past_error_logs=False,
    use_screenshot=True,
    use_som=False,
    openai_vision_detail="auto",
)

DEFAULT_ACTION_FLAGS = dp.ActionFlags(
    action_set=HighLevelActionSetArgs(subsets=["coord"]),
    long_description=True,
    individual_examples=False,
)

DEFAULT_PROMPT_FLAGS = PromptFlags(
    obs=DEFAULT_OBS_FLAGS,
    action=DEFAULT_ACTION_FLAGS,
    use_thinking=True,
    use_concrete_example=False,
    use_abstract_example=True,
    enable_chat=False,
    extra_instructions=None,
)

VISUAL_AGENT_4o = VisualAgentArgs(
    chat_model_args=CHAT_MODEL_ARGS_DICT["openai/gpt-4o-2024-05-13"],
    flags=DEFAULT_PROMPT_FLAGS,
)

VISUAL_AGENT_CLAUDE_3_5 = VisualAgentArgs(
    chat_model_args=CHAT_MODEL_ARGS_DICT["openrouter/anthropic/claude-3.5-sonnet:beta"],
    flags=DEFAULT_PROMPT_FLAGS,
)


"""
GenericAgent implementation for AgentLab

This module defines a `GenericAgent` class and its associated arguments for use in the AgentLab framework. \
The `GenericAgent` class is designed to interact with a chat-based model to determine actions based on \
observations. It includes methods for preprocessing observations, generating actions, and managing internal \
state such as plans, memories, and thoughts. The `GenericAgentArgs` class provides configuration options for \
the agent, including model arguments and flags for various behaviors.
"""

from dataclasses import asdict, dataclass

import bgym
from bgym import Benchmark
from browsergym.experiments.agent import Agent, AgentInfo

from agentlab.agents import dynamic_prompting as dp
from agentlab.agents.agent_args import AgentArgs
from agentlab.llm.chat_api import BaseModelArgs
from agentlab.llm.llm_utils import Discussion, ParseError, SystemMessage, retry
from agentlab.llm.tracking import cost_tracker_decorator

from .visual_agent_prompts import MainPrompt, PromptFlags

@dataclass
class VisualAgentArgs(AgentArgs):
    chat_model_args: BaseModelArgs = None
    flags: PromptFlags = None
    max_retry: int = 4

    def __post_init__(self):
        try:  # some attributes might be missing temporarily due to args.CrossProd for hyperparameter generation
            self.agent_name = f"VisualAgent-{self.chat_model_args.model_name}".replace("/", "_")
        except AttributeError:
            pass

    def set_benchmark(self, benchmark: Benchmark, demo_mode):
        """Override Some flags based on the benchmark."""
        self.flags.obs.use_tabs = benchmark.is_multi_tab

    def set_reproducibility_mode(self):
        self.chat_model_args.temperature = 0

    def prepare(self):
        return self.chat_model_args.prepare_server()

    def close(self):
        return self.chat_model_args.close_server()

    def make_agent(self):
        return VisualAgent(
            chat_model_args=self.chat_model_args, flags=self.flags, max_retry=self.max_retry
        )

class VisualAgent(Agent):

    def __init__(
        self,
        chat_model_args: BaseModelArgs,
        flags: PromptFlags,
        max_retry: int = 4,
    ):

        self.chat_llm = chat_model_args.make_model()
        self.chat_model_args = chat_model_args
        self.max_retry = max_retry

        self.flags = flags
        self.action_set = self.flags.action.action_set.make_action_set()
        self._obs_preprocessor = dp.make_obs_preprocessor(flags.obs)

        self.reset(seed=None)

    def obs_preprocessor(self, obs: dict) -> dict:
        return self._obs_preprocessor(obs)

    @cost_tracker_decorator
    def get_action(self, obs):

        main_prompt = MainPrompt(
            action_set=self.action_set,
            obs=obs,
            actions=self.actions,
            thoughts=self.thoughts,
            flags=self.flags,
        )

        system_prompt = SystemMessage(dp.SystemPrompt().prompt)
        try:
            # TODO, we would need to further shrink the prompt if the retry
            # cause it to be too long

            chat_messages = Discussion([system_prompt, main_prompt.prompt])
            ans_dict = retry(
                self.chat_llm,
                chat_messages,
                n_retry=self.max_retry,
                parser=main_prompt._parse_answer,
            )
            ans_dict["busted_retry"] = 0
            # inferring the number of retries, TODO: make this less hacky
            ans_dict["n_retry"] = (len(chat_messages) - 3) / 2
        except ParseError:
            ans_dict = dict(
                action=None,
                n_retry=self.max_retry + 1,
                busted_retry=1,
            )

        stats = self.chat_llm.get_stats()
        stats["n_retry"] = ans_dict["n_retry"]
        stats["busted_retry"] = ans_dict["busted_retry"]

        self.actions.append(ans_dict["action"])
        self.thoughts.append(ans_dict.get("think", None))

        agent_info = AgentInfo(
            think=ans_dict.get("think", None),
            chat_messages=chat_messages,
            stats=stats,
            extra_info={"chat_model_args": asdict(self.chat_model_args)},
        )
        return ans_dict["action"], agent_info

    def reset(self, seed=None):
        self.seed = seed
        self.thoughts = []
        self.actions = []


VisualAgentArgs

state such as plans, memories, and thoughts. The `GenericAgentArgs` class provides configuration options for \
the agent, including model arguments and flags for various behaviors.
"""

from dataclasses import asdict, dataclass

import bgym
from bgym import Benchmark
from browsergym.experiments.agent import Agent, AgentInfo

from agentlab.agents import dynamic_prompting as dp
from agentlab.agents.agent_args import AgentArgs
from agentlab.llm.chat_api import BaseModelArgs
from agentlab.llm.llm_utils import Discussion, ParseError, SystemMessage, retry
from agentlab.llm.tracking import cost_tracker_decorator

from .visual_agent_prompts import MainPrompt, PromptFlags

@dataclass
class VisualAgentArgs(AgentArgs):
    chat_model_args: BaseModelArgs = None
    flags: PromptFlags = None
    max_retry: int = 4

    def __post_init__(self):
        try:  # some attributes might be missing temporarily due to args.CrossProd for hyperparameter generation
            self.agent_name = f"VisualAgent-{self.chat_model_args.model_name}".replace("/", "_")
        except AttributeError:
            pass

    def set_benchmark(self, benchmark: Benchmark, demo_mode):
        """Override Some flags based on the benchmark."""
        self.flags.obs.use_tabs = benchmark.is_multi_tab

    def set_reproducibility_mode(self):
        self.chat_model_args.temperature = 0

    def prepare(self):
        return self.chat_model_args.prepare_server()

    def close(self):
        return self.chat_model_args.close_server()

    def make_agent(self):
        return VisualAgent(
            chat_model_args=self.chat_model_args, flags=self.flags, max_retry=self.max_retry
        )

class VisualAgent(Agent):

    def __init__(
        self,
        chat_model_args: BaseModelArgs,
        flags: PromptFlags,
        max_retry: int = 4,
    ):

        self.chat_llm = chat_model_args.make_model()
        self.chat_model_args = chat_model_args
        self.max_retry = max_retry

        self.flags = flags
        self.action_set = self.flags.action.action_set.make_action_set()
        self._obs_preprocessor = dp.make_obs_preprocessor(flags.obs)

        self.reset(seed=None)


VisualAgent

    def set_benchmark(self, benchmark: Benchmark, demo_mode):
        """Override Some flags based on the benchmark."""
        self.flags.obs.use_tabs = benchmark.is_multi_tab

    def set_reproducibility_mode(self):
        self.chat_model_args.temperature = 0

    def prepare(self):
        return self.chat_model_args.prepare_server()

    def close(self):
        return self.chat_model_args.close_server()

    def make_agent(self):
        return VisualAgent(
            chat_model_args=self.chat_model_args, flags=self.flags, max_retry=self.max_retry
        )

class VisualAgent(Agent):

    def __init__(
        self,
        chat_model_args: BaseModelArgs,
        flags: PromptFlags,
        max_retry: int = 4,
    ):

        self.chat_llm = chat_model_args.make_model()
        self.chat_model_args = chat_model_args
        self.max_retry = max_retry

        self.flags = flags
        self.action_set = self.flags.action.action_set.make_action_set()
        self._obs_preprocessor = dp.make_obs_preprocessor(flags.obs)

        self.reset(seed=None)

    def obs_preprocessor(self, obs: dict) -> dict:
        return self._obs_preprocessor(obs)

    @cost_tracker_decorator
    def get_action(self, obs):

        main_prompt = MainPrompt(
            action_set=self.action_set,
            obs=obs,
            actions=self.actions,
            thoughts=self.thoughts,
            flags=self.flags,
        )

        system_prompt = SystemMessage(dp.SystemPrompt().prompt)
        try:
            # TODO, we would need to further shrink the prompt if the retry
            # cause it to be too long

            chat_messages = Discussion([system_prompt, main_prompt.prompt])
            ans_dict = retry(
                self.chat_llm,
                chat_messages,
                n_retry=self.max_retry,
                parser=main_prompt._parse_answer,
            )
            ans_dict["busted_retry"] = 0
            # inferring the number of retries, TODO: make this less hacky
            ans_dict["n_retry"] = (len(chat_messages) - 3) / 2
        except ParseError:
            ans_dict = dict(
                action=None,
                n_retry=self.max_retry + 1,
                busted_retry=1,
            )

        stats = self.chat_llm.get_stats()
        stats["n_retry"] = ans_dict["n_retry"]
        stats["busted_retry"] = ans_dict["busted_retry"]

        self.actions.append(ans_dict["action"])
        self.thoughts.append(ans_dict.get("think", None))

        agent_info = AgentInfo(
            think=ans_dict.get("think", None),
            chat_messages=chat_messages,
            stats=stats,
            extra_info={"chat_model_args": asdict(self.chat_model_args)},
        )
        return ans_dict["action"], agent_info

    def reset(self, seed=None):
        self.seed = seed
        self.thoughts = []
        self.actions = []


import bgym
from bgym import Benchmark
from browsergym.experiments.agent import Agent, AgentInfo

from agentlab.agents import dynamic_prompting as dp
from agentlab.agents.agent_args import AgentArgs
from agentlab.llm.chat_api import BaseModelArgs
from agentlab.llm.llm_utils import Discussion, ParseError, SystemMessage, retry
from agentlab.llm.tracking import cost_tracker_decorator

from .visual_agent_prompts import MainPrompt, PromptFlags

@dataclass
class VisualAgentArgs(AgentArgs):
    chat_model_args: BaseModelArgs = None
    flags: PromptFlags = None
    max_retry: int = 4

    def __post_init__(self):
        try:  # some attributes might be missing temporarily due to args.CrossProd for hyperparameter generation
            self.agent_name = f"VisualAgent-{self.chat_model_args.model_name}".replace("/", "_")
        except AttributeError:
            pass

    def set_benchmark(self, benchmark: Benchmark, demo_mode):
        """Override Some flags based on the benchmark."""
        self.flags.obs.use_tabs = benchmark.is_multi_tab

    def set_reproducibility_mode(self):
        self.chat_model_args.temperature = 0

    def prepare(self):
        return self.chat_model_args.prepare_server()

    def close(self):
        return self.chat_model_args.close_server()

    def make_agent(self):
        return VisualAgent(
            chat_model_args=self.chat_model_args, flags=self.flags, max_retry=self.max_retry
        )


from agentlab.agents.agent_args import AgentArgs
from agentlab.llm.chat_api import BaseModelArgs
from agentlab.llm.llm_utils import Discussion, ParseError, SystemMessage, retry
from agentlab.llm.tracking import cost_tracker_decorator

from .visual_agent_prompts import MainPrompt, PromptFlags

@dataclass
class VisualAgentArgs(AgentArgs):
    chat_model_args: BaseModelArgs = None
    flags: PromptFlags = None
    max_retry: int = 4

    def __post_init__(self):
        try:  # some attributes might be missing temporarily due to args.CrossProd for hyperparameter generation
            self.agent_name = f"VisualAgent-{self.chat_model_args.model_name}".replace("/", "_")
        except AttributeError:
            pass

    def set_benchmark(self, benchmark: Benchmark, demo_mode):
        """Override Some flags based on the benchmark."""
        self.flags.obs.use_tabs = benchmark.is_multi_tab

    def set_reproducibility_mode(self):
        self.chat_model_args.temperature = 0

    def prepare(self):
        return self.chat_model_args.prepare_server()

    def close(self):
        return self.chat_model_args.close_server()

    def make_agent(self):
        return VisualAgent(
            chat_model_args=self.chat_model_args, flags=self.flags, max_retry=self.max_retry
        )

class VisualAgent(Agent):

    def __init__(
        self,


from .visual_agent_prompts import MainPrompt, PromptFlags

@dataclass
class VisualAgentArgs(AgentArgs):
    chat_model_args: BaseModelArgs = None
    flags: PromptFlags = None
    max_retry: int = 4

    def __post_init__(self):
        try:  # some attributes might be missing temporarily due to args.CrossProd for hyperparameter generation
            self.agent_name = f"VisualAgent-{self.chat_model_args.model_name}".replace("/", "_")
        except AttributeError:
            pass

    def set_benchmark(self, benchmark: Benchmark, demo_mode):
        """Override Some flags based on the benchmark."""
        self.flags.obs.use_tabs = benchmark.is_multi_tab

    def set_reproducibility_mode(self):
        self.chat_model_args.temperature = 0

    def prepare(self):
        return self.chat_model_args.prepare_server()

    def close(self):
        return self.chat_model_args.close_server()

    def make_agent(self):
        return VisualAgent(
            chat_model_args=self.chat_model_args, flags=self.flags, max_retry=self.max_retry
        )

class VisualAgent(Agent):

    def __init__(
        self,
        chat_model_args: BaseModelArgs,
        flags: PromptFlags,
        max_retry: int = 4,


prepare

@dataclass
class VisualAgentArgs(AgentArgs):
    chat_model_args: BaseModelArgs = None
    flags: PromptFlags = None
    max_retry: int = 4

    def __post_init__(self):
        try:  # some attributes might be missing temporarily due to args.CrossProd for hyperparameter generation
            self.agent_name = f"VisualAgent-{self.chat_model_args.model_name}".replace("/", "_")
        except AttributeError:
            pass

    def set_benchmark(self, benchmark: Benchmark, demo_mode):
        """Override Some flags based on the benchmark."""
        self.flags.obs.use_tabs = benchmark.is_multi_tab

    def set_reproducibility_mode(self):
        self.chat_model_args.temperature = 0

    def prepare(self):
        return self.chat_model_args.prepare_server()

    def close(self):
        return self.chat_model_args.close_server()

    def make_agent(self):
        return VisualAgent(
            chat_model_args=self.chat_model_args, flags=self.flags, max_retry=self.max_retry
        )

class VisualAgent(Agent):

    def __init__(
        self,
        chat_model_args: BaseModelArgs,
        flags: PromptFlags,
        max_retry: int = 4,
    ):

        self.chat_llm = chat_model_args.make_model()


close

    chat_model_args: BaseModelArgs = None
    flags: PromptFlags = None
    max_retry: int = 4

    def __post_init__(self):
        try:  # some attributes might be missing temporarily due to args.CrossProd for hyperparameter generation
            self.agent_name = f"VisualAgent-{self.chat_model_args.model_name}".replace("/", "_")
        except AttributeError:
            pass

    def set_benchmark(self, benchmark: Benchmark, demo_mode):
        """Override Some flags based on the benchmark."""
        self.flags.obs.use_tabs = benchmark.is_multi_tab

    def set_reproducibility_mode(self):
        self.chat_model_args.temperature = 0

    def prepare(self):
        return self.chat_model_args.prepare_server()

    def close(self):
        return self.chat_model_args.close_server()

    def make_agent(self):
        return VisualAgent(
            chat_model_args=self.chat_model_args, flags=self.flags, max_retry=self.max_retry
        )

class VisualAgent(Agent):

    def __init__(
        self,
        chat_model_args: BaseModelArgs,
        flags: PromptFlags,
        max_retry: int = 4,
    ):

        self.chat_llm = chat_model_args.make_model()
        self.chat_model_args = chat_model_args
        self.max_retry = max_retry


    def __post_init__(self):
        try:  # some attributes might be missing temporarily due to args.CrossProd for hyperparameter generation
            self.agent_name = f"VisualAgent-{self.chat_model_args.model_name}".replace("/", "_")
        except AttributeError:
            pass

    def set_benchmark(self, benchmark: Benchmark, demo_mode):
        """Override Some flags based on the benchmark."""
        self.flags.obs.use_tabs = benchmark.is_multi_tab

    def set_reproducibility_mode(self):
        self.chat_model_args.temperature = 0

    def prepare(self):
        return self.chat_model_args.prepare_server()

    def close(self):
        return self.chat_model_args.close_server()

    def make_agent(self):
        return VisualAgent(
            chat_model_args=self.chat_model_args, flags=self.flags, max_retry=self.max_retry
        )

class VisualAgent(Agent):

    def __init__(
        self,
        chat_model_args: BaseModelArgs,
        flags: PromptFlags,
        max_retry: int = 4,
    ):

        self.chat_llm = chat_model_args.make_model()
        self.chat_model_args = chat_model_args
        self.max_retry = max_retry

        self.flags = flags
        self.action_set = self.flags.action.action_set.make_action_set()
        self._obs_preprocessor = dp.make_obs_preprocessor(flags.obs)

        self.reset(seed=None)


        """Override Some flags based on the benchmark."""
        self.flags.obs.use_tabs = benchmark.is_multi_tab

    def set_reproducibility_mode(self):
        self.chat_model_args.temperature = 0

    def prepare(self):
        return self.chat_model_args.prepare_server()

    def close(self):
        return self.chat_model_args.close_server()

    def make_agent(self):
        return VisualAgent(
            chat_model_args=self.chat_model_args, flags=self.flags, max_retry=self.max_retry
        )

class VisualAgent(Agent):

    def __init__(
        self,
        chat_model_args: BaseModelArgs,
        flags: PromptFlags,
        max_retry: int = 4,
    ):

        self.chat_llm = chat_model_args.make_model()
        self.chat_model_args = chat_model_args
        self.max_retry = max_retry

        self.flags = flags
        self.action_set = self.flags.action.action_set.make_action_set()
        self._obs_preprocessor = dp.make_obs_preprocessor(flags.obs)

        self.reset(seed=None)

    def obs_preprocessor(self, obs: dict) -> dict:
        return self._obs_preprocessor(obs)

    @cost_tracker_decorator
    def get_action(self, obs):

        main_prompt = MainPrompt(
            action_set=self.action_set,
            obs=obs,
            actions=self.actions,
            thoughts=self.thoughts,
            flags=self.flags,
        )

        system_prompt = SystemMessage(dp.SystemPrompt().prompt)
        try:
            # TODO, we would need to further shrink the prompt if the retry
            # cause it to be too long


class VisualAgent(Agent):

    def __init__(
        self,
        chat_model_args: BaseModelArgs,
        flags: PromptFlags,
        max_retry: int = 4,
    ):

        self.chat_llm = chat_model_args.make_model()
        self.chat_model_args = chat_model_args
        self.max_retry = max_retry

        self.flags = flags
        self.action_set = self.flags.action.action_set.make_action_set()
        self._obs_preprocessor = dp.make_obs_preprocessor(flags.obs)

        self.reset(seed=None)

    def obs_preprocessor(self, obs: dict) -> dict:
        return self._obs_preprocessor(obs)

    @cost_tracker_decorator
    def get_action(self, obs):

        main_prompt = MainPrompt(
            action_set=self.action_set,
            obs=obs,
            actions=self.actions,
            thoughts=self.thoughts,
            flags=self.flags,
        )

        system_prompt = SystemMessage(dp.SystemPrompt().prompt)
        try:
            # TODO, we would need to further shrink the prompt if the retry
            # cause it to be too long

            chat_messages = Discussion([system_prompt, main_prompt.prompt])
            ans_dict = retry(
                self.chat_llm,


        self,
        chat_model_args: BaseModelArgs,
        flags: PromptFlags,
        max_retry: int = 4,
    ):

        self.chat_llm = chat_model_args.make_model()
        self.chat_model_args = chat_model_args
        self.max_retry = max_retry

        self.flags = flags
        self.action_set = self.flags.action.action_set.make_action_set()
        self._obs_preprocessor = dp.make_obs_preprocessor(flags.obs)

        self.reset(seed=None)

    def obs_preprocessor(self, obs: dict) -> dict:
        return self._obs_preprocessor(obs)

    @cost_tracker_decorator
    def get_action(self, obs):

        main_prompt = MainPrompt(
            action_set=self.action_set,
            obs=obs,
            actions=self.actions,
            thoughts=self.thoughts,
            flags=self.flags,
        )

        system_prompt = SystemMessage(dp.SystemPrompt().prompt)
        try:
            # TODO, we would need to further shrink the prompt if the retry
            # cause it to be too long

            chat_messages = Discussion([system_prompt, main_prompt.prompt])
            ans_dict = retry(
                self.chat_llm,
                chat_messages,
                n_retry=self.max_retry,
                parser=main_prompt._parse_answer,
            )
            ans_dict["busted_retry"] = 0
            # inferring the number of retries, TODO: make this less hacky
            ans_dict["n_retry"] = (len(chat_messages) - 3) / 2
        except ParseError:
            ans_dict = dict(
                action=None,
                n_retry=self.max_retry + 1,
                busted_retry=1,
            )

        stats = self.chat_llm.get_stats()
        stats["n_retry"] = ans_dict["n_retry"]
        stats["busted_retry"] = ans_dict["busted_retry"]

        self.actions.append(ans_dict["action"])
        self.thoughts.append(ans_dict.get("think", None))

        agent_info = AgentInfo(
            think=ans_dict.get("think", None),
            chat_messages=chat_messages,
            stats=stats,
            extra_info={"chat_model_args": asdict(self.chat_model_args)},
        )
        return ans_dict["action"], agent_info

    def reset(self, seed=None):
        self.seed = seed
        self.thoughts = []
        self.actions = []


reset

                action=None,
                n_retry=self.max_retry + 1,
                busted_retry=1,
            )

        stats = self.chat_llm.get_stats()
        stats["n_retry"] = ans_dict["n_retry"]
        stats["busted_retry"] = ans_dict["busted_retry"]

        self.actions.append(ans_dict["action"])
        self.thoughts.append(ans_dict.get("think", None))

        agent_info = AgentInfo(
            think=ans_dict.get("think", None),
            chat_messages=chat_messages,
            stats=stats,
            extra_info={"chat_model_args": asdict(self.chat_model_args)},
        )
        return ans_dict["action"], agent_info

    def reset(self, seed=None):
        self.seed = seed
        self.thoughts = []
        self.actions = []


"""
Prompt builder for GenericAgent

It is based on the dynamic_prompting module from the agentlab package.
"""

import logging
from dataclasses import dataclass
import bgym

from browsergym.core.action.base import AbstractActionSet

from agentlab.agents import dynamic_prompting as dp
from agentlab.llm.llm_utils import BaseMessage, HumanMessage, image_to_jpg_base64_url

@dataclass
class PromptFlags(dp.Flags):
    """
    A class to represent various flags used to control features in an application.
    """

    obs: dp.ObsFlags = None
    action: dp.ActionFlags = None
    use_thinking: bool = True
    use_concrete_example: bool = False
    use_abstract_example: bool = True
    enable_chat: bool = False
    extra_instructions: str | None = None

class SystemPrompt(dp.PromptElement):
    _prompt = """\
You are an agent trying to solve a web task based on the content of the page and
user instructions. You can interact with the page and explore, and send messages to the user. Each time you
submit an action it will be sent to the browser and you will receive a new page."""

def make_instructions(obs: dict, from_chat: bool, extra_instructions: str | None):
    """Convenient wrapper to extract instructions from either goal or chat"""
    if from_chat:
        instructions = dp.ChatInstructions(
            obs["chat_messages"], extra_instructions=extra_instructions
        )
    else:
        if sum([msg["role"] == "user" for msg in obs.get("chat_messages", [])]) > 1:
            logging.warning(
                "Agent is in goal mode, but multiple user messages are present in the chat. Consider switching to `enable_chat=True`."
            )
        instructions = dp.GoalInstructions(
            obs["goal_object"], extra_instructions=extra_instructions
        )
    return instructions

class History(dp.PromptElement):
    """
    Format the actions and thoughts of previous steps."""

    def __init__(self, actions, thoughts) -> None:
        super().__init__()
        prompt_elements = []
        for i, (action, thought) in enumerate(zip(actions, thoughts)):
            prompt_elements.append(
                f"""
## Step {i}
### Thoughts:
{thought}
### Action:
{action}
"""
            )
        self._prompt = "\n".join(prompt_elements) + "\n"

class Observation(dp.PromptElement):
    """Observation of the current step.

    Contains the html, the accessibility tree and the error logs.
    """

    def __init__(self, obs, flags: dp.ObsFlags) -> None:
        super().__init__()
        self.flags = flags
        self.obs = obs

        # for a multi-tab browser, we need to show the current tab
        self.tabs = dp.Tabs(
            obs,
            visible=lambda: flags.use_tabs,
            prefix="## ",
        )

        # if an error is present, we need to show it
        self.error = dp.Error(
            obs["last_action_error"],
            visible=lambda: flags.use_error_logs and obs["last_action_error"],
            prefix="## ",
        )

    @property
    def _prompt(self) -> str:
        return f"""
# Observation of current step:
{self.tabs.prompt}{self.error.prompt}

"""

    def add_screenshot(self, prompt: BaseMessage) -> BaseMessage:
        if self.flags.use_screenshot:
            if self.flags.use_som:
                screenshot = self.obs["screenshot_som"]
                prompt.add_text(
                    "\n## Screenshot:\nHere is a screenshot of the page, it is annotated with bounding boxes and corresponding bids:"
                )
            else:
                screenshot = self.obs["screenshot"]
                prompt.add_text("\n## Screenshot:\nHere is a screenshot of the page:")
            img_url = image_to_jpg_base64_url(screenshot)
            prompt.add_image(img_url, detail=self.flags.openai_vision_detail)
        return prompt

class MainPrompt(dp.PromptElement):

    def __init__(
        self,
        action_set: AbstractActionSet,
        obs: dict,
        actions: list[str],
        thoughts: list[str],
        flags: PromptFlags,
    ) -> None:
        super().__init__()
        self.flags = flags
        self.history = History(actions, thoughts)
        self.instructions = make_instructions(obs, flags.enable_chat, flags.extra_instructions)
        self.obs = Observation(obs, self.flags.obs)

        self.action_prompt = dp.ActionPrompt(action_set, action_flags=flags.action)
        self.think = dp.Think(visible=lambda: flags.use_thinking)

    @property
    def _prompt(self) -> HumanMessage:
        prompt = HumanMessage(self.instructions.prompt)
        prompt.add_text(
            f"""\
{self.obs.prompt}\
{self.history.prompt}\
{self.action_prompt.prompt}\
{self.think.prompt}\
"""
        )

        if self.flags.use_abstract_example:
            prompt.add_text(
                f"""
# Abstract Example

Here is an abstract version of the answer with description of the content of
each tag. Make sure you follow this structure, but replace the content with your
answer:
{self.think.abstract_ex}\
{self.action_prompt.abstract_ex}\
"""
            )

        if self.flags.use_concrete_example:
            prompt.add_text(
                f"""
# Concrete Example

Here is a concrete example of how to format your answer.
Make sure to follow the template with proper tags:
{self.think.concrete_ex}\
{self.action_prompt.concrete_ex}\
"""
            )
        return self.obs.add_screenshot(prompt)

    def _parse_answer(self, text_answer):
        ans_dict = {}
        ans_dict.update(self.think.parse_answer(text_answer))
        ans_dict.update(self.action_prompt.parse_answer(text_answer))
        return ans_dict


PromptFlags

"""
Prompt builder for GenericAgent

It is based on the dynamic_prompting module from the agentlab package.
"""

import logging
from dataclasses import dataclass
import bgym

from browsergym.core.action.base import AbstractActionSet

from agentlab.agents import dynamic_prompting as dp
from agentlab.llm.llm_utils import BaseMessage, HumanMessage, image_to_jpg_base64_url

@dataclass
class PromptFlags(dp.Flags):
    """
    A class to represent various flags used to control features in an application.
    """

    obs: dp.ObsFlags = None
    action: dp.ActionFlags = None
    use_thinking: bool = True
    use_concrete_example: bool = False
    use_abstract_example: bool = True
    enable_chat: bool = False
    extra_instructions: str | None = None

class SystemPrompt(dp.PromptElement):
    _prompt = """\
You are an agent trying to solve a web task based on the content of the page and
user instructions. You can interact with the page and explore, and send messages to the user. Each time you
submit an action it will be sent to the browser and you will receive a new page."""

def make_instructions(obs: dict, from_chat: bool, extra_instructions: str | None):
    """Convenient wrapper to extract instructions from either goal or chat"""
    if from_chat:
        instructions = dp.ChatInstructions(
            obs["chat_messages"], extra_instructions=extra_instructions
        )
    else:
        if sum([msg["role"] == "user" for msg in obs.get("chat_messages", [])]) > 1:
            logging.warning(
                "Agent is in goal mode, but multiple user messages are present in the chat. Consider switching to `enable_chat=True`."
            )


SystemPrompt

from agentlab.agents import dynamic_prompting as dp
from agentlab.llm.llm_utils import BaseMessage, HumanMessage, image_to_jpg_base64_url

@dataclass
class PromptFlags(dp.Flags):
    """
    A class to represent various flags used to control features in an application.
    """

    obs: dp.ObsFlags = None
    action: dp.ActionFlags = None
    use_thinking: bool = True
    use_concrete_example: bool = False
    use_abstract_example: bool = True
    enable_chat: bool = False
    extra_instructions: str | None = None

class SystemPrompt(dp.PromptElement):
    _prompt = """\
You are an agent trying to solve a web task based on the content of the page and
user instructions. You can interact with the page and explore, and send messages to the user. Each time you
submit an action it will be sent to the browser and you will receive a new page."""

def make_instructions(obs: dict, from_chat: bool, extra_instructions: str | None):
    """Convenient wrapper to extract instructions from either goal or chat"""
    if from_chat:
        instructions = dp.ChatInstructions(
            obs["chat_messages"], extra_instructions=extra_instructions
        )
    else:
        if sum([msg["role"] == "user" for msg in obs.get("chat_messages", [])]) > 1:
            logging.warning(
                "Agent is in goal mode, but multiple user messages are present in the chat. Consider switching to `enable_chat=True`."
            )
        instructions = dp.GoalInstructions(
            obs["goal_object"], extra_instructions=extra_instructions
        )
    return instructions

class History(dp.PromptElement):


    """
    A class to represent various flags used to control features in an application.
    """

    obs: dp.ObsFlags = None
    action: dp.ActionFlags = None
    use_thinking: bool = True
    use_concrete_example: bool = False
    use_abstract_example: bool = True
    enable_chat: bool = False
    extra_instructions: str | None = None

class SystemPrompt(dp.PromptElement):
    _prompt = """\
You are an agent trying to solve a web task based on the content of the page and
user instructions. You can interact with the page and explore, and send messages to the user. Each time you
submit an action it will be sent to the browser and you will receive a new page."""

def make_instructions(obs: dict, from_chat: bool, extra_instructions: str | None):
    """Convenient wrapper to extract instructions from either goal or chat"""
    if from_chat:
        instructions = dp.ChatInstructions(
            obs["chat_messages"], extra_instructions=extra_instructions
        )
    else:
        if sum([msg["role"] == "user" for msg in obs.get("chat_messages", [])]) > 1:
            logging.warning(
                "Agent is in goal mode, but multiple user messages are present in the chat. Consider switching to `enable_chat=True`."
            )
        instructions = dp.GoalInstructions(
            obs["goal_object"], extra_instructions=extra_instructions
        )
    return instructions

class History(dp.PromptElement):
    """
    Format the actions and thoughts of previous steps."""

    def __init__(self, actions, thoughts) -> None:
        super().__init__()
        prompt_elements = []
        for i, (action, thought) in enumerate(zip(actions, thoughts)):
            prompt_elements.append(
                f"""
## Step {i}
### Thoughts:
{thought}
### Action:
{action}
"""
            )
        self._prompt = "\n".join(prompt_elements) + "\n"


History

submit an action it will be sent to the browser and you will receive a new page."""

def make_instructions(obs: dict, from_chat: bool, extra_instructions: str | None):
    """Convenient wrapper to extract instructions from either goal or chat"""
    if from_chat:
        instructions = dp.ChatInstructions(
            obs["chat_messages"], extra_instructions=extra_instructions
        )
    else:
        if sum([msg["role"] == "user" for msg in obs.get("chat_messages", [])]) > 1:
            logging.warning(
                "Agent is in goal mode, but multiple user messages are present in the chat. Consider switching to `enable_chat=True`."
            )
        instructions = dp.GoalInstructions(
            obs["goal_object"], extra_instructions=extra_instructions
        )
    return instructions

class History(dp.PromptElement):
    """
    Format the actions and thoughts of previous steps."""

    def __init__(self, actions, thoughts) -> None:
        super().__init__()
        prompt_elements = []
        for i, (action, thought) in enumerate(zip(actions, thoughts)):
            prompt_elements.append(
                f"""
## Step {i}
### Thoughts:
{thought}
### Action:
{action}
"""
            )
        self._prompt = "\n".join(prompt_elements) + "\n"

class Observation(dp.PromptElement):
    """Observation of the current step.

    Contains the html, the accessibility tree and the error logs.
    """

    def __init__(self, obs, flags: dp.ObsFlags) -> None:
        super().__init__()
        self.flags = flags
        self.obs = obs

        # for a multi-tab browser, we need to show the current tab
        self.tabs = dp.Tabs(
            obs,
            visible=lambda: flags.use_tabs,
            prefix="## ",
        )


Observation

class History(dp.PromptElement):
    """
    Format the actions and thoughts of previous steps."""

    def __init__(self, actions, thoughts) -> None:
        super().__init__()
        prompt_elements = []
        for i, (action, thought) in enumerate(zip(actions, thoughts)):
            prompt_elements.append(
                f"""
## Step {i}
### Thoughts:
{thought}
### Action:
{action}
"""
            )
        self._prompt = "\n".join(prompt_elements) + "\n"

class Observation(dp.PromptElement):
    """Observation of the current step.

    Contains the html, the accessibility tree and the error logs.
    """

    def __init__(self, obs, flags: dp.ObsFlags) -> None:
        super().__init__()
        self.flags = flags
        self.obs = obs

        # for a multi-tab browser, we need to show the current tab
        self.tabs = dp.Tabs(
            obs,
            visible=lambda: flags.use_tabs,
            prefix="## ",
        )

        # if an error is present, we need to show it
        self.error = dp.Error(
            obs["last_action_error"],
            visible=lambda: flags.use_error_logs and obs["last_action_error"],
            prefix="## ",
        )

    @property
    def _prompt(self) -> str:
        return f"""
# Observation of current step:
{self.tabs.prompt}{self.error.prompt}

"""

    def add_screenshot(self, prompt: BaseMessage) -> BaseMessage:
        if self.flags.use_screenshot:
            if self.flags.use_som:
                screenshot = self.obs["screenshot_som"]
                prompt.add_text(
                    "\n## Screenshot:\nHere is a screenshot of the page, it is annotated with bounding boxes and corresponding bids:"
                )
            else:
                screenshot = self.obs["screenshot"]
                prompt.add_text("\n## Screenshot:\nHere is a screenshot of the page:")
            img_url = image_to_jpg_base64_url(screenshot)
            prompt.add_image(img_url, detail=self.flags.openai_vision_detail)
        return prompt

class MainPrompt(dp.PromptElement):

    def __init__(
        self,
        action_set: AbstractActionSet,
        obs: dict,
        actions: list[str],
        thoughts: list[str],
        flags: PromptFlags,
    ) -> None:
        super().__init__()
        self.flags = flags
        self.history = History(actions, thoughts)
        self.instructions = make_instructions(obs, flags.enable_chat, flags.extra_instructions)
        self.obs = Observation(obs, self.flags.obs)

        self.action_prompt = dp.ActionPrompt(action_set, action_flags=flags.action)
        self.think = dp.Think(visible=lambda: flags.use_thinking)


MainPrompt

# Observation of current step:
{self.tabs.prompt}{self.error.prompt}

"""

    def add_screenshot(self, prompt: BaseMessage) -> BaseMessage:
        if self.flags.use_screenshot:
            if self.flags.use_som:
                screenshot = self.obs["screenshot_som"]
                prompt.add_text(
                    "\n## Screenshot:\nHere is a screenshot of the page, it is annotated with bounding boxes and corresponding bids:"
                )
            else:
                screenshot = self.obs["screenshot"]
                prompt.add_text("\n## Screenshot:\nHere is a screenshot of the page:")
            img_url = image_to_jpg_base64_url(screenshot)
            prompt.add_image(img_url, detail=self.flags.openai_vision_detail)
        return prompt

class MainPrompt(dp.PromptElement):

    def __init__(
        self,
        action_set: AbstractActionSet,
        obs: dict,
        actions: list[str],
        thoughts: list[str],
        flags: PromptFlags,
    ) -> None:
        super().__init__()
        self.flags = flags
        self.history = History(actions, thoughts)
        self.instructions = make_instructions(obs, flags.enable_chat, flags.extra_instructions)
        self.obs = Observation(obs, self.flags.obs)

        self.action_prompt = dp.ActionPrompt(action_set, action_flags=flags.action)
        self.think = dp.Think(visible=lambda: flags.use_thinking)

    @property
    def _prompt(self) -> HumanMessage:
        prompt = HumanMessage(self.instructions.prompt)
        prompt.add_text(
            f"""\
{self.obs.prompt}\
{self.history.prompt}\
{self.action_prompt.prompt}\
{self.think.prompt}\
"""
        )

        if self.flags.use_abstract_example:
            prompt.add_text(
                f"""
# Abstract Example

Here is an abstract version of the answer with description of the content of
each tag. Make sure you follow this structure, but replace the content with your
answer:
{self.think.abstract_ex}\
{self.action_prompt.abstract_ex}\
"""
            )

        if self.flags.use_concrete_example:
            prompt.add_text(
                f"""
# Concrete Example

Here is a concrete example of how to format your answer.
Make sure to follow the template with proper tags:
{self.think.concrete_ex}\
{self.action_prompt.concrete_ex}\
"""
            )
        return self.obs.add_screenshot(prompt)

    def _parse_answer(self, text_answer):
        ans_dict = {}
        ans_dict.update(self.think.parse_answer(text_answer))
        ans_dict.update(self.action_prompt.parse_answer(text_answer))
        return ans_dict


"""

    def add_screenshot(self, prompt: BaseMessage) -> BaseMessage:
        if self.flags.use_screenshot:
            if self.flags.use_som:
                screenshot = self.obs["screenshot_som"]
                prompt.add_text(
                    "\n## Screenshot:\nHere is a screenshot of the page, it is annotated with bounding boxes and corresponding bids:"
                )
            else:
                screenshot = self.obs["screenshot"]
                prompt.add_text("\n## Screenshot:\nHere is a screenshot of the page:")
            img_url = image_to_jpg_base64_url(screenshot)
            prompt.add_image(img_url, detail=self.flags.openai_vision_detail)
        return prompt

class MainPrompt(dp.PromptElement):

    def __init__(
        self,
        action_set: AbstractActionSet,
        obs: dict,
        actions: list[str],
        thoughts: list[str],
        flags: PromptFlags,
    ) -> None:
        super().__init__()
        self.flags = flags
        self.history = History(actions, thoughts)
        self.instructions = make_instructions(obs, flags.enable_chat, flags.extra_instructions)
        self.obs = Observation(obs, self.flags.obs)

        self.action_prompt = dp.ActionPrompt(action_set, action_flags=flags.action)
        self.think = dp.Think(visible=lambda: flags.use_thinking)

    @property
    def _prompt(self) -> HumanMessage:
        prompt = HumanMessage(self.instructions.prompt)
        prompt.add_text(
            f"""\
{self.obs.prompt}\
{self.history.prompt}\
{self.action_prompt.prompt}\
{self.think.prompt}\
"""
        )

        if self.flags.use_abstract_example:
            prompt.add_text(
                f"""
# Abstract Example

Here is an abstract version of the answer with description of the content of
each tag. Make sure you follow this structure, but replace the content with your


class MainPrompt(dp.PromptElement):

    def __init__(
        self,
        action_set: AbstractActionSet,
        obs: dict,
        actions: list[str],
        thoughts: list[str],
        flags: PromptFlags,
    ) -> None:
        super().__init__()
        self.flags = flags
        self.history = History(actions, thoughts)
        self.instructions = make_instructions(obs, flags.enable_chat, flags.extra_instructions)
        self.obs = Observation(obs, self.flags.obs)

        self.action_prompt = dp.ActionPrompt(action_set, action_flags=flags.action)
        self.think = dp.Think(visible=lambda: flags.use_thinking)

    @property
    def _prompt(self) -> HumanMessage:
        prompt = HumanMessage(self.instructions.prompt)
        prompt.add_text(
            f"""\
{self.obs.prompt}\
{self.history.prompt}\
{self.action_prompt.prompt}\
{self.think.prompt}\
"""
        )

        if self.flags.use_abstract_example:
            prompt.add_text(
                f"""
# Abstract Example

Here is an abstract version of the answer with description of the content of
each tag. Make sure you follow this structure, but replace the content with your
answer:
{self.think.abstract_ex}\
{self.action_prompt.abstract_ex}\
"""
            )

        if self.flags.use_concrete_example:
            prompt.add_text(
                f"""
# Concrete Example

Here is a concrete example of how to format your answer.
Make sure to follow the template with proper tags:
{self.think.concrete_ex}\
{self.action_prompt.concrete_ex}\
"""
            )
        return self.obs.add_screenshot(prompt)

    def _parse_answer(self, text_answer):
        ans_dict = {}
        ans_dict.update(self.think.parse_answer(text_answer))
        ans_dict.update(self.action_prompt.parse_answer(text_answer))
        return ans_dict


            obs,
            visible=lambda: flags.use_tabs,
            prefix="## ",
        )

        # if an error is present, we need to show it
        self.error = dp.Error(
            obs["last_action_error"],
            visible=lambda: flags.use_error_logs and obs["last_action_error"],
            prefix="## ",
        )

    @property
    def _prompt(self) -> str:
        return f"""
# Observation of current step:
{self.tabs.prompt}{self.error.prompt}

"""

    def add_screenshot(self, prompt: BaseMessage) -> BaseMessage:
        if self.flags.use_screenshot:
            if self.flags.use_som:
                screenshot = self.obs["screenshot_som"]
                prompt.add_text(
                    "\n## Screenshot:\nHere is a screenshot of the page, it is annotated with bounding boxes and corresponding bids:"
                )
            else:
                screenshot = self.obs["screenshot"]
                prompt.add_text("\n## Screenshot:\nHere is a screenshot of the page:")
            img_url = image_to_jpg_base64_url(screenshot)
            prompt.add_image(img_url, detail=self.flags.openai_vision_detail)
        return prompt

class MainPrompt(dp.PromptElement):

    def __init__(
        self,
        action_set: AbstractActionSet,
        obs: dict,
        actions: list[str],
        thoughts: list[str],
        flags: PromptFlags,
    ) -> None:
        super().__init__()
        self.flags = flags
        self.history = History(actions, thoughts)
        self.instructions = make_instructions(obs, flags.enable_chat, flags.extra_instructions)
        self.obs = Observation(obs, self.flags.obs)

        self.action_prompt = dp.ActionPrompt(action_set, action_flags=flags.action)
        self.think = dp.Think(visible=lambda: flags.use_thinking)


each tag. Make sure you follow this structure, but replace the content with your
answer:
{self.think.abstract_ex}\
{self.action_prompt.abstract_ex}\
"""
            )

        if self.flags.use_concrete_example:
            prompt.add_text(
                f"""
# Concrete Example

Here is a concrete example of how to format your answer.
Make sure to follow the template with proper tags:
{self.think.concrete_ex}\
{self.action_prompt.concrete_ex}\
"""
            )
        return self.obs.add_screenshot(prompt)

    def _parse_answer(self, text_answer):
        ans_dict = {}
        ans_dict.update(self.think.parse_answer(text_answer))
        ans_dict.update(self.action_prompt.parse_answer(text_answer))
        return ans_dict


src.agentlab.agents.tapeagent.agent

import logging
from dataclasses import dataclass
from typing import Literal

import bgym
import hydra
from omegaconf import DictConfig
from pydantic import Field
from tapeagents.agent import Agent
from tapeagents.core import Action, Observation, StopStep, TapeMetadata, Thought
from tapeagents.core import Tape as BaseTape

from agentlab.agents.agent_args import AgentArgs

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

class ExtendedMetadata(TapeMetadata):
    name: str = ""
    task: dict = {}
    terminated: bool = False
    truncated: bool = False
    reward: float = 0.0
    attempt_number: int = 0
    other: dict = {}

class Tape(BaseTape):
    metadata: ExtendedMetadata = Field(default_factory=ExtendedMetadata)  # type: ignore

def load_config(config_name: str) -> DictConfig:
    with hydra.initialize(config_path="conf", version_base="1.1"):
        config = hydra.compose(config_name=config_name)
    return config

@dataclass
class TapeAgentArgs(AgentArgs):
    config: DictConfig = None  # type: ignore

    def make_agent(self) -> bgym.Agent:
        agent: Agent = hydra.utils.instantiate(self.config.agent)
        return TapeAgent(agent=agent)

@dataclass
class TapeAgentInfo(bgym.AgentInfo):
    thoughts: list[Thought] = None  # type: ignore

class DictObservation(Observation):
    """
    Container for wrapping old dict observation into new Observation class.
    """

    kind: Literal["dict_observation"] = "dict_observation"  # type: ignore
    content: str

class TapeAgent(bgym.Agent):
    agent: Agent
    tape: Tape

    def __init__(self, agent: Agent):
        super().__init__()
        self.agent = agent
        self.tape = Tape(steps=[])

    def obs_preprocessor(self, obs: Observation | list[Observation]) -> list[Observation]:
        if isinstance(obs, Observation):
            obs = [obs]
        assert isinstance(obs, list), f"Expected list of Observations, got {type(obs)}"
        logger.info(f"Observations: {[type(o).__name__ for o in obs]}")
        return obs

    def get_action(self, obs: Observation | list[Observation]) -> tuple[Action, TapeAgentInfo]:
        self.tape += obs  # type: ignore
        thoughts: list[Thought] = []
        action = None
        while not action:
            for event in self.agent.run(self.tape):
                if not event.step:
                    continue
                self.tape = self.tape.append(event.step)
                if isinstance(event.step, Thought):
                    thoughts.append(event.step)
                    logger.info(f"Thought: {event.step.llm_view()}")
                elif isinstance(event.step, Action) and not action:  # we use first action only
                    action = event.step
                    logger.info(f"Action: {action.llm_view()}")
                else:
                    # there could be control flow steps for switching nodes and if clauses
                    logger.info(f"Other step: {type(event.step)}")
        logger.info(f"Tape after run: ({len(self.tape)}) {[type(s).__name__ for s in self.tape]}")
        return (action, TapeAgentInfo(thoughts=thoughts))

    @property
    def final_tape(self) -> Tape:
        truncated = not any([isinstance(s, StopStep) for s in self.tape.steps])
        self.tape.metadata = ExtendedMetadata(author=self.agent.name, truncated=truncated)
        return self.tape


ExtendedMetadata

import logging
from dataclasses import dataclass
from typing import Literal

import bgym
import hydra
from omegaconf import DictConfig
from pydantic import Field
from tapeagents.agent import Agent
from tapeagents.core import Action, Observation, StopStep, TapeMetadata, Thought
from tapeagents.core import Tape as BaseTape

from agentlab.agents.agent_args import AgentArgs

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

class ExtendedMetadata(TapeMetadata):
    name: str = ""
    task: dict = {}
    terminated: bool = False
    truncated: bool = False
    reward: float = 0.0
    attempt_number: int = 0
    other: dict = {}

class Tape(BaseTape):
    metadata: ExtendedMetadata = Field(default_factory=ExtendedMetadata)  # type: ignore

def load_config(config_name: str) -> DictConfig:
    with hydra.initialize(config_path="conf", version_base="1.1"):
        config = hydra.compose(config_name=config_name)
    return config

@dataclass
class TapeAgentArgs(AgentArgs):
    config: DictConfig = None  # type: ignore

    def make_agent(self) -> bgym.Agent:
        agent: Agent = hydra.utils.instantiate(self.config.agent)
        return TapeAgent(agent=agent)


Tape

from tapeagents.agent import Agent
from tapeagents.core import Action, Observation, StopStep, TapeMetadata, Thought
from tapeagents.core import Tape as BaseTape

from agentlab.agents.agent_args import AgentArgs

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

class ExtendedMetadata(TapeMetadata):
    name: str = ""
    task: dict = {}
    terminated: bool = False
    truncated: bool = False
    reward: float = 0.0
    attempt_number: int = 0
    other: dict = {}

class Tape(BaseTape):
    metadata: ExtendedMetadata = Field(default_factory=ExtendedMetadata)  # type: ignore

def load_config(config_name: str) -> DictConfig:
    with hydra.initialize(config_path="conf", version_base="1.1"):
        config = hydra.compose(config_name=config_name)
    return config

@dataclass
class TapeAgentArgs(AgentArgs):
    config: DictConfig = None  # type: ignore

    def make_agent(self) -> bgym.Agent:
        agent: Agent = hydra.utils.instantiate(self.config.agent)
        return TapeAgent(agent=agent)

@dataclass
class TapeAgentInfo(bgym.AgentInfo):
    thoughts: list[Thought] = None  # type: ignore


from agentlab.agents.agent_args import AgentArgs

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

class ExtendedMetadata(TapeMetadata):
    name: str = ""
    task: dict = {}
    terminated: bool = False
    truncated: bool = False
    reward: float = 0.0
    attempt_number: int = 0
    other: dict = {}

class Tape(BaseTape):
    metadata: ExtendedMetadata = Field(default_factory=ExtendedMetadata)  # type: ignore

def load_config(config_name: str) -> DictConfig:
    with hydra.initialize(config_path="conf", version_base="1.1"):
        config = hydra.compose(config_name=config_name)
    return config

@dataclass
class TapeAgentArgs(AgentArgs):
    config: DictConfig = None  # type: ignore

    def make_agent(self) -> bgym.Agent:
        agent: Agent = hydra.utils.instantiate(self.config.agent)
        return TapeAgent(agent=agent)

@dataclass
class TapeAgentInfo(bgym.AgentInfo):
    thoughts: list[Thought] = None  # type: ignore

class DictObservation(Observation):
    """
    Container for wrapping old dict observation into new Observation class.
    """


TapeAgentArgs

    name: str = ""
    task: dict = {}
    terminated: bool = False
    truncated: bool = False
    reward: float = 0.0
    attempt_number: int = 0
    other: dict = {}

class Tape(BaseTape):
    metadata: ExtendedMetadata = Field(default_factory=ExtendedMetadata)  # type: ignore

def load_config(config_name: str) -> DictConfig:
    with hydra.initialize(config_path="conf", version_base="1.1"):
        config = hydra.compose(config_name=config_name)
    return config

@dataclass
class TapeAgentArgs(AgentArgs):
    config: DictConfig = None  # type: ignore

    def make_agent(self) -> bgym.Agent:
        agent: Agent = hydra.utils.instantiate(self.config.agent)
        return TapeAgent(agent=agent)

@dataclass
class TapeAgentInfo(bgym.AgentInfo):
    thoughts: list[Thought] = None  # type: ignore

class DictObservation(Observation):
    """
    Container for wrapping old dict observation into new Observation class.
    """

    kind: Literal["dict_observation"] = "dict_observation"  # type: ignore
    content: str

class TapeAgent(bgym.Agent):
    agent: Agent
    tape: Tape


TapeAgentInfo

class Tape(BaseTape):
    metadata: ExtendedMetadata = Field(default_factory=ExtendedMetadata)  # type: ignore

def load_config(config_name: str) -> DictConfig:
    with hydra.initialize(config_path="conf", version_base="1.1"):
        config = hydra.compose(config_name=config_name)
    return config

@dataclass
class TapeAgentArgs(AgentArgs):
    config: DictConfig = None  # type: ignore

    def make_agent(self) -> bgym.Agent:
        agent: Agent = hydra.utils.instantiate(self.config.agent)
        return TapeAgent(agent=agent)

@dataclass
class TapeAgentInfo(bgym.AgentInfo):
    thoughts: list[Thought] = None  # type: ignore

class DictObservation(Observation):
    """
    Container for wrapping old dict observation into new Observation class.
    """

    kind: Literal["dict_observation"] = "dict_observation"  # type: ignore
    content: str

class TapeAgent(bgym.Agent):
    agent: Agent
    tape: Tape

    def __init__(self, agent: Agent):
        super().__init__()
        self.agent = agent
        self.tape = Tape(steps=[])


DictObservation

def load_config(config_name: str) -> DictConfig:
    with hydra.initialize(config_path="conf", version_base="1.1"):
        config = hydra.compose(config_name=config_name)
    return config

@dataclass
class TapeAgentArgs(AgentArgs):
    config: DictConfig = None  # type: ignore

    def make_agent(self) -> bgym.Agent:
        agent: Agent = hydra.utils.instantiate(self.config.agent)
        return TapeAgent(agent=agent)

@dataclass
class TapeAgentInfo(bgym.AgentInfo):
    thoughts: list[Thought] = None  # type: ignore

class DictObservation(Observation):
    """
    Container for wrapping old dict observation into new Observation class.
    """

    kind: Literal["dict_observation"] = "dict_observation"  # type: ignore
    content: str

class TapeAgent(bgym.Agent):
    agent: Agent
    tape: Tape

    def __init__(self, agent: Agent):
        super().__init__()
        self.agent = agent
        self.tape = Tape(steps=[])

    def obs_preprocessor(self, obs: Observation | list[Observation]) -> list[Observation]:
        if isinstance(obs, Observation):
            obs = [obs]
        assert isinstance(obs, list), f"Expected list of Observations, got {type(obs)}"
        logger.info(f"Observations: {[type(o).__name__ for o in obs]}")
        return obs

    def get_action(self, obs: Observation | list[Observation]) -> tuple[Action, TapeAgentInfo]:
        self.tape += obs  # type: ignore