import base64
import copy
import io
import re
from dataclasses import Field, asdict, dataclass
from typing import Dict, List

import bgym
import numpy as np
from browsergym.experiments.agent import AgentInfo
from PIL import Image

from agentlab.agents import dynamic_prompting as dp
from agentlab.agents.agent_utils import overlay_action
from agentlab.agents.generic_agent.generic_agent import GenericAgent, GenericAgentArgs
from agentlab.agents.generic_agent.generic_agent_prompt import MainPrompt
from agentlab.agents.hitl_agent.hint_labelling import (
    HintLabeling,
    HintLabelingInputs,
)
from agentlab.llm.llm_utils import (
    Discussion,
    HumanMessage,
    SystemMessage,
    img_to_base_64,
)
from agentlab.llm.tracking import cost_tracker_decorator

class CandidatesGeneration(dp.PromptElement):
    # Ask for multiple alternatives; each candidate must contain <think> and <action>.
    def __init__(self, hint: list[str] | None = None, n_candidates=3) -> None:
        self.hint = hint
        self.n_candidates = n_candidates
        self.hint_prompt = "\n".join(f"{i}. {c}" for i, c in enumerate(hint, 1)) if hint else ""
        super().__init__(True)
        self._prompt = [
            dict(
                type="text",
                text=f"""
    You are a web agent. Propose {self.n_candidates} alternative next steps for the current page.
    {('Use the Hints:' + self.hint_prompt) if self.hint else ""}\n
    Return EACH candidate wrapped as numbered tags:
    <candidate_generation_1>...</candidate_generation_1>
    <candidate_generation_2>...</candidate_generation_2>

    Inside every candidate you MUST include:
    <think>...why this action is appropriate now...</think>
    <action>...ONE atomic, executable action string...</action>

    Do not include any extra text outside the candidate tags.
    Use this format:
    <candidate_generation_1>
    <think>Explain why Candidate One is chosen</think>
    <action>Candidate One Action</action>
    </candidate_generation_1>

    <candidate_generation_2>
    <think>Explain why Candidate Two is chosen</think>
    <action>Candidate Two Action</action>
    </candidate_generation_2>
    # Example
    <candidate_generation_1>
    <think>The login button is visible and proceeding will reveal the auth form.</think>
    <action>click(role="button", name="Log in")</action>
    </candidate_generation_1>

    <candidate_generation_2>
    <think>User might need to enter email first; the email field is focused and visible.</think>
    <action>fill(bid="a112", text="user@example.com")</action>
    </candidate_generation_2>
    """,
            )
        ]

    # Regex patterns for numbered candidates only
    _NUM_BLOCK = re.compile(
        r"<\s*candidate[_ ]generation[_ ](?P<idx>[0-9]+)\s*>(?P<body>.*?)<\s*/\s*candidate[_ ]generation[_ ](?P=idx)\s*>",
        flags=re.IGNORECASE | re.DOTALL,
    )
    _THINK_PATTERN = re.compile(
        r"<\s*think\s*>(?P<think>.*?)<\s*/\s*think\s*>",
        flags=re.IGNORECASE | re.DOTALL,
    )
    _ACTION_PATTERN = re.compile(
        r"<\s*action\s*>(?P<action>.*?)<\s*/\s*action\s*>",
        flags=re.IGNORECASE | re.DOTALL,
    )

    def _parse_answer(self, text_answer: str) -> Dict[str, Dict[str, str]]:
        """Extract up to n_candidates candidates, using numbered tags only.

        Args:
            text_answer: The text response containing candidate generation tags.

        Returns:
            Dictionary mapping candidate names to their think and action content.
            Format: {"candidate_generation_1": {"think": "...", "action": "..."}, ...}
        """
        result = {
            f"candidate_generation_{i+1}": {"think": "", "action": ""}
            for i in range(self.n_candidates)
        }

        if not isinstance(text_answer, str):
            return result

        matches: List[re.Match] = list(self._NUM_BLOCK.finditer(text_answer))
        # Sort by numeric index
        matches_sorted = sorted(matches, key=lambda m: int(m.group("idx")))
        for i, m in enumerate(matches_sorted[: self.n_candidates]):
            body = m.group("body").strip()
            think_m = self._THINK_PATTERN.search(body)
            action_m = self._ACTION_PATTERN.search(body)
            result[f"candidate_generation_{i+1}"] = {
                "think": (think_m.group("think").strip() if think_m else ""),
                "action": (action_m.group("action").strip() if action_m else ""),
            }

        return result

@dataclass
class MultipleProposalGenericAgentArgs(GenericAgentArgs):

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

    def __post_init__(self):
        """Prefix subagent name with 'HITL-'."""
        super().__post_init__()
        if hasattr(self, "agent_name") and self.agent_name:
            self.agent_name = "HITL-" + self.agent_name

class MultipleProposalGenericAgent(GenericAgent):

    def __init__(
        self,
        chat_model_args,
        flags,
        max_retry: int = 4,
    ):
        super().__init__(chat_model_args, flags, max_retry)
        self.ui = None  # Single HintLabeling instance

    def get_candidate_generation(
        self,
        sys_prompt: SystemMessage,
        human_prompt: HumanMessage,
        hint: list[str] | None = None,
        n_candidates=3,
    ) -> tuple[Dict[str, Dict[str, str]], Discussion]:

        cg = CandidatesGeneration(hint=hint, n_candidates=n_candidates)
        candidates_prompt = HumanMessage(cg.prompt)
        chat_messages = Discussion([sys_prompt, human_prompt, candidates_prompt])
        output = self.chat_llm(chat_messages)
        candidates = cg._parse_answer(output["content"])
        self.step_n_human_intervention_rounds += 1
        msg_to_add_to_xray = Discussion([sys_prompt, human_prompt])

        return candidates, msg_to_add_to_xray

    @cost_tracker_decorator
    def get_action(self, obs):
        # reset vars
        step_hint = []
        self.step_n_human_intervention_rounds = 0
        self.obs_history.append(obs)
        main_prompt = MainPrompt(
            action_set=self.action_set,
            obs_history=self.obs_history,
            actions=self.actions,
            memories=self.memories,
            thoughts=self.thoughts,
            previous_plan=self.plan,
            step=self.plan_step,
            flags=self.flags,
        )

        max_prompt_tokens, max_trunc_itr = self._get_maxes()

        system_prompt = SystemMessage(dp.SystemPrompt().prompt)

        human_prompt = dp.fit_tokens(
            shrinkable=main_prompt,
            max_prompt_tokens=max_prompt_tokens,
            model_name=self.chat_model_args.model_name,
            max_iterations=max_trunc_itr,
            additional_prompts=system_prompt,
        )
        # Initialize UI once outside the loop
        if self.ui is None:
            self.ui = HintLabeling(headless=False)
            # Show initial waiting state
            initial_inputs = HintLabelingInputs(
                goal=(
                    obs.get("goal_object", [{}])[0].get("text", "")
                    if obs.get("goal_object")
                    else ""
                ),
                error_feedback="",
                screenshot=(img_to_base_64(obs["screenshot"]) if "screenshot" in obs else ""),
                screenshots=[],  # no overlay screenshots yet
                axtree=obs.get("axtree_txt", ""),
                hints=[],
                suggestions=[],  # no suggestions yet
            )
            self.ui.update_context(initial_inputs)

        # Generate first candidates
        candidates, chat_messages = self.get_candidate_generation(
            sys_prompt=system_prompt,
            human_prompt=human_prompt,
            hint=step_hint if step_hint else None,
        )
        suggestions = [
            {
                "id": key.split("_")[-1],
                "action": candidate["action"],
                "think": candidate["think"],
            }
            for key, candidate in candidates.items()
        ]
        # List of Images as base64 - create overlay screenshots for each suggestion
        screenshots = [overlay_action(obs, choice["action"]) for choice in suggestions]

        while True:
            try:
                hint_labeling_inputs = HintLabelingInputs(
                    goal=(
                        obs.get("goal_object", [{}])[0].get("text", "")
                        if obs.get("goal_object")
                        else ""
                    ),
                    error_feedback=obs.get("last_action_error", ""),
                    screenshot=(img_to_base_64(obs["screenshot"]) if "screenshot" in obs else ""),
                    screenshots=screenshots,  # list of overlay screenshots for hover
                    axtree=obs.get("axtree_txt", ""),
                    hints=step_hint,
                    suggestions=suggestions,
                )

                self.ui.update_context(hint_labeling_inputs)
                response = self.ui.wait_for_response(timeout=600)

                if response["type"] == "reprompt":
                    new_hints = response["payload"].get("hints", [])
                    step_hint = list(new_hints) if isinstance(new_hints, list) else step_hint
                    candidates, chat_messages = self.get_candidate_generation(
                        sys_prompt=system_prompt,
                        human_prompt=human_prompt,
                        hint=step_hint if step_hint else None,
                    )
                    suggestions = [
                        {
                            "id": key.split("_")[-1],
                            "action": candidate["action"],
                            "think": candidate["think"],
                        }
                        for key, candidate in candidates.items()
                    ]
                    # Regenerate screenshots for new suggestions
                    screenshots = [overlay_action(obs, choice["action"]) for choice in suggestions]
                    # Continue the loop to show new suggestions
                elif response["type"] == "step":
                    selected_action = response["payload"]["action"]
                    choice_idx = None
                    for i, candidate in enumerate(suggestions, 1):
                        if candidate["action"] == selected_action:
                            choice_idx = i
                            break
                    if choice_idx is None:
                        choice_idx = 1
                    ans_dict = candidates[f"candidate_generation_{choice_idx}"]
                    break
                else:
                    ans_dict = candidates["candidate_generation_1"]
                    break

            except KeyboardInterrupt:
                print("User cancelled the operation")
                if self.ui:
                    self.ui.close()
                raise
            except Exception as e:
                print(f"Error in human intervention UI: {e}")
                if self.ui:
                    self.ui.close()
                    self.ui = None
                # Raise exception instead of falling back to console input
                raise RuntimeError(f"Human intervention UI failed: {e}") from e

        # TODO: Refactor as discussed with ALAC.
        stats = self.chat_llm.get_stats()
        self.plan = ans_dict.get("plan", self.plan)
        self.plan_step = ans_dict.get("step", self.plan_step)
        self.actions.append(ans_dict["action"])
        self.memories.append(ans_dict.get("memory", None))
        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),
                "step_hints": step_hint,
                "n_human_intervention_rounds": self.step_n_human_intervention_rounds,
                "candidates": candidates,
                "suggestions": suggestions,
            },
        )
        return ans_dict["action"], agent_info

def get_base_agent(llm_config):
    """Creates and returns a MultipleProposalGenericAgentArgs instance with
    specified LLM configuration from CHAT_MODEL_ARGS_DICT.

    Args:
        llm_config: The LLM configuration key to use from CHAT_MODEL_ARGS_DICT.

    Returns:
        MultipleProposalGenericAgentArgs: Configured agent arguments instance.
    """

    from agentlab.agents.generic_agent.tmlr_config import BASE_FLAGS
    from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT

    return MultipleProposalGenericAgentArgs(
        chat_model_args=CHAT_MODEL_ARGS_DICT[llm_config],
        flags=BASE_FLAGS,
    )

