browser_agent.py

# -*- coding: utf-8 -*-
"""Browser Agent"""
# pylint: disable=W0212
# pylint: disable=too-many-lines
import os
from typing import Type, Optional
import asyncio
from pydantic import BaseModel
from core.agent import CoreMixin
from core.memory import MemoryMixin
from core.evaluation import EvaluationMixin
from core.utils import UtilsMixin
from search.functions import SearchMixin
from agentscope.agent import ReActAgent
from agentscope.formatter import FormatterBase
from agentscope.memory import MemoryBase
from agentscope.message import (
    Msg,
    ToolResultBlock,
)
from agentscope.model import ChatModelBase
from agentscope.tool import Toolkit
from agentscope.token import TokenCounterBase, OpenAITokenCounter

with open(
    "build_in_prompt/browser_agent_sys_prompt.md",
    "r",
    encoding="utf-8",
) as f:
    _BROWSER_AGENT_DEFAULT_SYS_PROMPT = f.read()
with open(
    "build_in_prompt/browser_agent_reasoning_prompt.md",
    "r",
    encoding="utf-8",
) as f:
    _BROWSER_AGENT_DEFAULT_REASONING_PROMPT = f.read()
with open(
    "build_in_prompt/browser_agent_task_decomposition_prompt.md",  # noqa: E501 pylint: disable=C0301
    "r",
    encoding="utf-8",
) as f:
    _BROWSER_AGENT_DEFAULT_TASK_DECOMPOSITION_PROMPT = f.read()