HUMAN_GUIDED_GENERIC_AGENT = get_base_agent("openai/gpt-5-mini-2025-08-07")

if __name__ == "__main__":
    import logging

    from agentlab.agents.hitl_agent.generic_human_guided_agent import (
        HUMAN_GUIDED_GENERIC_AGENT,
    )
    from agentlab.experiments.study import Study

    agent_configs = [HUMAN_GUIDED_GENERIC_AGENT]
    benchmark = bgym.DEFAULT_BENCHMARKS["miniwob"]()
    benchmark = benchmark.subset_from_glob("task_name", "*book*")
    benchmark.env_args_list = benchmark.env_args_list[3:4]

    for env_args in benchmark.env_args_list:
        env_args.max_steps = 100  # max human steps
        env_args.headless = True

    Study(agent_configs, benchmark, logging_level=logging.WARNING).run(
        n_jobs=1,
        parallel_backend="sequential",
        n_relaunch=1,
    )


CandidatesGeneration

from browsergym.experiments.agent import AgentInfo
from PIL import Image

from agentlab.agents import dynamic_prompting as dp
from agentlab.agents.agent_utils import overlay_action
from agentlab.agents.generic_agent.generic_agent import GenericAgent, GenericAgentArgs
from agentlab.agents.generic_agent.generic_agent_prompt import MainPrompt
from agentlab.agents.hitl_agent.hint_labelling import (
    HintLabeling,
    HintLabelingInputs,
)
from agentlab.llm.llm_utils import (
    Discussion,
    HumanMessage,
    SystemMessage,
    img_to_base_64,
)
from agentlab.llm.tracking import cost_tracker_decorator

class CandidatesGeneration(dp.PromptElement):
    # Ask for multiple alternatives; each candidate must contain <think> and <action>.
    def __init__(self, hint: list[str] | None = None, n_candidates=3) -> None:
        self.hint = hint
        self.n_candidates = n_candidates
        self.hint_prompt = "\n".join(f"{i}. {c}" for i, c in enumerate(hint, 1)) if hint else ""
        super().__init__(True)
        self._prompt = [
            dict(
                type="text",
                text=f"""
    You are a web agent. Propose {self.n_candidates} alternative next steps for the current page.
    {('Use the Hints:' + self.hint_prompt) if self.hint else ""}\n
    Return EACH candidate wrapped as numbered tags:
    <candidate_generation_1>...</candidate_generation_1>
    <candidate_generation_2>...</candidate_generation_2>

    Inside every candidate you MUST include:
    <think>...why this action is appropriate now...</think>
    <action>...ONE atomic, executable action string...</action>

    Do not include any extra text outside the candidate tags.
    Use this format:
    <candidate_generation_1>
    <think>Explain why Candidate One is chosen</think>
    <action>Candidate One Action</action>
    </candidate_generation_1>

    <candidate_generation_2>
    <think>Explain why Candidate Two is chosen</think>
    <action>Candidate Two Action</action>
    </candidate_generation_2>
    # Example
    <candidate_generation_1>
    <think>The login button is visible and proceeding will reveal the auth form.</think>
    <action>click(role="button", name="Log in")</action>
    </candidate_generation_1>

    <candidate_generation_2>
    <think>User might need to enter email first; the email field is focused and visible.</think>
    <action>fill(bid="a112", text="user@example.com")</action>
    </candidate_generation_2>
    """,
            )
        ]

    # Regex patterns for numbered candidates only
    _NUM_BLOCK = re.compile(
        r"<\s*candidate[_ ]generation[_ ](?P<idx>[0-9]+)\s*>(?P<body>.*?)<\s*/\s*candidate[_ ]generation[_ ](?P=idx)\s*>",
        flags=re.IGNORECASE | re.DOTALL,
    )
    _THINK_PATTERN = re.compile(
        r"<\s*think\s*>(?P<think>.*?)<\s*/\s*think\s*>",
        flags=re.IGNORECASE | re.DOTALL,
    )
    _ACTION_PATTERN = re.compile(
        r"<\s*action\s*>(?P<action>.*?)<\s*/\s*action\s*>",
        flags=re.IGNORECASE | re.DOTALL,
    )

    def _parse_answer(self, text_answer: str) -> Dict[str, Dict[str, str]]:
        """Extract up to n_candidates candidates, using numbered tags only.

        Args:
            text_answer: The text response containing candidate generation tags.

        Returns:
            Dictionary mapping candidate names to their think and action content.
            Format: {"candidate_generation_1": {"think": "...", "action": "..."}, ...}
        """
        result = {
            f"candidate_generation_{i+1}": {"think": "", "action": ""}
            for i in range(self.n_candidates)
        }

        if not isinstance(text_answer, str):
            return result

        matches: List[re.Match] = list(self._NUM_BLOCK.finditer(text_answer))
        # Sort by numeric index
        matches_sorted = sorted(matches, key=lambda m: int(m.group("idx")))
        for i, m in enumerate(matches_sorted[: self.n_candidates]):
            body = m.group("body").strip()
            think_m = self._THINK_PATTERN.search(body)
            action_m = self._ACTION_PATTERN.search(body)
            result[f"candidate_generation_{i+1}"] = {
                "think": (think_m.group("think").strip() if think_m else ""),
                "action": (action_m.group("action").strip() if action_m else ""),
            }

        return result

@dataclass
class MultipleProposalGenericAgentArgs(GenericAgentArgs):

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

    def __post_init__(self):
        """Prefix subagent name with 'HITL-'."""
        super().__post_init__()
        if hasattr(self, "agent_name") and self.agent_name:
            self.agent_name = "HITL-" + self.agent_name

class MultipleProposalGenericAgent(GenericAgent):

    def __init__(


MultipleProposalGenericAgentArgs

        if not isinstance(text_answer, str):
            return result

        matches: List[re.Match] = list(self._NUM_BLOCK.finditer(text_answer))
        # Sort by numeric index
        matches_sorted = sorted(matches, key=lambda m: int(m.group("idx")))
        for i, m in enumerate(matches_sorted[: self.n_candidates]):
            body = m.group("body").strip()
            think_m = self._THINK_PATTERN.search(body)
            action_m = self._ACTION_PATTERN.search(body)
            result[f"candidate_generation_{i+1}"] = {
                "think": (think_m.group("think").strip() if think_m else ""),
                "action": (action_m.group("action").strip() if action_m else ""),
            }

        return result

@dataclass
class MultipleProposalGenericAgentArgs(GenericAgentArgs):

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

    def __post_init__(self):
        """Prefix subagent name with 'HITL-'."""
        super().__post_init__()
        if hasattr(self, "agent_name") and self.agent_name:
            self.agent_name = "HITL-" + self.agent_name

class MultipleProposalGenericAgent(GenericAgent):

    def __init__(
        self,
        chat_model_args,
        flags,
        max_retry: int = 4,
    ):
        super().__init__(chat_model_args, flags, max_retry)
        self.ui = None  # Single HintLabeling instance

    def get_candidate_generation(
        self,
        sys_prompt: SystemMessage,
        human_prompt: HumanMessage,
        hint: list[str] | None = None,
        n_candidates=3,
    ) -> tuple[Dict[str, Dict[str, str]], Discussion]:


MultipleProposalGenericAgent

            }

        return result

@dataclass
class MultipleProposalGenericAgentArgs(GenericAgentArgs):

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

    def __post_init__(self):
        """Prefix subagent name with 'HITL-'."""
        super().__post_init__()
        if hasattr(self, "agent_name") and self.agent_name:
            self.agent_name = "HITL-" + self.agent_name

class MultipleProposalGenericAgent(GenericAgent):

    def __init__(
        self,
        chat_model_args,
        flags,
        max_retry: int = 4,
    ):
        super().__init__(chat_model_args, flags, max_retry)
        self.ui = None  # Single HintLabeling instance

    def get_candidate_generation(
        self,
        sys_prompt: SystemMessage,
        human_prompt: HumanMessage,
        hint: list[str] | None = None,
        n_candidates=3,
    ) -> tuple[Dict[str, Dict[str, str]], Discussion]:

        cg = CandidatesGeneration(hint=hint, n_candidates=n_candidates)
        candidates_prompt = HumanMessage(cg.prompt)
        chat_messages = Discussion([sys_prompt, human_prompt, candidates_prompt])
        output = self.chat_llm(chat_messages)
        candidates = cg._parse_answer(output["content"])
        self.step_n_human_intervention_rounds += 1
        msg_to_add_to_xray = Discussion([sys_prompt, human_prompt])

        return candidates, msg_to_add_to_xray

    @cost_tracker_decorator
    def get_action(self, obs):
        # reset vars
        step_hint = []
        self.step_n_human_intervention_rounds = 0
        self.obs_history.append(obs)
        main_prompt = MainPrompt(
            action_set=self.action_set,
            obs_history=self.obs_history,
            actions=self.actions,
            memories=self.memories,
            thoughts=self.thoughts,
            previous_plan=self.plan,
            step=self.plan_step,
            flags=self.flags,
        )

        max_prompt_tokens, max_trunc_itr = self._get_maxes()

        system_prompt = SystemMessage(dp.SystemPrompt().prompt)

        human_prompt = dp.fit_tokens(
            shrinkable=main_prompt,
            max_prompt_tokens=max_prompt_tokens,
            model_name=self.chat_model_args.model_name,
            max_iterations=max_trunc_itr,
            additional_prompts=system_prompt,
        )
        # Initialize UI once outside the loop
        if self.ui is None:
            self.ui = HintLabeling(headless=False)
            # Show initial waiting state
            initial_inputs = HintLabelingInputs(
                goal=(
                    obs.get("goal_object", [{}])[0].get("text", "")
                    if obs.get("goal_object")
                    else ""
                ),
                error_feedback="",
                screenshot=(img_to_base_64(obs["screenshot"]) if "screenshot" in obs else ""),
                screenshots=[],  # no overlay screenshots yet
                axtree=obs.get("axtree_txt", ""),
                hints=[],
                suggestions=[],  # no suggestions yet
            )
            self.ui.update_context(initial_inputs)

        # Generate first candidates
        candidates, chat_messages = self.get_candidate_generation(
            sys_prompt=system_prompt,
            human_prompt=human_prompt,
            hint=step_hint if step_hint else None,
        )
        suggestions = [
            {
                "id": key.split("_")[-1],
                "action": candidate["action"],
                "think": candidate["think"],
            }
            for key, candidate in candidates.items()
        ]
        # List of Images as base64 - create overlay screenshots for each suggestion
        screenshots = [overlay_action(obs, choice["action"]) for choice in suggestions]

        while True:
            try:
                hint_labeling_inputs = HintLabelingInputs(
                    goal=(
                        obs.get("goal_object", [{}])[0].get("text", "")
                        if obs.get("goal_object")
                        else ""
                    ),
                    error_feedback=obs.get("last_action_error", ""),
                    screenshot=(img_to_base_64(obs["screenshot"]) if "screenshot" in obs else ""),
                    screenshots=screenshots,  # list of overlay screenshots for hover
                    axtree=obs.get("axtree_txt", ""),
                    hints=step_hint,
                    suggestions=suggestions,
                )

                self.ui.update_context(hint_labeling_inputs)
                response = self.ui.wait_for_response(timeout=600)

                if response["type"] == "reprompt":
                    new_hints = response["payload"].get("hints", [])
                    step_hint = list(new_hints) if isinstance(new_hints, list) else step_hint
                    candidates, chat_messages = self.get_candidate_generation(
                        sys_prompt=system_prompt,
                        human_prompt=human_prompt,
                        hint=step_hint if step_hint else None,
                    )
                    suggestions = [
                        {
                            "id": key.split("_")[-1],
                            "action": candidate["action"],
                            "think": candidate["think"],
                        }
                        for key, candidate in candidates.items()
                    ]
                    # Regenerate screenshots for new suggestions
                    screenshots = [overlay_action(obs, choice["action"]) for choice in suggestions]
                    # Continue the loop to show new suggestions
                elif response["type"] == "step":
                    selected_action = response["payload"]["action"]
                    choice_idx = None
                    for i, candidate in enumerate(suggestions, 1):
                        if candidate["action"] == selected_action:
                            choice_idx = i
                            break
                    if choice_idx is None:
                        choice_idx = 1
                    ans_dict = candidates[f"candidate_generation_{choice_idx}"]
                    break
                else:
                    ans_dict = candidates["candidate_generation_1"]
                    break

            except KeyboardInterrupt:
                print("User cancelled the operation")
                if self.ui:
                    self.ui.close()
                raise
            except Exception as e:
                print(f"Error in human intervention UI: {e}")
                if self.ui:
                    self.ui.close()
                    self.ui = None
                # Raise exception instead of falling back to console input
                raise RuntimeError(f"Human intervention UI failed: {e}") from e

        # TODO: Refactor as discussed with ALAC.
        stats = self.chat_llm.get_stats()
        self.plan = ans_dict.get("plan", self.plan)
        self.plan_step = ans_dict.get("step", self.plan_step)
        self.actions.append(ans_dict["action"])
        self.memories.append(ans_dict.get("memory", None))
        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),
                "step_hints": step_hint,
                "n_human_intervention_rounds": self.step_n_human_intervention_rounds,
                "candidates": candidates,
                "suggestions": suggestions,
            },
        )
        return ans_dict["action"], agent_info

def get_base_agent(llm_config):
    """Creates and returns a MultipleProposalGenericAgentArgs instance with
    specified LLM configuration from CHAT_MODEL_ARGS_DICT.

    Args:
        llm_config: The LLM configuration key to use from CHAT_MODEL_ARGS_DICT.

    Returns:
        MultipleProposalGenericAgentArgs: Configured agent arguments instance.
    """

    from agentlab.agents.generic_agent.tmlr_config import BASE_FLAGS
    from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT

    return MultipleProposalGenericAgentArgs(
        chat_model_args=CHAT_MODEL_ARGS_DICT[llm_config],
        flags=BASE_FLAGS,
    )


        self.plan = ans_dict.get("plan", self.plan)
        self.plan_step = ans_dict.get("step", self.plan_step)
        self.actions.append(ans_dict["action"])
        self.memories.append(ans_dict.get("memory", None))
        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),
                "step_hints": step_hint,
                "n_human_intervention_rounds": self.step_n_human_intervention_rounds,
                "candidates": candidates,
                "suggestions": suggestions,
            },
        )
        return ans_dict["action"], agent_info

def get_base_agent(llm_config):
    """Creates and returns a MultipleProposalGenericAgentArgs instance with
    specified LLM configuration from CHAT_MODEL_ARGS_DICT.

    Args:
        llm_config: The LLM configuration key to use from CHAT_MODEL_ARGS_DICT.

    Returns:
        MultipleProposalGenericAgentArgs: Configured agent arguments instance.
    """

    from agentlab.agents.generic_agent.tmlr_config import BASE_FLAGS
    from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT

    return MultipleProposalGenericAgentArgs(
        chat_model_args=CHAT_MODEL_ARGS_DICT[llm_config],
        flags=BASE_FLAGS,
    )

HUMAN_GUIDED_GENERIC_AGENT = get_base_agent("openai/gpt-5-mini-2025-08-07")

if __name__ == "__main__":
    import logging

    from agentlab.agents.hitl_agent.generic_human_guided_agent import (
        HUMAN_GUIDED_GENERIC_AGENT,
    )
    from agentlab.experiments.study import Study

    agent_configs = [HUMAN_GUIDED_GENERIC_AGENT]
    benchmark = bgym.DEFAULT_BENCHMARKS["miniwob"]()
    benchmark = benchmark.subset_from_glob("task_name", "*book*")
    benchmark.env_args_list = benchmark.env_args_list[3:4]

    for env_args in benchmark.env_args_list:
        env_args.max_steps = 100  # max human steps
        env_args.headless = True


        return result

@dataclass
class MultipleProposalGenericAgentArgs(GenericAgentArgs):

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

    def __post_init__(self):
        """Prefix subagent name with 'HITL-'."""
        super().__post_init__()
        if hasattr(self, "agent_name") and self.agent_name:
            self.agent_name = "HITL-" + self.agent_name