class BrowserAgent(
    ReActAgent,
    CoreMixin,
    MemoryMixin,
    EvaluationMixin,
    UtilsMixin,
    SearchMixin,
):
    """
    Browser Agent that extends ReActAgent with browser-specific capabilities.

    The agent leverages MCP (Model Context Protocol) servers to access browser
    tools with Playwright, enabling sophisticated web automation tasks.

    Example:
        .. code-block:: python

            agent = BrowserAgent(
                name="web_navigator",
                model=my_chat_model,
                formatter=my_formatter,
                memory=my_memory,
                toolkit=browser_toolkit,
                start_url="https://example.com"
            )

            response = await agent.reply("Search for Python tutorials")
    """

    # ReActAgent precedes the mixins in the MRO and also defines _acting;
    # bind the browser-specific implementation explicitly so tool-result
    # cleaning and subtask-manager memory pruning take effect.
    _acting = UtilsMixin._acting

    def __init__(
        self,
        name: str,
        model: ChatModelBase,
        formatter: FormatterBase,
        memory: MemoryBase,
        toolkit: Toolkit,
        sys_prompt: str = _BROWSER_AGENT_DEFAULT_SYS_PROMPT,
        max_iters: int = 50,
        start_url: Optional[str] = "https://www.google.com",
        reasoning_prompt: str = _BROWSER_AGENT_DEFAULT_REASONING_PROMPT,
        task_decomposition_prompt: str = (
            _BROWSER_AGENT_DEFAULT_TASK_DECOMPOSITION_PROMPT
        ),
        token_counter: TokenCounterBase = OpenAITokenCounter("gpt-4o"),
        max_mem_length: int = 20,
        use_search_reply: bool = False,
        search_branch_factor: int = 5,
        search_max_depth: int = 5,
        replay_mode: str = "mixed",
        use_background_reasoning: bool = True,
        page_action_memory_path: Optional[str] = (
            "./logs/page_action_memory.json"
        ),
    ) -> None:
        """Initialize the Browser Agent."""
        # Initialize core functionality
        CoreMixin.__init__(
            self,
            name=name,
            model=model,
            formatter=formatter,
            memory=memory,
            toolkit=toolkit,
            sys_prompt=sys_prompt,
            max_iters=max_iters,
            start_url=start_url,
            reasoning_prompt=reasoning_prompt,
            task_decomposition_prompt=task_decomposition_prompt,
            token_counter=token_counter,
            max_mem_length=max_mem_length,
            use_search_reply=use_search_reply,
            search_branch_factor=search_branch_factor,
            search_max_depth=search_max_depth,
            replay_mode=replay_mode,
            use_background_reasoning=use_background_reasoning,
            page_action_memory_path=page_action_memory_path,
        )

        super().__init__(
            name=name,
            sys_prompt=sys_prompt,
            model=model,
            formatter=formatter,
            memory=memory,
            toolkit=toolkit,
            max_iters=max_iters,
        )

        self.toolkit.register_tool_function(self.browser_subtask_manager)

        if (
            self.model.model_name.startswith("qvq")
            or "-vl" in self.model.model_name
        ):
            # If the model supports multimodal input,
            # prepare a directory for screenshots
            screenshot_dir = os.path.join(
                "./logs/screenshots/",
                "tmp" + "_browser_agent",
            )
            os.makedirs(screenshot_dir, exist_ok=True)
            self.screenshot_dir = screenshot_dir

    async def reply(
        self,
        msg: Msg | list[Msg] | None = None,
        structured_model: Type[BaseModel] | None = None,
    ) -> Msg:
        """
        Process a message and return a response.

        Args:
            msg (`Msg | list[Msg] | None`, optional):
                The input message(s) to the agent.
            structured_model (`Type[BaseModel] | None`, optional):
                The required structured output model. If provided, the agent
                is expected to generate structured output in the `metadata`
                field of the output message.

        Returns:
            Msg: The response message.
        """
        self.init_query = (
            msg.content
            if isinstance(msg, Msg)
            else msg[0].content if isinstance(msg, list) else ""
        )
        if self.start_url and not self._has_initial_navigated:
            await self._navigate_to_start_url()
            self._has_initial_navigated = True

        # Choose between default reply and search reply based on the parameter
        if self.use_search_reply:
            return await self._reply_with_search(msg, structured_model)
        else:
            return await self._default_reply(msg, structured_model)

    async def _default_reply(
        self,
        msg: Msg | list[Msg] | None = None,
        structured_model: Type[BaseModel] | None = None,
    ) -> Msg:
        """Default reply method using the original reasoning-acting loop."""
        msg = await self._task_decomposition_and_reformat(msg)
        # original reply function
        await self.memory.add(msg)
        self._required_structured_model = structured_model
        # Record structured output model if provided
        if structured_model:
            self.toolkit.set_extended_model(
                self.finish_function_name,
                structured_model,
            )
        # The reasoning-acting loop
        reply_msg = None
        for iter_n in range(self.max_iters):
            self.iter_n = iter_n + 1
            await self._summarize_mem()
            observe_msg = await self._build_observation()
            msg_reasoning = await self._reasoning(observe_msg)

            futures = [
                self._acting(tool_call)
                for tool_call in msg_reasoning.get_content_blocks(
                    "tool_use",
                )
            ]

            # Parallel tool calls or not
            if self.parallel_tool_calls:
                acting_responses = await asyncio.gather(*futures)

            else:
                # Sequential tool calls
                acting_responses = [await _ for _ in futures]

            # Find the first non-None replying message from the acting
            for acting_msg in acting_responses:
                reply_msg = reply_msg or acting_msg

            if reply_msg:
                break

        # When the maximum iterations are reached
        if not reply_msg:
            await self._memory_summarizing()

        await self.memory.add(reply_msg)
        return reply_msg

    async def _reply_with_search(
        self,
        msg: Msg | list[Msg] | None = None,
        structured_model: Type[BaseModel] | None = None,
    ) -> Msg:
        """Search reply method using the search-based approach."""
        msg = await self._task_decomposition_and_reformat(msg)
        self._required_structured_model = structured_model

        # Record structured output model if provided
        if structured_model:
            self.toolkit.set_extended_model(
                self.finish_function_name,
                structured_model,
            )

        # Use the search method from SearchMixin
        return await SearchMixin._reply_with_search(self, msg)

    async def _reasoning(
        self,
        observe_msg: Msg | None = None,
    ) -> Msg:
        """Perform the reasoning process."""
        # Get current URL for the page action memory lookup
        current_url = await self._extract_current_url()

        # Prepare messages including the page action memory
        msgs = [
            Msg("system", self.sys_prompt, "system"),
            *await self.memory.get_memory(),
        ]

        # Add explored actions of the current page if available
        if (
            current_url in self.search_history
            and self.search_history[current_url]
        ):
            search_history_msg = Msg(
                "system",
                f"Search history for {current_url}: "
                f"{self.search_history[current_url]}",
                "system",
            )
            msgs.append(search_history_msg)

        msgs.append(observe_msg)
        prompt = await self.formatter.format(msgs=msgs)

        res = await self.model(
            prompt,
            tools=self.toolkit.get_json_schemas(),
        )
        # handle output from the model
        interrupted_by_user = False
        msg = None
        try:
            if self.model.stream:
                msg = Msg(self.name, [], "assistant")
                async for content_chunk in res:
                    msg.content = content_chunk.content
                    await self.print(msg, False)
                await self.print(msg, True)

            else:
                msg = Msg(self.name, list(res.content), "assistant")
                await self.print(msg, True)

            return msg

        except asyncio.CancelledError as e:
            interrupted_by_user = True
            raise e from None

        finally:
            await self.memory.add(msg)
            await self._update_chunk_observation_status(
                output_msg=msg,
            )
            # Post-process for user interruption
            if interrupted_by_user and msg:
                # Fake tool results
                tool_use_blocks: list = (
                    msg.get_content_blocks(  # pylint: disable=E1133
                        "tool_use",
                    )
                )
                for tool_call in tool_use_blocks:  # pylint: disable=E1133
                    msg_res = Msg(
                        "system",
                        [
                            ToolResultBlock(
                                type="tool_result",
                                id=tool_call["id"],
                                name=tool_call["name"],
                                output="The tool call has been interrupted "
                                "by the user.",
                            ),
                        ],
                        "system",
                    )
                    await self.memory.add(msg_res)
                    await self.print(msg_res, True)

    async def _build_observation(
        self,
    ) -> Msg:
        """Get a snapshot in text before reasoning"""

        image_path: Optional[str] = None
        if (
            self.model.model_name.startswith("qvq")
            or "-vl" in self.model.model_name
        ):
            # If the model supports multimodal input, take a screenshot
            # and pass it to the observation message
            img_path = os.path.join(
                self.screenshot_dir,
                f"screenshot_{self.iter_n}.png",
            )
            # if the img_path already exists,
            # reuse it instead of taking a screenshot again
            if os.path.exists(img_path):
                image_path = img_path
            else:
                image_path = await self._get_screenshot(img_path)
        if not self.chunk_continue_status:
            self.snapshot_in_chunk = await self._get_snapshot_in_text()

        observe_msg = self.observe_by_chunk(image_path)
        return observe_msg