class MultipleProposalGenericAgent(GenericAgent):

    def __init__(
        self,
        chat_model_args,
        flags,
        max_retry: int = 4,
    ):
        super().__init__(chat_model_args, flags, max_retry)
        self.ui = None  # Single HintLabeling instance

    def get_candidate_generation(
        self,
        sys_prompt: SystemMessage,
        human_prompt: HumanMessage,
        hint: list[str] | None = None,
        n_candidates=3,
    ) -> tuple[Dict[str, Dict[str, str]], Discussion]:

        cg = CandidatesGeneration(hint=hint, n_candidates=n_candidates)
        candidates_prompt = HumanMessage(cg.prompt)
        chat_messages = Discussion([sys_prompt, human_prompt, candidates_prompt])
        output = self.chat_llm(chat_messages)
        candidates = cg._parse_answer(output["content"])
        self.step_n_human_intervention_rounds += 1
        msg_to_add_to_xray = Discussion([sys_prompt, human_prompt])

        return candidates, msg_to_add_to_xray

    @cost_tracker_decorator


    <action>fill(bid="a112", text="user@example.com")</action>
    </candidate_generation_2>
    """,
            )
        ]

    # Regex patterns for numbered candidates only
    _NUM_BLOCK = re.compile(
        r"<\s*candidate[_ ]generation[_ ](?P<idx>[0-9]+)\s*>(?P<body>.*?)<\s*/\s*candidate[_ ]generation[_ ](?P=idx)\s*>",
        flags=re.IGNORECASE | re.DOTALL,
    )
    _THINK_PATTERN = re.compile(
        r"<\s*think\s*>(?P<think>.*?)<\s*/\s*think\s*>",
        flags=re.IGNORECASE | re.DOTALL,
    )
    _ACTION_PATTERN = re.compile(
        r"<\s*action\s*>(?P<action>.*?)<\s*/\s*action\s*>",
        flags=re.IGNORECASE | re.DOTALL,
    )

    def _parse_answer(self, text_answer: str) -> Dict[str, Dict[str, str]]:
        """Extract up to n_candidates candidates, using numbered tags only.

        Args:
            text_answer: The text response containing candidate generation tags.

        Returns:
            Dictionary mapping candidate names to their think and action content.
            Format: {"candidate_generation_1": {"think": "...", "action": "..."}, ...}
        """
        result = {
            f"candidate_generation_{i+1}": {"think": "", "action": ""}
            for i in range(self.n_candidates)
        }

        if not isinstance(text_answer, str):
            return result

        matches: List[re.Match] = list(self._NUM_BLOCK.finditer(text_answer))
        # Sort by numeric index
        matches_sorted = sorted(matches, key=lambda m: int(m.group("idx")))
        for i, m in enumerate(matches_sorted[: self.n_candidates]):
            body = m.group("body").strip()
            think_m = self._THINK_PATTERN.search(body)
            action_m = self._ACTION_PATTERN.search(body)
            result[f"candidate_generation_{i+1}"] = {
                "think": (think_m.group("think").strip() if think_m else ""),
                "action": (action_m.group("action").strip() if action_m else ""),
            }

        return result

@dataclass
class MultipleProposalGenericAgentArgs(GenericAgentArgs):

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

    def __post_init__(self):
        """Prefix subagent name with 'HITL-'."""
        super().__post_init__()
        if hasattr(self, "agent_name") and self.agent_name:
            self.agent_name = "HITL-" + self.agent_name

class MultipleProposalGenericAgent(GenericAgent):

    def __init__(


            return result

        matches: List[re.Match] = list(self._NUM_BLOCK.finditer(text_answer))
        # Sort by numeric index
        matches_sorted = sorted(matches, key=lambda m: int(m.group("idx")))
        for i, m in enumerate(matches_sorted[: self.n_candidates]):
            body = m.group("body").strip()
            think_m = self._THINK_PATTERN.search(body)
            action_m = self._ACTION_PATTERN.search(body)
            result[f"candidate_generation_{i+1}"] = {
                "think": (think_m.group("think").strip() if think_m else ""),
                "action": (action_m.group("action").strip() if action_m else ""),
            }

        return result

@dataclass
class MultipleProposalGenericAgentArgs(GenericAgentArgs):

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

    def __post_init__(self):
        """Prefix subagent name with 'HITL-'."""
        super().__post_init__()
        if hasattr(self, "agent_name") and self.agent_name:
            self.agent_name = "HITL-" + self.agent_name

class MultipleProposalGenericAgent(GenericAgent):

    def __init__(
        self,
        chat_model_args,
        flags,
        max_retry: int = 4,
    ):
        super().__init__(chat_model_args, flags, max_retry)
        self.ui = None  # Single HintLabeling instance

    def get_candidate_generation(


        for i, m in enumerate(matches_sorted[: self.n_candidates]):
            body = m.group("body").strip()
            think_m = self._THINK_PATTERN.search(body)
            action_m = self._ACTION_PATTERN.search(body)
            result[f"candidate_generation_{i+1}"] = {
                "think": (think_m.group("think").strip() if think_m else ""),
                "action": (action_m.group("action").strip() if action_m else ""),
            }

        return result

@dataclass
class MultipleProposalGenericAgentArgs(GenericAgentArgs):

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

    def __post_init__(self):
        """Prefix subagent name with 'HITL-'."""
        super().__post_init__()
        if hasattr(self, "agent_name") and self.agent_name:
            self.agent_name = "HITL-" + self.agent_name

class MultipleProposalGenericAgent(GenericAgent):

    def __init__(
        self,
        chat_model_args,
        flags,
        max_retry: int = 4,
    ):
        super().__init__(chat_model_args, flags, max_retry)
        self.ui = None  # Single HintLabeling instance

    def get_candidate_generation(
        self,
        sys_prompt: SystemMessage,
        human_prompt: HumanMessage,
        hint: list[str] | None = None,
        n_candidates=3,
    ) -> tuple[Dict[str, Dict[str, str]], Discussion]:


        )

    def __post_init__(self):
        """Prefix subagent name with 'HITL-'."""
        super().__post_init__()
        if hasattr(self, "agent_name") and self.agent_name:
            self.agent_name = "HITL-" + self.agent_name

class MultipleProposalGenericAgent(GenericAgent):

    def __init__(
        self,
        chat_model_args,
        flags,
        max_retry: int = 4,
    ):
        super().__init__(chat_model_args, flags, max_retry)
        self.ui = None  # Single HintLabeling instance

    def get_candidate_generation(
        self,
        sys_prompt: SystemMessage,
        human_prompt: HumanMessage,
        hint: list[str] | None = None,
        n_candidates=3,
    ) -> tuple[Dict[str, Dict[str, str]], Discussion]:

        cg = CandidatesGeneration(hint=hint, n_candidates=n_candidates)
        candidates_prompt = HumanMessage(cg.prompt)
        chat_messages = Discussion([sys_prompt, human_prompt, candidates_prompt])
        output = self.chat_llm(chat_messages)
        candidates = cg._parse_answer(output["content"])
        self.step_n_human_intervention_rounds += 1
        msg_to_add_to_xray = Discussion([sys_prompt, human_prompt])

        return candidates, msg_to_add_to_xray

    @cost_tracker_decorator
    def get_action(self, obs):
        # reset vars
        step_hint = []
        self.step_n_human_intervention_rounds = 0
        self.obs_history.append(obs)
        main_prompt = MainPrompt(
            action_set=self.action_set,
            obs_history=self.obs_history,
            actions=self.actions,
            memories=self.memories,
            thoughts=self.thoughts,
            previous_plan=self.plan,
            step=self.plan_step,
            flags=self.flags,
        )

        max_prompt_tokens, max_trunc_itr = self._get_maxes()


    def get_candidate_generation(
        self,
        sys_prompt: SystemMessage,
        human_prompt: HumanMessage,
        hint: list[str] | None = None,
        n_candidates=3,
    ) -> tuple[Dict[str, Dict[str, str]], Discussion]:

        cg = CandidatesGeneration(hint=hint, n_candidates=n_candidates)
        candidates_prompt = HumanMessage(cg.prompt)
        chat_messages = Discussion([sys_prompt, human_prompt, candidates_prompt])
        output = self.chat_llm(chat_messages)
        candidates = cg._parse_answer(output["content"])
        self.step_n_human_intervention_rounds += 1
        msg_to_add_to_xray = Discussion([sys_prompt, human_prompt])

        return candidates, msg_to_add_to_xray

    @cost_tracker_decorator
    def get_action(self, obs):
        # reset vars
        step_hint = []
        self.step_n_human_intervention_rounds = 0
        self.obs_history.append(obs)
        main_prompt = MainPrompt(
            action_set=self.action_set,
            obs_history=self.obs_history,
            actions=self.actions,
            memories=self.memories,
            thoughts=self.thoughts,
            previous_plan=self.plan,
            step=self.plan_step,
            flags=self.flags,
        )

        max_prompt_tokens, max_trunc_itr = self._get_maxes()

        system_prompt = SystemMessage(dp.SystemPrompt().prompt)

        human_prompt = dp.fit_tokens(
            shrinkable=main_prompt,
            max_prompt_tokens=max_prompt_tokens,
            model_name=self.chat_model_args.model_name,
            max_iterations=max_trunc_itr,
            additional_prompts=system_prompt,
        )
        # Initialize UI once outside the loop
        if self.ui is None:
            self.ui = HintLabeling(headless=False)
            # Show initial waiting state
            initial_inputs = HintLabelingInputs(
                goal=(
                    obs.get("goal_object", [{}])[0].get("text", "")
                    if obs.get("goal_object")
                    else ""
                ),
                error_feedback="",
                screenshot=(img_to_base_64(obs["screenshot"]) if "screenshot" in obs else ""),
                screenshots=[],  # no overlay screenshots yet
                axtree=obs.get("axtree_txt", ""),
                hints=[],
                suggestions=[],  # no suggestions yet
            )
            self.ui.update_context(initial_inputs)

        # Generate first candidates
        candidates, chat_messages = self.get_candidate_generation(
            sys_prompt=system_prompt,
            human_prompt=human_prompt,
            hint=step_hint if step_hint else None,
        )
        suggestions = [
            {
                "id": key.split("_")[-1],
                "action": candidate["action"],
                "think": candidate["think"],
            }
            for key, candidate in candidates.items()
        ]
        # List of Images as base64 - create overlay screenshots for each suggestion
        screenshots = [overlay_action(obs, choice["action"]) for choice in suggestions]

        while True:
            try:
                hint_labeling_inputs = HintLabelingInputs(
                    goal=(
                        obs.get("goal_object", [{}])[0].get("text", "")
                        if obs.get("goal_object")
                        else ""
                    ),
                    error_feedback=obs.get("last_action_error", ""),
                    screenshot=(img_to_base_64(obs["screenshot"]) if "screenshot" in obs else ""),
                    screenshots=screenshots,  # list of overlay screenshots for hover
                    axtree=obs.get("axtree_txt", ""),
                    hints=step_hint,
                    suggestions=suggestions,
                )

                self.ui.update_context(hint_labeling_inputs)
                response = self.ui.wait_for_response(timeout=600)

                if response["type"] == "reprompt":
                    new_hints = response["payload"].get("hints", [])
                    step_hint = list(new_hints) if isinstance(new_hints, list) else step_hint
                    candidates, chat_messages = self.get_candidate_generation(
                        sys_prompt=system_prompt,
                        human_prompt=human_prompt,
                        hint=step_hint if step_hint else None,
                    )
                    suggestions = [
                        {
                            "id": key.split("_")[-1],
                            "action": candidate["action"],
                            "think": candidate["think"],
                        }
                        for key, candidate in candidates.items()
                    ]
                    # Regenerate screenshots for new suggestions
                    screenshots = [overlay_action(obs, choice["action"]) for choice in suggestions]
                    # Continue the loop to show new suggestions
                elif response["type"] == "step":
                    selected_action = response["payload"]["action"]
                    choice_idx = None
                    for i, candidate in enumerate(suggestions, 1):
                        if candidate["action"] == selected_action:
                            choice_idx = i
                            break
                    if choice_idx is None:
                        choice_idx = 1
                    ans_dict = candidates[f"candidate_generation_{choice_idx}"]
                    break
                else:
                    ans_dict = candidates["candidate_generation_1"]
                    break

            except KeyboardInterrupt:
                print("User cancelled the operation")
                if self.ui:
                    self.ui.close()
                raise
            except Exception as e:
                print(f"Error in human intervention UI: {e}")
                if self.ui:
                    self.ui.close()
                    self.ui = None
                # Raise exception instead of falling back to console input
                raise RuntimeError(f"Human intervention UI failed: {e}") from e

        # TODO: Refactor as discussed with ALAC.
        stats = self.chat_llm.get_stats()
        self.plan = ans_dict.get("plan", self.plan)
        self.plan_step = ans_dict.get("step", self.plan_step)
        self.actions.append(ans_dict["action"])
        self.memories.append(ans_dict.get("memory", None))
        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),
                "step_hints": step_hint,
                "n_human_intervention_rounds": self.step_n_human_intervention_rounds,
                "candidates": candidates,
                "suggestions": suggestions,
            },
        )
        return ans_dict["action"], agent_info

def get_base_agent(llm_config):
    """Creates and returns a MultipleProposalGenericAgentArgs instance with
    specified LLM configuration from CHAT_MODEL_ARGS_DICT.

    Args:
        llm_config: The LLM configuration key to use from CHAT_MODEL_ARGS_DICT.

    Returns:
        MultipleProposalGenericAgentArgs: Configured agent arguments instance.
    """

    from agentlab.agents.generic_agent.tmlr_config import BASE_FLAGS
    from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT

    return MultipleProposalGenericAgentArgs(
        chat_model_args=CHAT_MODEL_ARGS_DICT[llm_config],
        flags=BASE_FLAGS,
    )


import fnmatch
import json
import logging
import os
import random
import time
from abc import ABC, abstractmethod
from collections import defaultdict
from copy import copy
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Literal

import bgym
import numpy as np
import pandas as pd
import requests
from bgym import Benchmark as BgymBenchmark
from browsergym.core.observation import extract_screenshot
from browsergym.utils.obs import (
    flatten_axtree_to_str,
    flatten_dom_to_str,
    overlay_som,
    prune_html,
)

from agentlab.agents.agent_args import AgentArgs
from agentlab.benchmarks.abstract_env import AbstractBenchmark as AgentLabBenchmark
from agentlab.benchmarks.osworld import OSWorldActionSet
from agentlab.llm.base_api import BaseModelArgs
from agentlab.llm.llm_utils import image_to_png_base64_url
from agentlab.llm.response_api import (
    APIPayload,
    ClaudeResponseModelArgs,
    LLMOutput,
    MessageBuilder,
    OpenAIChatModelArgs,
    OpenAIResponseModelArgs,
    OpenRouterModelArgs,
    ToolCalls,
)
from agentlab.llm.tracking import cost_tracker_decorator

logger = logging.getLogger(__name__)

@dataclass
class Block(ABC):
    def _init(self):
        """Initialize the block."""
        pass

    def make(self) -> "Block":
        """Returns a copy so the init can start adding some stuff to `self` without changing the
        original datatclass that should only contain a config.
        The aim is avoid having 2 calss definition for each block, e.g. Block and BlockArgs.

        Returns:
            Block: A copy of the current block instance with initialization applied.
        """
        block = self.__class__(**asdict(self))
        block._init()
        return block

    @abstractmethod
    def apply(self, llm, messages: list[MessageBuilder], **kwargs):
        pass

@dataclass
class MsgGroup:
    name: str = None
    messages: list[MessageBuilder] = field(default_factory=list)
    summary: MessageBuilder = None

class StructuredDiscussion:
    """
    A structured discussion that groups messages into named groups with a potential summary for each group.

    When the discussion is flattened, only the last `keep_last_n_obs` groups are kept in the final list,
    the other groups are replaced by their summaries if they have one.
    """

    def __init__(self, keep_last_n_obs=None):
        self.groups: list[MsgGroup] = []
        self.keep_last_n_obs: int | None = keep_last_n_obs

    def append(self, message: MessageBuilder):
        """Append a message to the last group."""
        self.groups[-1].messages.append(message)

    def new_group(self, name: str = None):
        """Start a new group of messages."""
        if name is None:
            name = f"group_{len(self.groups)}"
        self.groups.append(MsgGroup(name))

    def flatten(self) -> list[MessageBuilder]:
        """Flatten the groups into a single list of messages."""

        keep_last_n_obs = self.keep_last_n_obs or len(self.groups)
        messages = []
        for i, group in enumerate(self.groups):
            is_tail = i >= len(self.groups) - keep_last_n_obs

            if not is_tail and group.summary is not None:
                messages.append(group.summary)
            else:
                messages.extend(group.messages)
            # Mark all summarized messages for caching
            if i == len(self.groups) - keep_last_n_obs:
                for msg in messages:  # unset previous cache breakpoints
                    msg._cache_breakpoint = False
                # set new cache breakpoint
                messages[i].mark_all_previous_msg_for_caching()
        return messages

    def set_last_summary(self, summary: MessageBuilder):
        # append None to summaries until we reach the current group index
        self.groups[-1].summary = summary

    def get_last_summary(self) -> MessageBuilder | None:
        """Get the last summary message."""
        if len(self.groups) == 0:
            return None
        return self.groups[-1].summary

    def is_goal_set(self) -> bool:
        """Check if the goal is set in the first group."""
        return len(self.groups) > 0

SYS_MSG = """You are a web agent. Based on the observation, you will decide which action to take to accomplish your goal.
You strive for excellence and need to be as meticulous as possible. Make sure to explore when not sure.
"""

@dataclass
class Goal(Block):
    """Block to add the goal to the messages."""

    goal_as_system_msg: bool = True

    def apply(
        self, llm, discussion: StructuredDiscussion, obs: dict, sys_msg: str = SYS_MSG
    ) -> dict:
        system_message = llm.msg.system().add_text(sys_msg)
        discussion.append(system_message)

        if self.goal_as_system_msg:
            goal_message = llm.msg.system()
        else:
            goal_message = llm.msg.user()

        goal_message.add_text("# Goal:\n")
        for content in obs["goal_object"]:
            if content["type"] == "text":
                goal_message.add_text(content["text"])
            elif content["type"] == "image_url":
                goal_message.add_image(content["image_url"])
        discussion.append(goal_message)

AXTREE_NOTE = """
AXTree extracts most of the interactive elements of the DOM in a tree structure. It may also contain information that is not visible in the screenshot.
A line starting with [bid] is a node in the AXTree. It is a unique alpha-numeric identifier to be used when calling tools, e.g, click(bid="a253"). Make sure to include letters and numbers in the bid.
"""

@dataclass
class Obs(Block):
    """Block to add the observation to the messages."""

    use_last_error: bool = True
    use_screenshot: bool = True
    use_axtree: bool = False
    use_dom: bool = False
    use_som: bool = False
    use_tabs: bool = False
    # add_mouse_pointer: bool = False
    use_zoomed_webpage: bool = False
    skip_preprocessing: bool = False

    def apply(
        self, llm, discussion: StructuredDiscussion, obs: dict, last_llm_output: LLMOutput
    ) -> dict:
        obs_msg = llm.msg.user()
        tool_calls = last_llm_output.tool_calls
        if self.use_last_error:
            if obs["last_action_error"] != "":
                obs_msg.add_text(f"Last action error:\n{obs['last_action_error']}")

        if self.use_screenshot:
            if self.use_som:
                screenshot = obs["screenshot_som"]
            else:
                screenshot = obs["screenshot"]

            # if self.add_mouse_pointer:
            #     screenshot = np.array(
            #         agent_utils.add_mouse_pointer_from_action(
            #             Image.fromarray(obs["screenshot"]), obs["last_action"]
            #         )
            #     )

            obs_msg.add_image(image_to_png_base64_url(screenshot))
        if self.use_axtree:
            obs_msg.add_text(f"AXTree:\n{AXTREE_NOTE}\n{obs['axtree_txt']}")
        if self.use_dom:
            obs_msg.add_text(f"DOM:\n{obs['pruned_html']}")
        if self.use_tabs:
            obs_msg.add_text(_format_tabs(obs))

        discussion.append(obs_msg)

        if tool_calls:
            for call in tool_calls:
                call.response_text("See Observation")
            tool_response = llm.msg.add_responded_tool_calls(tool_calls)
            discussion.append(tool_response)

        return obs_msg

def _format_tabs(obs):
    """Format the open tabs in a llm-readable way."""
    prompt_pieces = ["Currently open tabs:"]
    for page_index, (page_url, page_title) in enumerate(
        zip(obs["open_pages_urls"], obs["open_pages_titles"])
    ):
        active_or_not = " (active tab)" if page_index == obs["active_page_index"] else ""
        prompt_piece = f"""\
Tab {page_index}{active_or_not}:
    Title: {page_title}
    URL: {page_url}
"""
        prompt_pieces.append(prompt_piece)
    return "\n".join(prompt_pieces)

@dataclass
class GeneralHints(Block):
    use_hints: bool = True

    def apply(self, llm, discussion: StructuredDiscussion) -> dict:
        if not self.use_hints:
            return

        hints = []

        hints.append(
            """Use ControlOrMeta instead of Control and Meta for keyboard shortcuts, to be cross-platform compatible. E.g. use ControlOrMeta for mutliple selection in lists.\n"""
        )

        discussion.append(llm.msg.user().add_text("\n".join(hints)))

@dataclass
class Summarizer(Block):
    """Block to summarize the last action and the current state of the environment."""

    do_summary: bool = False
    high_details: bool = True

    def apply(self, llm, discussion: StructuredDiscussion) -> dict:
        if not self.do_summary:
            return

        msg = llm.msg.user().add_text("""Summarize\n""")

        discussion.append(msg)

        summary_response = llm(APIPayload(messages=discussion.flatten()))

        summary_msg = llm.msg.assistant().add_text(summary_response.think)
        discussion.append(summary_msg)
        discussion.set_last_summary(summary_msg)
        return summary_msg

    def apply_init(self, llm, discussion: StructuredDiscussion) -> dict:
        """Initialize the summarizer block."""
        if not self.do_summary:
            return

        system_msg = llm.msg.system()
        if self.high_details:
            # Add a system message to the LLM to indicate that it should summarize
            system_msg.add_text(
                """# Summarizer instructions:\nWhen asked to summarize, do the following:
1) Summarize the effect of the last action, with attention to details.
2) Give a semantic description of the current state of the environment, with attention to details. If there was a repeating mistake, mention the cause of it.
3) Reason about the overall task at a high level.
4) What hint can be relevant for the next action? Only chose from the hints provided in the task description. Or select none.
5) Reason about the next action to take, based on the current state and the goal.
"""
            )
        else:
            system_msg.add_text(
                """When asked to summarize, give a semantic description of the current state of the environment."""
            )
        discussion.append(system_msg)

@dataclass
class TaskHint(Block):
    use_task_hint: bool = True
    hint_db_rel_path: str = "hint_db.csv"
    hint_retrieval_mode: Literal["direct", "llm", "emb"] = "direct"
    top_n: int = 4  # Number of top hints to return when using embedding retrieval
    embedder_model: str = "Qwen/Qwen3-Embedding-0.6B"  # Model for embedding hints
    embedder_server: str = "http://localhost:5000"
    llm_prompt: str = """We're choosing hints to help solve the following task:\n{goal}.\n
You need to choose the most relevant hints topic from the following list:\n\nHint topics:\n{topics}\n
Choose hint topic for the task and return only its number, e.g. 1. If you don't know the answer, return -1."""

    def _init(self):
        """Initialize the block."""
        if Path(self.hint_db_rel_path).is_absolute():
            hint_db_path = Path(self.hint_db_rel_path)
        else:
            hint_db_path = Path(__file__).parent / self.hint_db_rel_path
        self.hint_db = pd.read_csv(hint_db_path, header=0, index_col=None, dtype=str)
        if self.hint_retrieval_mode == "emb":
            self.encode_hints()

    def oai_embed(self, text: str):
        response = self._oai_emb.create(input=text, model="text-embedding-3-small")
        return response.data[0].embedding

    def encode_hints(self):
        self.uniq_hints = self.hint_db.drop_duplicates(subset=["hint"], keep="first")
        logger.info(
            f"Encoding {len(self.uniq_hints)} unique hints with semantic keys using {self.embedder_model} model."
        )
        hints = self.uniq_hints["hint"].tolist()
        semantic_keys = self.uniq_hints["semantic_keys"].tolist()
        lines = [f"{k}: {h}" for h, k in zip(hints, semantic_keys)]
        emb_path = f"{self.hint_db_rel_path}.embs.npy"
        assert os.path.exists(emb_path), f"Embedding file not found: {emb_path}"
        logger.info(f"Loading hint embeddings from: {emb_path}")
        emb_dict = np.load(emb_path, allow_pickle=True).item()
        self.hint_embeddings = np.array([emb_dict[k] for k in lines])
        logger.info(f"Loaded hint embeddings shape: {self.hint_embeddings.shape}")

    def apply(self, llm, discussion: StructuredDiscussion, task_name: str) -> dict:
        if not self.use_task_hint:
            return {}

        goal = "\n".join([c.get("text", "") for c in discussion.groups[0].messages[1].content])
        task_hints = self.choose_hints(llm, task_name, goal)

        hints = []
        for hint in task_hints:
            hint = hint.strip()
            if hint:
                hints.append(f"- {hint}")

        if len(hints) > 0:
            hints_str = (
                "# Hints:\nHere are some hints for the task you are working on:\n"
                + "\n".join(hints)
            )
            msg = llm.msg.user().add_text(hints_str)

            discussion.append(msg)

    def choose_hints(self, llm, task_name: str, goal: str) -> list[str]:
        """Choose hints based on the task name."""
        if self.hint_retrieval_mode == "llm":
            return self.choose_hints_llm(llm, goal)
        elif self.hint_retrieval_mode == "direct":
            return self.choose_hints_direct(task_name)
        elif self.hint_retrieval_mode == "emb":
            return self.choose_hints_emb(goal)
        else:
            raise ValueError(f"Unknown hint retrieval mode: {self.hint_retrieval_mode}")

    def choose_hints_llm(self, llm, goal: str) -> list[str]:
        """Choose hints using LLM to filter the hints."""
        topic_to_hints = defaultdict(list)
        for i, row in self.hint_db.iterrows():
            topic_to_hints[row["semantic_keys"]].append(i)
        hint_topics = list(topic_to_hints.keys())
        topics = "\n".join([f"{i}. {h}" for i, h in enumerate(hint_topics)])
        prompt = self.llm_prompt.format(goal=goal, topics=topics)
        response = llm(APIPayload(messages=[llm.msg.user().add_text(prompt)]))
        try:
            hint_topic_idx = json.loads(response.think)
            if hint_topic_idx < 0 or hint_topic_idx >= len(hint_topics):
                logger.error(f"Wrong LLM hint id response: {response.think}, no hints")
                return []
            hint_topic = hint_topics[hint_topic_idx]
            hint_indices = topic_to_hints[hint_topic]
            df = self.hint_db.iloc[hint_indices].copy()
            df = df.drop_duplicates(subset=["hint"], keep="first")  # leave only unique hints
            hints = df["hint"].tolist()
            logger.debug(f"LLM hint topic {hint_topic_idx}, chosen hints: {df['hint'].tolist()}")
        except json.JSONDecodeError:
            logger.error(f"Failed to parse LLM hint id response: {response.think}, no hints")
            hints = []
        return hints

    def choose_hints_emb(self, goal: str) -> list[str]:
        """Choose hints using embeddings to filter the hints."""
        goal_embeddings = self._encode([goal], prompt="task description")
        similarities = self._similarity(goal_embeddings.tolist(), self.hint_embeddings.tolist())
        top_indices = similarities.argsort()[0][-self.top_n :].tolist()
        logger.info(f"Top hint indices based on embedding similarity: {top_indices}")
        hints = self.uniq_hints.iloc[top_indices]
        logger.info(f"Embedding-based hints chosen: {hints}")
        return hints["hint"].tolist()

    def _encode(self, texts: list[str], prompt: str = "", timeout: int = 10, max_retries: int = 5):
        """Call the encode API endpoint with timeout and retries"""
        for attempt in range(max_retries):
            try:
                response = requests.post(
                    f"{self.embedder_server}/encode",
                    json={"texts": texts, "prompt": prompt},
                    timeout=timeout,
                )
                embs = response.json()["embeddings"]
                return np.asarray(embs)
            except (requests.exceptions.RequestException, requests.exceptions.Timeout) as e:
                if attempt == max_retries - 1:
                    raise e
                time.sleep(random.uniform(1, timeout))
                continue

    def _similarity(
        self, texts1: list[str], texts2: list[str], timeout: int = 2, max_retries: int = 5
    ):
        """Call the similarity API endpoint with timeout and retries"""
        for attempt in range(max_retries):
            try:
                response = requests.post(
                    f"{self.embedder_server}/similarity",
                    json={"texts1": texts1, "texts2": texts2},
                    timeout=timeout,
                )
                similarities = response.json()["similarities"]
                return np.asarray(similarities)
            except (requests.exceptions.RequestException, requests.exceptions.Timeout) as e:
                if attempt == max_retries - 1:
                    raise e
                time.sleep(random.uniform(1, timeout))
                continue

    def choose_hints_direct(self, task_name: str) -> list[str]:
        hints = self.hint_db[
            self.hint_db["task_name"].apply(lambda x: fnmatch.fnmatch(x, task_name))
        ]
        return hints["hint"].tolist()

@dataclass
class PromptConfig:
    tag_screenshot: bool = True  # Whether to tag the screenshot with the last action.
    goal: Goal = None
    obs: Obs = None
    summarizer: Summarizer = None
    general_hints: GeneralHints = None
    task_hint: TaskHint = None
    keep_last_n_obs: int = 1
    multiaction: bool = False
    action_subsets: tuple[str] = None

@dataclass
class ToolUseAgentArgs(AgentArgs):
    model_args: BaseModelArgs = None
    config: PromptConfig = None
    use_raw_page_output: bool = False  # This attribute is used in loop.py to setup the env.
    action_set: bgym.AbstractActionSet | None = None

    def __post_init__(self):
        try:
            self.agent_name = f"ToolUse-{self.model_args.model_name}".replace("/", "_")
        except AttributeError:
            pass

    def make_agent(self) -> bgym.Agent:
        if self.config is None:
            self.config = DEFAULT_PROMPT_CONFIG
        return ToolUseAgent(
            model_args=self.model_args,  # type: ignore
            config=self.config,
            action_set=self.action_set,
        )

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

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

    def set_benchmark(self, benchmark: AgentLabBenchmark | BgymBenchmark, demo_mode: bool):
        """Set benchmark specific flags."""
        benchmark_name = benchmark.name
        if benchmark_name == "osworld":
            self.config.obs.skip_preprocessing = True

class ToolUseAgent(bgym.Agent):
    def __init__(
        self,
        model_args: OpenAIResponseModelArgs,
        config: PromptConfig = None,
        action_set: bgym.AbstractActionSet | None = None,
    ):
        self.model_args = model_args
        self.config = config
        self.action_set: bgym.AbstractActionSet = action_set or bgym.HighLevelActionSet(
            self.config.action_subsets,
            multiaction=self.config.multiaction,  # type: ignore
        )
        self.tools = self.action_set.to_tool_description(api=model_args.api)

        self.call_ids = []

        self.llm = model_args.make_model()
        self.msg_builder = model_args.get_message_builder()
        self.llm.msg = self.msg_builder

        self.task_hint = self.config.task_hint.make()
        self.obs_block = self.config.obs.make()

        self.discussion = StructuredDiscussion(self.config.keep_last_n_obs)
        self.last_response: LLMOutput = LLMOutput()
        self._responses: list[LLMOutput] = []

    def o
# ... truncated ...


Block

from agentlab.benchmarks.abstract_env import AbstractBenchmark as AgentLabBenchmark
from agentlab.benchmarks.osworld import OSWorldActionSet
from agentlab.llm.base_api import BaseModelArgs
from agentlab.llm.llm_utils import image_to_png_base64_url
from agentlab.llm.response_api import (
    APIPayload,
    ClaudeResponseModelArgs,
    LLMOutput,
    MessageBuilder,
    OpenAIChatModelArgs,
    OpenAIResponseModelArgs,
    OpenRouterModelArgs,
    ToolCalls,
)
from agentlab.llm.tracking import cost_tracker_decorator

logger = logging.getLogger(__name__)

@dataclass
class Block(ABC):
    def _init(self):
        """Initialize the block."""
        pass

    def make(self) -> "Block":
        """Returns a copy so the init can start adding some stuff to `self` without changing the
        original datatclass that should only contain a config.
        The aim is avoid having 2 calss definition for each block, e.g. Block and BlockArgs.

        Returns:
            Block: A copy of the current block instance with initialization applied.
        """
        block = self.__class__(**asdict(self))
        block._init()
        return block

    @abstractmethod
    def apply(self, llm, messages: list[MessageBuilder], **kwargs):
        pass

@dataclass
class MsgGroup:
    name: str = None
    messages: list[MessageBuilder] = field(default_factory=list)
    summary: MessageBuilder = None

class StructuredDiscussion:
    """
    A structured discussion that groups messages into named groups with a potential summary for each group.

    When the discussion is flattened, only the last `keep_last_n_obs` groups are kept in the final list,
    the other groups are replaced by their summaries if they have one.
    """

    def __init__(self, keep_last_n_obs=None):
        self.groups: list[MsgGroup] = []
        self.keep_last_n_obs: int | None = keep_last_n_obs


MsgGroup

        pass

    def make(self) -> "Block":
        """Returns a copy so the init can start adding some stuff to `self` without changing the
        original datatclass that should only contain a config.
        The aim is avoid having 2 calss definition for each block, e.g. Block and BlockArgs.

        Returns:
            Block: A copy of the current block instance with initialization applied.
        """
        block = self.__class__(**asdict(self))
        block._init()
        return block

    @abstractmethod
    def apply(self, llm, messages: list[MessageBuilder], **kwargs):
        pass

@dataclass
class MsgGroup:
    name: str = None
    messages: list[MessageBuilder] = field(default_factory=list)
    summary: MessageBuilder = None

class StructuredDiscussion:
    """
    A structured discussion that groups messages into named groups with a potential summary for each group.

    When the discussion is flattened, only the last `keep_last_n_obs` groups are kept in the final list,
    the other groups are replaced by their summaries if they have one.
    """

    def __init__(self, keep_last_n_obs=None):
        self.groups: list[MsgGroup] = []
        self.keep_last_n_obs: int | None = keep_last_n_obs

    def append(self, message: MessageBuilder):
        """Append a message to the last group."""
        self.groups[-1].messages.append(message)

    def new_group(self, name: str = None):
        """Start a new group of messages."""


StructuredDiscussion

        Returns:
            Block: A copy of the current block instance with initialization applied.
        """
        block = self.__class__(**asdict(self))
        block._init()
        return block

    @abstractmethod
    def apply(self, llm, messages: list[MessageBuilder], **kwargs):
        pass

@dataclass
class MsgGroup:
    name: str = None
    messages: list[MessageBuilder] = field(default_factory=list)
    summary: MessageBuilder = None

class StructuredDiscussion:
    """
    A structured discussion that groups messages into named groups with a potential summary for each group.

    When the discussion is flattened, only the last `keep_last_n_obs` groups are kept in the final list,
    the other groups are replaced by their summaries if they have one.
    """

    def __init__(self, keep_last_n_obs=None):
        self.groups: list[MsgGroup] = []
        self.keep_last_n_obs: int | None = keep_last_n_obs

    def append(self, message: MessageBuilder):
        """Append a message to the last group."""
        self.groups[-1].messages.append(message)

    def new_group(self, name: str = None):
        """Start a new group of messages."""
        if name is None:
            name = f"group_{len(self.groups)}"
        self.groups.append(MsgGroup(name))

    def flatten(self) -> list[MessageBuilder]:
        """Flatten the groups into a single list of messages."""

        keep_last_n_obs = self.keep_last_n_obs or len(self.groups)
        messages = []
        for i, group in enumerate(self.groups):
            is_tail = i >= len(self.groups) - keep_last_n_obs

            if not is_tail and group.summary is not None:
                messages.append(group.summary)
            else:
                messages.extend(group.messages)
            # Mark all summarized messages for caching
            if i == len(self.groups) - keep_last_n_obs:
                for msg in messages:  # unset previous cache breakpoints
                    msg._cache_breakpoint = False
                # set new cache breakpoint
                messages[i].mark_all_previous_msg_for_caching()
        return messages

    def set_last_summary(self, summary: MessageBuilder):
        # append None to summaries until we reach the current group index
        self.groups[-1].summary = summary

    def get_last_summary(self) -> MessageBuilder | None:
        """Get the last summary message."""
        if len(self.groups) == 0:
            return None
        return self.groups[-1].summary

    def is_goal_set(self) -> bool:
        """Check if the goal is set in the first group."""
        return len(self.groups) > 0

SYS_MSG = """You are a web agent. Based on the observation, you will decide which action to take to accomplish your goal.
You strive for excellence and need to be as meticulous as possible. Make sure to explore when not sure.
"""

@dataclass
class Goal(Block):
    """Block to add the goal to the messages."""

    goal_as_system_msg: bool = True

    def apply(
        self, llm, discussion: StructuredDiscussion, obs: dict, sys_msg: str = SYS_MSG
    ) -> dict:
        system_message = llm.msg.system().add_text(sys_msg)
        discussion.append(system_message)

        if self.goal_as_system_msg:


Goal

        # append None to summaries until we reach the current group index
        self.groups[-1].summary = summary

    def get_last_summary(self) -> MessageBuilder | None:
        """Get the last summary message."""
        if len(self.groups) == 0:
            return None
        return self.groups[-1].summary

    def is_goal_set(self) -> bool:
        """Check if the goal is set in the first group."""
        return len(self.groups) > 0

SYS_MSG = """You are a web agent. Based on the observation, you will decide which action to take to accomplish your goal.
You strive for excellence and need to be as meticulous as possible. Make sure to explore when not sure.
"""

@dataclass
class Goal(Block):
    """Block to add the goal to the messages."""

    goal_as_system_msg: bool = True

    def apply(
        self, llm, discussion: StructuredDiscussion, obs: dict, sys_msg: str = SYS_MSG
    ) -> dict:
        system_message = llm.msg.system().add_text(sys_msg)
        discussion.append(system_message)

        if self.goal_as_system_msg:
            goal_message = llm.msg.system()
        else:
            goal_message = llm.msg.user()

        goal_message.add_text("# Goal:\n")
        for content in obs["goal_object"]:
            if content["type"] == "text":
                goal_message.add_text(content["text"])
            elif content["type"] == "image_url":
                goal_message.add_image(content["image_url"])
        discussion.append(goal_message)

AXTREE_NOTE = """
AXTree extracts most of the interactive elements of the DOM in a tree structure. It may also contain information that is not visible in the screenshot.
A line starting with [bid] is a node in the AXTree. It is a unique alpha-numeric identifier to be used when calling tools, e.g, click(bid="a253"). Make sure to include letters and numbers in the bid.
"""

@dataclass
class Obs(Block):
    """Block to add the observation to the messages."""

    use_last_error: bool = True
    use_screenshot: bool = True
    use_axtree: bool = False
    use_dom: bool = False
    use_som: bool = False
    use_tabs: bool = False
    # add_mouse_pointer: bool = False
    use_zoomed_webpage: bool = False


Obs

            goal_message = llm.msg.system()
        else:
            goal_message = llm.msg.user()

        goal_message.add_text("# Goal:\n")
        for content in obs["goal_object"]:
            if content["type"] == "text":
                goal_message.add_text(content["text"])
            elif content["type"] == "image_url":
                goal_message.add_image(content["image_url"])
        discussion.append(goal_message)

AXTREE_NOTE = """
AXTree extracts most of the interactive elements of the DOM in a tree structure. It may also contain information that is not visible in the screenshot.
A line starting with [bid] is a node in the AXTree. It is a unique alpha-numeric identifier to be used when calling tools, e.g, click(bid="a253"). Make sure to include letters and numbers in the bid.
"""

@dataclass
class Obs(Block):
    """Block to add the observation to the messages."""

    use_last_error: bool = True
    use_screenshot: bool = True
    use_axtree: bool = False
    use_dom: bool = False
    use_som: bool = False
    use_tabs: bool = False
    # add_mouse_pointer: bool = False
    use_zoomed_webpage: bool = False
    skip_preprocessing: bool = False

    def apply(
        self, llm, discussion: StructuredDiscussion, obs: dict, last_llm_output: LLMOutput
    ) -> dict:
        obs_msg = llm.msg.user()
        tool_calls = last_llm_output.tool_calls
        if self.use_last_error:
            if obs["last_action_error"] != "":
                obs_msg.add_text(f"Last action error:\n{obs['last_action_error']}")

        if self.use_screenshot:
            if self.use_som:
                screenshot = obs["screenshot_som"]
            else:
                screenshot = obs["screenshot"]

            # if self.add_mouse_pointer:
            #     screenshot = np.array(
            #         agent_utils.add_mouse_pointer_from_action(
            #             Image.fromarray(obs["screenshot"]), obs["last_action"]
            #         )
            #     )

            obs_msg.add_image(image_to_png_base64_url(screenshot))
        if self.use_axtree:
            obs_msg.add_text(f"AXTree:\n{AXTREE_NOTE}\n{obs['axtree_txt']}")
        if self.use_dom:
            obs_msg.add_text(f"DOM:\n{obs['pruned_html']}")
        if self.use_tabs:
            obs_msg.add_text(_format_tabs(obs))

        discussion.append(obs_msg)

        if tool_calls:
            for call in tool_calls:
                call.response_text("See Observation")
            tool_response = llm.msg.add_responded_tool_calls(tool_calls)
            discussion.append(tool_response)

        return obs_msg

def _format_tabs(obs):
    """Format the open tabs in a llm-readable way."""
    prompt_pieces = ["Currently open tabs:"]
    for page_index, (page_url, page_title) in enumerate(
        zip(obs["open_pages_urls"], obs["open_pages_titles"])
    ):
        active_or_not = " (active tab)" if page_index == obs["active_page_index"] else ""
        prompt_piece = f"""\
Tab {page_index}{active_or_not}:
    Title: {page_title}
    URL: {page_url}
"""
        prompt_pieces.append(prompt_piece)
    return "\n".join(prompt_pieces)

@dataclass
class GeneralHints(Block):


            obs_msg.add_image(image_to_png_base64_url(screenshot))
        if self.use_axtree:
            obs_msg.add_text(f"AXTree:\n{AXTREE_NOTE}\n{obs['axtree_txt']}")
        if self.use_dom:
            obs_msg.add_text(f"DOM:\n{obs['pruned_html']}")
        if self.use_tabs:
            obs_msg.add_text(_format_tabs(obs))

        discussion.append(obs_msg)

        if tool_calls:
            for call in tool_calls:
                call.response_text("See Observation")
            tool_response = llm.msg.add_responded_tool_calls(tool_calls)
            discussion.append(tool_response)

        return obs_msg

def _format_tabs(obs):
    """Format the open tabs in a llm-readable way."""
    prompt_pieces = ["Currently open tabs:"]
    for page_index, (page_url, page_title) in enumerate(
        zip(obs["open_pages_urls"], obs["open_pages_titles"])
    ):
        active_or_not = " (active tab)" if page_index == obs["active_page_index"] else ""
        prompt_piece = f"""\
Tab {page_index}{active_or_not}:
    Title: {page_title}
    URL: {page_url}
"""
        prompt_pieces.append(prompt_piece)
    return "\n".join(prompt_pieces)

@dataclass
class GeneralHints(Block):
    use_hints: bool = True

    def apply(self, llm, discussion: StructuredDiscussion) -> dict:
        if not self.use_hints:
            return

        hints = []

        hints.append(
            """Use ControlOrMeta instead of Control and Meta for keyboard shortcuts, to be cross-platform compatible. E.g. use ControlOrMeta for mutliple selection in lists.\n"""
        )

        discussion.append(llm.msg.user().add_text("\n".join(hints)))

@dataclass


GeneralHints

        return obs_msg

def _format_tabs(obs):
    """Format the open tabs in a llm-readable way."""
    prompt_pieces = ["Currently open tabs:"]
    for page_index, (page_url, page_title) in enumerate(
        zip(obs["open_pages_urls"], obs["open_pages_titles"])
    ):
        active_or_not = " (active tab)" if page_index == obs["active_page_index"] else ""
        prompt_piece = f"""\
Tab {page_index}{active_or_not}:
    Title: {page_title}
    URL: {page_url}
"""
        prompt_pieces.append(prompt_piece)
    return "\n".join(prompt_pieces)

@dataclass
class GeneralHints(Block):
    use_hints: bool = True

    def apply(self, llm, discussion: StructuredDiscussion) -> dict:
        if not self.use_hints:
            return

        hints = []

        hints.append(
            """Use ControlOrMeta instead of Control and Meta for keyboard shortcuts, to be cross-platform compatible. E.g. use ControlOrMeta for mutliple selection in lists.\n"""
        )

        discussion.append(llm.msg.user().add_text("\n".join(hints)))

@dataclass
class Summarizer(Block):
    """Block to summarize the last action and the current state of the environment."""

    do_summary: bool = False
    high_details: bool = True

    def apply(self, llm, discussion: StructuredDiscussion) -> dict:
        if not self.do_summary:
            return

        msg = llm.msg.user().add_text("""Summarize\n""")

        discussion.append(msg)

        summary_response = llm(APIPayload(messages=discussion.flatten()))

        summary_msg = llm.msg.assistant().add_text(summary_response.think)


Summarizer

@dataclass
class GeneralHints(Block):
    use_hints: bool = True

    def apply(self, llm, discussion: StructuredDiscussion) -> dict:
        if not self.use_hints:
            return

        hints = []

        hints.append(
            """Use ControlOrMeta instead of Control and Meta for keyboard shortcuts, to be cross-platform compatible. E.g. use ControlOrMeta for mutliple selection in lists.\n"""
        )

        discussion.append(llm.msg.user().add_text("\n".join(hints)))

@dataclass
class Summarizer(Block):
    """Block to summarize the last action and the current state of the environment."""

    do_summary: bool = False
    high_details: bool = True

    def apply(self, llm, discussion: StructuredDiscussion) -> dict:
        if not self.do_summary:
            return

        msg = llm.msg.user().add_text("""Summarize\n""")

        discussion.append(msg)

        summary_response = llm(APIPayload(messages=discussion.flatten()))

        summary_msg = llm.msg.assistant().add_text(summary_response.think)
        discussion.append(summary_msg)
        discussion.set_last_summary(summary_msg)
        return summary_msg

    def apply_init(self, llm, discussion: StructuredDiscussion) -> dict:
        """Initialize the summarizer block."""
        if not self.do_summary:
            return

        system_msg = llm.msg.system()
        if self.high_details:
            # Add a system message to the LLM to indicate that it should summarize
            system_msg.add_text(
                """# Summarizer instructions:\nWhen asked to summarize, do the following:
1) Summarize the effect of the last action, with attention to details.
2) Give a semantic description of the current state of the environment, with attention to details. If there was a repeating mistake, mention the cause of it.
3) Reason about the overall task at a high level.
4) What hint can be relevant for the next action? Only chose from the hints provided in the task description. Or select none.
5) Reason about the next action to take, based on the current state and the goal.
"""
            )
        else:
            system_msg.add_text(
                """When asked to summarize, give a semantic description of the current state of the environment."""
            )
        discussion.append(system_msg)

@dataclass
class TaskHint(Block):
    use_task_hint: bool = True
    hint_db_rel_path: str = "hint_db.csv"
    hint_retrieval_mode: Literal["direct", "llm", "emb"] = "direct"
    top_n: int = 4  # Number of top hints to return when using embedding retrieval
    embedder_model: str = "Qwen/Qwen3-Embedding-0.6B"  # Model for embedding hints
    embedder_server: str = "http://localhost:5000"
    llm_prompt: str = """We're choosing hints to help solve the following task:\n{goal}.\n
You need to choose the most relevant hints topic from the following list:\n\nHint topics:\n{topics}\n
Choose hint topic for the task and return only its number, e.g. 1. If you don't know the answer, return -1."""

    def _init(self):
        """Initialize the block."""
        if Path(self.hint_db_rel_path).is_absolute():
            hint_db_path = Path(self.hint_db_rel_path)
        else:
            hint_db_path = Path(__file__).parent / self.hint_db_rel_path


TaskHint

        system_msg = llm.msg.system()
        if self.high_details:
            # Add a system message to the LLM to indicate that it should summarize
            system_msg.add_text(
                """# Summarizer instructions:\nWhen asked to summarize, do the following:
1) Summarize the effect of the last action, with attention to details.
2) Give a semantic description of the current state of the environment, with attention to details. If there was a repeating mistake, mention the cause of it.
3) Reason about the overall task at a high level.
4) What hint can be relevant for the next action? Only chose from the hints provided in the task description. Or select none.
5) Reason about the next action to take, based on the current state and the goal.
"""
            )
        else:
            system_msg.add_text(
                """When asked to summarize, give a semantic description of the current state of the environment."""
            )
        discussion.append(system_msg)

@dataclass
class TaskHint(Block):
    use_task_hint: bool = True
    hint_db_rel_path: str = "hint_db.csv"
    hint_retrieval_mode: Literal["direct", "llm", "emb"] = "direct"
    top_n: int = 4  # Number of top hints to return when using embedding retrieval
    embedder_model: str = "Qwen/Qwen3-Embedding-0.6B"  # Model for embedding hints
    embedder_server: str = "http://localhost:5000"
    llm_prompt: str = """We're choosing hints to help solve the following task:\n{goal}.\n
You need to choose the most relevant hints topic from the following list:\n\nHint topics:\n{topics}\n
Choose hint topic for the task and return only its number, e.g. 1. If you don't know the answer, return -1."""

    def _init(self):
        """Initialize the block."""
        if Path(self.hint_db_rel_path).is_absolute():
            hint_db_path = Path(self.hint_db_rel_path)
        else:
            hint_db_path = Path(__file__).parent / self.hint_db_rel_path
        self.hint_db = pd.read_csv(hint_db_path, header=0, index_col=None, dtype=str)
        if self.hint_retrieval_mode == "emb":
            self.encode_hints()

    def oai_embed(self, text: str):
        response = self._oai_emb.create(input=text, model="text-embedding-3-small")
        return response.data[0].embedding

    def encode_hints(self):
        self.uniq_hints = self.hint_db.drop_duplicates(subset=["hint"], keep="first")
        logger.info(
            f"Encoding {len(self.uniq_hints)} unique hints with semantic keys using {self.embedder_model} model."
        )
        hints = self.uniq_hints["hint"].tolist()
        semantic_keys = self.uniq_hints["semantic_keys"].tolist()
        lines = [f"{k}: {h}" for h, k in zip(hints, semantic_keys)]
        emb_path = f"{self.hint_db_rel_path}.embs.npy"
        assert os.path.exists(emb_path), f"Embedding file not found: {emb_path}"
        logger.info(f"Loading hint embeddings from: {emb_path}")
        emb_dict = np.load(emb_path, allow_pickle=True).item()
        self.hint_embeddings = np.array([emb_dict[k] for k in lines])
        logger.info(f"Loaded hint embeddings shape: {self.hint_embeddings.shape}")

    def apply(self, llm, discussion: StructuredDiscussion, task_name: str) -> dict:
        if not self.use_task_hint:
            return {}

        goal = "\n".join([c.get("text", "") for c in discussion.groups[0].messages[1].content])
        task_hints = self.choose_hints(llm, task_name, goal)

        hints = []
        for hint in task_hints:
            hint = hint.strip()
            if hint:
                hints.append(f"- {hint}")

        if len(hints) > 0:
            hints_str = (
                "# Hints:\nHere are some hints for the task you are working on:\n"
                + "\n".join(hints)
            )
            msg = llm.msg.user().add_text(hints_str)

            discussion.append(msg)

    def choose_hints(self, llm, task_name: str, goal: str) -> list[str]:
        """Choose hints based on the task name."""
        if self.hint_retrieval_mode == "llm":
            return self.choose_hints_llm(llm, goal)
        elif self.hint_retrieval_mode == "direct":
            return self.choose_hints_direct(task_name)
        elif self.hint_retrieval_mode == "emb":
            return self.choose_hints_emb(goal)
        else:
            raise ValueError(f"Unknown hint retrieval mode: {self.hint_retrieval_mode}")

    def choose_hints_llm(self, llm, goal: str) -> list[str]:
        """Choose hints using LLM to filter the hints."""
        topic_to_hints = defaultdict(list)
        for i, row in self.hint_db.iterrows():
            topic_to_hints[row["semantic_keys"]].append(i)
        hint_topics = list(topic_to_hints.keys())
        topics = "\n".join([f"{i}. {h}" for i, h in enumerate(hint_topics)])
        prompt = self.llm_prompt.format(goal=goal, topics=topics)
        response = llm(APIPayload(messages=[llm.msg.user().add_text(prompt)]))
        try:
            hint_topic_idx = json.loads(response.think)
            if hint_topic_idx < 0 or hint_topic_idx >= len(hint_topics):
                logger.error(f"Wrong LLM hint id response: {response.think}, no hints")
                return []
            hint_topic = hint_topics[hint_topic_idx]
            hint_indices = topic_to_hints[hint_topic]
            df = self.hint_db.iloc[hint_indices].copy()
            df = df.drop_duplicates(subset=["hint"], keep="first")  # leave only unique hints
            hints = df["hint"].tolist()
            logger.debug(f"LLM hint topic {hint_topic_idx}, chosen hints: {df['hint'].tolist()}")
        except json.JSONDecodeError:
            logger.error(f"Failed to parse LLM hint id response: {response.think}, no hints")
            hints = []
        return hints

    def choose_hints_emb(self, goal: str) -> list[str]:
        """Choose hints using embeddings to filter the hints."""
        goal_embeddings = self._encode([goal], prompt="task description")
        similarities = self._similarity(goal_embeddings.tolist(), self.hint_embeddings.tolist())
        top_indices = similarities.argsort()[0][-self.top_n :].tolist()
        logger.info(f"Top hint indices based on embedding similarity: {top_indices}")
        hints = self.uniq_hints.iloc[top_indices]
        logger.info(f"Embedding-based hints chosen: {hints}")
        return hints["hint"].tolist()

    def _encode(self, texts: list[str], prompt: str = "", timeout: int = 10, max_retries: int = 5):
        """Call the encode API endpoint with timeout and retries"""
        for attempt in range(max_retries):
            try:
                response = requests.post(
                    f"{self.embedder_server}/encode",
                    json={"texts": texts, "prompt": prompt},
                    timeout=timeout,
                )
                embs = response.json()["embeddings"]
                return np.asarray(embs)
            except (requests.exceptions.RequestException, requests.exceptions.Timeout) as e:
                if attempt == max_retries - 1:
                    raise e
                time.sleep(random.uniform(1, timeout))
                continue

    def _similarity(
        self, texts1: list[str], texts2: list[str], timeout: int = 2, max_retries: int = 5
    ):
        """Call the similarity API endpoint with timeout and retries"""
        for attempt in range(max_retries):
            try:
                response = requests.post(
                    f"{self.embedder_server}/similarity",
                    json={"texts1": texts1, "texts2": texts2},
                    timeout=timeout,
                )
                similarities = response.json()["similarities"]
                return np.asarray(similarities)
            except (requests.exceptions.RequestException, requests.exceptions.Timeout) as e:
                if attempt == max_retries - 1:
                    raise e
                time.sleep(random.uniform(1, timeout))
                continue

    def choose_hints_direct(self, task_name: str) -> list[str]:
        hints = self.hint_db[
            self.hint_db["task_name"].apply(lambda x: fnmatch.fnmatch(x, task_name))
        ]
        return hints["hint"].tolist()

@dataclass
class PromptConfig:
    tag_screenshot: bool = True  # Whether to tag the screenshot with the last action.
    goal: Goal = None
    obs: Obs = None
    summarizer: Summarizer = None
    general_hints: GeneralHints = None
    task_hint: TaskHint = None
    keep_last_n_obs: int = 1
    multiaction: bool = False
    action_subsets: tuple[str] = None

@dataclass
class ToolUseAgentArgs(AgentArgs):
    model_args: BaseModelArgs = None
    config: PromptConfig = None
    use_raw_page_output: bool = False  # This attribute is used in loop.py to setup the env.


PromptConfig

                    f"{self.embedder_server}/similarity",
                    json={"texts1": texts1, "texts2": texts2},
                    timeout=timeout,
                )
                similarities = response.json()["similarities"]
                return np.asarray(similarities)
            except (requests.exceptions.RequestException, requests.exceptions.Timeout) as e:
                if attempt == max_retries - 1:
                    raise e
                time.sleep(random.uniform(1, timeout))
                continue

    def choose_hints_direct(self, task_name: str) -> list[str]:
        hints = self.hint_db[
            self.hint_db["task_name"].apply(lambda x: fnmatch.fnmatch(x, task_name))
        ]
        return hints["hint"].tolist()

@dataclass
class PromptConfig:
    tag_screenshot: bool = True  # Whether to tag the screenshot with the last action.
    goal: Goal = None
    obs: Obs = None
    summarizer: Summarizer = None
    general_hints: GeneralHints = None
    task_hint: TaskHint = None
    keep_last_n_obs: int = 1
    multiaction: bool = False
    action_subsets: tuple[str] = None

@dataclass
class ToolUseAgentArgs(AgentArgs):
    model_args: BaseModelArgs = None
    config: PromptConfig = None
    use_raw_page_output: bool = False  # This attribute is used in loop.py to setup the env.
    action_set: bgym.AbstractActionSet | None = None

    def __post_init__(self):
        try:
            self.agent_name = f"ToolUse-{self.model_args.model_name}".replace("/", "_")
        except AttributeError:
            pass

    def make_agent(self) -> bgym.Agent:
        if self.config is None:
            self.config = DEFAULT_PROMPT_CONFIG
        return ToolUseAgent(
            model_args=self.model_args,  # type: ignore


ToolUseAgentArgs

        hints = self.hint_db[
            self.hint_db["task_name"].apply(lambda x: fnmatch.fnmatch(x, task_name))
        ]
        return hints["hint"].tolist()

@dataclass
class PromptConfig:
    tag_screenshot: bool = True  # Whether to tag the screenshot with the last action.
    goal: Goal = None
    obs: Obs = None
    summarizer: Summarizer = None
    general_hints: GeneralHints = None
    task_hint: TaskHint = None
    keep_last_n_obs: int = 1
    multiaction: bool = False
    action_subsets: tuple[str] = None

@dataclass
class ToolUseAgentArgs(AgentArgs):
    model_args: BaseModelArgs = None
    config: PromptConfig = None
    use_raw_page_output: bool = False  # This attribute is used in loop.py to setup the env.
    action_set: bgym.AbstractActionSet | None = None

    def __post_init__(self):
        try:
            self.agent_name = f"ToolUse-{self.model_args.model_name}".replace("/", "_")
        except AttributeError:
            pass

    def make_agent(self) -> bgym.Agent:
        if self.config is None:
            self.config = DEFAULT_PROMPT_CONFIG
        return ToolUseAgent(
            model_args=self.model_args,  # type: ignore
            config=self.config,
            action_set=self.action_set,
        )

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

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

    def set_benchmark(self, benchmark: AgentLabBenchmark | BgymBenchmark, demo_mode: bool):
        """Set benchmark specific flags."""
        benchmark_name = benchmark.name
        if benchmark_name == "osworld":
            self.config.obs.skip_preprocessing = True

class ToolUseAgent(bgym.Agent):
    def __init__(
        self,
        model_args: OpenAIResponseModelArgs,
        config: PromptConfig = None,
        action_set: bgym.AbstractActionSet | None = None,
    ):
        self.model_args = model_args
        self.config = config
        self.action_set: bgym.AbstractActionSet = action_set or bgym.HighLevelActionSet(
            self.config.action_subsets,
            multiaction=self.config.multiaction,  # type: ignore
        )
        self.tools = self.action_set.to_tool_description(api=model_args.api)

        self.call_ids = []

        self.llm = model_args.make_model()