Upload 5 files

src/agent/__init__.py

@@ -0,0 +1,6 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2025/1/1
+# @Author  : wenshao
+# @Email   : wenshaoguo1026@gmail.com
+# @Project : browser-use-webui
+# @FileName: __init__.py.py

src/agent/custom_agent.py

@@ -1,480 +1,480 @@
-# -*- coding: utf-8 -*-
-# @Time    : 2025/1/2
-# @Author  : wenshao
-# @ProjectName: browser-use-webui
-# @FileName: custom_agent.py
-
-import json
-import logging
-import pdb
-import traceback
-from typing import Optional, Type
-from PIL import Image, ImageDraw, ImageFont
-import os
-import base64
-import io
-
-from browser_use.agent.prompts import SystemPrompt
-from browser_use.agent.service import Agent
-from browser_use.agent.views import (
-    ActionResult,
-    AgentHistoryList,
-    AgentOutput,
-    AgentHistory,
-)
-from browser_use.browser.browser import Browser
-from browser_use.browser.context import BrowserContext
-from browser_use.browser.views import BrowserStateHistory
-from browser_use.controller.service import Controller
-from browser_use.telemetry.views import (
-    AgentEndTelemetryEvent,
-    AgentRunTelemetryEvent,
-    AgentStepErrorTelemetryEvent,
-)
-from browser_use.utils import time_execution_async
-from langchain_core.language_models.chat_models import BaseChatModel
-from langchain_core.messages import (
-    BaseMessage,
-)
-from src.utils.agent_state import AgentState
-
-from .custom_massage_manager import CustomMassageManager
-from .custom_views import CustomAgentOutput, CustomAgentStepInfo
-
-logger = logging.getLogger(__name__)
-
-
-class CustomAgent(Agent):
-    def __init__(
-            self,
-            task: str,
-            llm: BaseChatModel,
-            add_infos: str = "",
-            browser: Browser | None = None,
-            browser_context: BrowserContext | None = None,
-            controller: Controller = Controller(),
-            use_vision: bool = True,
-            save_conversation_path: Optional[str] = None,
-            max_failures: int = 5,
-            retry_delay: int = 10,
-            system_prompt_class: Type[SystemPrompt] = SystemPrompt,
-            max_input_tokens: int = 128000,
-            validate_output: bool = False,
-            include_attributes: list[str] = [
-                "title",
-                "type",
-                "name",
-                "role",
-                "tabindex",
-                "aria-label",
-                "placeholder",
-                "value",
-                "alt",
-                "aria-expanded",
-            ],
-            max_error_length: int = 400,
-            max_actions_per_step: int = 10,
-            tool_call_in_content: bool = True,
-            agent_state: AgentState = None,
-    ):
-        super().__init__(
-            task=task,
-            llm=llm,
-            browser=browser,
-            browser_context=browser_context,
-            controller=controller,
-            use_vision=use_vision,
-            save_conversation_path=save_conversation_path,
-            max_failures=max_failures,
-            retry_delay=retry_delay,
-            system_prompt_class=system_prompt_class,
-            max_input_tokens=max_input_tokens,
-            validate_output=validate_output,
-            include_attributes=include_attributes,
-            max_error_length=max_error_length,
-            max_actions_per_step=max_actions_per_step,
-            tool_call_in_content=tool_call_in_content,
-        )
-        self.add_infos = add_infos
-        self.agent_state = agent_state
-        self.message_manager = CustomMassageManager(
-            llm=self.llm,
-            task=self.task,
-            action_descriptions=self.controller.registry.get_prompt_description(),
-            system_prompt_class=self.system_prompt_class,
-            max_input_tokens=self.max_input_tokens,
-            include_attributes=self.include_attributes,
-            max_error_length=self.max_error_length,
-            max_actions_per_step=self.max_actions_per_step,
-            tool_call_in_content=tool_call_in_content,
-        )
-
-    def _setup_action_models(self) -> None:
-        """Setup dynamic action models from controller's registry"""
-        # Get the dynamic action model from controller's registry
-        self.ActionModel = self.controller.registry.create_action_model()
-        # Create output model with the dynamic actions
-        self.AgentOutput = CustomAgentOutput.type_with_custom_actions(self.ActionModel)
-
-    def _log_response(self, response: CustomAgentOutput) -> None:
-        """Log the model's response"""
-        if "Success" in response.current_state.prev_action_evaluation:
-            emoji = "✅"
-        elif "Failed" in response.current_state.prev_action_evaluation:
-            emoji = "❌"
-        else:
-            emoji = "🤷"
-
-        logger.info(f"{emoji} Eval: {response.current_state.prev_action_evaluation}")
-        logger.info(f"🧠 New Memory: {response.current_state.important_contents}")
-        logger.info(f"⏳ Task Progress: {response.current_state.completed_contents}")
-        logger.info(f"🤔 Thought: {response.current_state.thought}")
-        logger.info(f"🎯 Summary: {response.current_state.summary}")
-        for i, action in enumerate(response.action):
-            logger.info(
-                f"🛠️  Action {i + 1}/{len(response.action)}: {action.model_dump_json(exclude_unset=True)}"
-            )
-
-    def update_step_info(
-            self, model_output: CustomAgentOutput, step_info: CustomAgentStepInfo = None
-    ):
-        """
-        update step info
-        """
-        if step_info is None:
-            return
-
-        step_info.step_number += 1
-        important_contents = model_output.current_state.important_contents
-        if (
-                important_contents
-                and "None" not in important_contents
-                and important_contents not in step_info.memory
-        ):
-            step_info.memory += important_contents + "\n"
-
-        completed_contents = model_output.current_state.completed_contents
-        if completed_contents and "None" not in completed_contents:
-            step_info.task_progress = completed_contents
-
-    @time_execution_async("--get_next_action")
-    async def get_next_action(self, input_messages: list[BaseMessage]) -> AgentOutput:
-        """Get next action from LLM based on current state"""
-        try:
-            structured_llm = self.llm.with_structured_output(self.AgentOutput, include_raw=True)
-            response: dict[str, Any] = await structured_llm.ainvoke(input_messages)  # type: ignore
-
-            parsed: AgentOutput = response['parsed']
-            # cut the number of actions to max_actions_per_step
-            parsed.action = parsed.action[: self.max_actions_per_step]
-            self._log_response(parsed)
-            self.n_steps += 1
-
-            return parsed
-        except Exception as e:
-            # If something goes wrong, try to invoke the LLM again without structured output,
-            # and Manually parse the response. Temporarily solution for DeepSeek
-            ret = self.llm.invoke(input_messages)
-            if isinstance(ret.content, list):
-                parsed_json = json.loads(ret.content[0].replace("```json", "").replace("```", ""))
-            else:
-                parsed_json = json.loads(ret.content.replace("```json", "").replace("```", ""))
-            parsed: AgentOutput = self.AgentOutput(**parsed_json)
-            if parsed is None:
-                raise ValueError(f'Could not parse response.')
-
-            # cut the number of actions to max_actions_per_step
-            parsed.action = parsed.action[: self.max_actions_per_step]
-            self._log_response(parsed)
-            self.n_steps += 1
-
-            return parsed
-
-    @time_execution_async("--step")
-    async def step(self, step_info: Optional[CustomAgentStepInfo] = None) -> None:
-        """Execute one step of the task"""
-        logger.info(f"\n📍 Step {self.n_steps}")
-        state = None
-        model_output = None
-        result: list[ActionResult] = []
-
-        try:
-            state = await self.browser_context.get_state(use_vision=self.use_vision)
-            self.message_manager.add_state_message(state, self._last_result, step_info)
-            input_messages = self.message_manager.get_messages()
-            model_output = await self.get_next_action(input_messages)
-            self.update_step_info(model_output, step_info)
-            logger.info(f"🧠 All Memory: {step_info.memory}")
-            self._save_conversation(input_messages, model_output)
-            self.message_manager._remove_last_state_message()  # we dont want the whole state in the chat history
-            self.message_manager.add_model_output(model_output)
-
-            result: list[ActionResult] = await self.controller.multi_act(
-                model_output.action, self.browser_context
-            )
-            self._last_result = result
-
-            if len(result) > 0 and result[-1].is_done:
-                logger.info(f"📄 Result: {result[-1].extracted_content}")
-
-            self.consecutive_failures = 0
-
-        except Exception as e:
-            result = self._handle_step_error(e)
-            self._last_result = result
-
-        finally:
-            if not result:
-                return
-            for r in result:
-                if r.error:
-                    self.telemetry.capture(
-                        AgentStepErrorTelemetryEvent(
-                            agent_id=self.agent_id,
-                            error=r.error,
-                        )
-                    )
-            if state:
-                self._make_history_item(model_output, state, result)
-    def create_history_gif(
-            self,
-            output_path: str = 'agent_history.gif',
-            duration: int = 3000,
-            show_goals: bool = True,
-            show_task: bool = True,
-            show_logo: bool = False,
-            font_size: int = 40,
-            title_font_size: int = 56,
-            goal_font_size: int = 44,
-            margin: int = 40,
-            line_spacing: float = 1.5,
-    ) -> None:
-        """Create a GIF from the agent's history with overlaid task and goal text."""
-        if not self.history.history:
-            logger.warning('No history to create GIF from')
-            return
-
-        images = []
-        # if history is empty or first screenshot is None, we can't create a gif
-        if not self.history.history or not self.history.history[0].state.screenshot:
-            logger.warning('No history or first screenshot to create GIF from')
-            return
-
-        # Try to load nicer fonts
-        try:
-            # Try different font options in order of preference
-            font_options = ['Helvetica', 'Arial', 'DejaVuSans', 'Verdana']
-            font_loaded = False
-
-            for font_name in font_options:
-                try:
-                    import platform
-                    if platform.system() == "Windows":
-                        # Need to specify the abs font path on Windows
-                        font_name = os.path.join(os.getenv("WIN_FONT_DIR", "C:\\Windows\\Fonts"), font_name + ".ttf")
-                    regular_font = ImageFont.truetype(font_name, font_size)
-                    title_font = ImageFont.truetype(font_name, title_font_size)
-                    goal_font = ImageFont.truetype(font_name, goal_font_size)
-                    font_loaded = True
-                    break
-                except OSError:
-                    continue
-
-            if not font_loaded:
-                raise OSError('No preferred fonts found')
-
-        except OSError:
-            regular_font = ImageFont.load_default()
-            title_font = ImageFont.load_default()
-
-            goal_font = regular_font
-
-        # Load logo if requested
-        logo = None
-        if show_logo:
-            try:
-                logo = Image.open('./static/browser-use.png')
-                # Resize logo to be small (e.g., 40px height)
-                logo_height = 150
-                aspect_ratio = logo.width / logo.height
-                logo_width = int(logo_height * aspect_ratio)
-                logo = logo.resize((logo_width, logo_height), Image.Resampling.LANCZOS)
-            except Exception as e:
-                logger.warning(f'Could not load logo: {e}')
-
-        # Create task frame if requested
-        if show_task and self.task:
-            task_frame = self._create_task_frame(
-                self.task,
-                self.history.history[0].state.screenshot,
-                title_font,
-                regular_font,
-                logo,
-                line_spacing,
-            )
-            images.append(task_frame)
-
-        # Process each history item
-        for i, item in enumerate(self.history.history, 1):
-            if not item.state.screenshot:
-                continue
-
-            # Convert base64 screenshot to PIL Image
-            img_data = base64.b64decode(item.state.screenshot)
-            image = Image.open(io.BytesIO(img_data))
-
-            if show_goals and item.model_output:
-                image = self._add_overlay_to_image(
-                    image=image,
-                    step_number=i,
-                    goal_text=item.model_output.current_state.thought,
-                    regular_font=regular_font,
-                    title_font=title_font,
-                    margin=margin,
-                    logo=logo,
-                )
-
-            images.append(image)
-
-        if images:
-            # Save the GIF
-            images[0].save(
-                output_path,
-                save_all=True,
-                append_images=images[1:],
-                duration=duration,
-                loop=0,
-                optimize=False,
-            )
-            logger.info(f'Created GIF at {output_path}')
-        else:
-            logger.warning('No images found in history to create GIF')
-
-    async def run(self, max_steps: int = 100) -> AgentHistoryList:
-        """Execute the task with maximum number of steps"""
-        try:
-            logger.info(f"🚀 Starting task: {self.task}")
-
-            self.telemetry.capture(
-                AgentRunTelemetryEvent(
-                    agent_id=self.agent_id,
-                    task=self.task,
-                )
-            )
-
-            step_info = CustomAgentStepInfo(
-                task=self.task,
-                add_infos=self.add_infos,
-                step_number=1,
-                max_steps=max_steps,
-                memory="",
-                task_progress="",
-            )
-
-            for step in range(max_steps):
-                # 1) Check if stop requested
-                if self.agent_state and self.agent_state.is_stop_requested():
-                    logger.info("🛑 Stop requested by user")
-                    self._create_stop_history_item()
-                    break
-
-                # 2) Store last valid state before step
-                if self.browser_context and self.agent_state:
-                    state = await self.browser_context.get_state(use_vision=self.use_vision)
-                    self.agent_state.set_last_valid_state(state)
-
-                if self._too_many_failures():
-                    break
-
-                # 3) Do the step
-                await self.step(step_info)
-
-                if self.history.is_done():
-                    if (
-                            self.validate_output and step < max_steps - 1
-                    ):  # if last step, we dont need to validate
-                        if not await self._validate_output():
-                            continue
-
-                    logger.info("✅ Task completed successfully")
-                    break
-            else:
-                logger.info("❌ Failed to complete task in maximum steps")
-
-            return self.history
-
-        finally:
-            self.telemetry.capture(
-                AgentEndTelemetryEvent(
-                    agent_id=self.agent_id,
-                    task=self.task,
-                    success=self.history.is_done(),
-                    steps=len(self.history.history),
-                )
-            )
-            if not self.injected_browser_context:
-                await self.browser_context.close()
-
-            if not self.injected_browser and self.browser:
-                await self.browser.close()
-
-            if self.generate_gif:
-                self.create_history_gif()
-
-    def _create_stop_history_item(self):
-        """Create a history item for when the agent is stopped."""
-        try:
-            # Attempt to retrieve the last valid state from agent_state
-            state = None
-            if self.agent_state:
-                last_state = self.agent_state.get_last_valid_state()
-                if last_state:
-                    # Convert to BrowserStateHistory
-                    state = BrowserStateHistory(
-                        url=getattr(last_state, 'url', ""),
-                        title=getattr(last_state, 'title', ""),
-                        tabs=getattr(last_state, 'tabs', []),
-                        interacted_element=[None],
-                        screenshot=getattr(last_state, 'screenshot', None)
-                    )
-                else:
-                    state = self._create_empty_state()
-            else:
-                state = self._create_empty_state()
-
-            # Create a final item in the agent history indicating done
-            stop_history = AgentHistory(
-                model_output=None,
-                state=state,
-                result=[ActionResult(extracted_content=None, error=None, is_done=True)]
-            )
-            self.history.history.append(stop_history)
-
-        except Exception as e:
-            logger.error(f"Error creating stop history item: {e}")
-            # Create empty state as fallback
-            state = self._create_empty_state()
-            stop_history = AgentHistory(
-                model_output=None,
-                state=state,
-                result=[ActionResult(extracted_content=None, error=None, is_done=True)]
-            )
-            self.history.history.append(stop_history)
-
-    def _convert_to_browser_state_history(self, browser_state):
-        return BrowserStateHistory(
-            url=getattr(browser_state, 'url', ""),
-            title=getattr(browser_state, 'title', ""),
-            tabs=getattr(browser_state, 'tabs', []),
-            interacted_element=[None],
-            screenshot=getattr(browser_state, 'screenshot', None)
-        )
-
-    def _create_empty_state(self):
-        return BrowserStateHistory(
-            url="",
-            title="",
-            tabs=[],
-            interacted_element=[None],
-            screenshot=None
-        )
+# -*- coding: utf-8 -*-
+# @Time    : 2025/1/2
+# @Author  : wenshao
+# @ProjectName: browser-use-webui
+# @FileName: custom_agent.py
+
+import json
+import logging
+import pdb
+import traceback
+from typing import Optional, Type
+from PIL import Image, ImageDraw, ImageFont
+import os
+import base64
+import io
+
+from browser_use.agent.prompts import SystemPrompt
+from browser_use.agent.service import Agent
+from browser_use.agent.views import (
+    ActionResult,
+    AgentHistoryList,
+    AgentOutput,
+    AgentHistory,
+)
+from browser_use.browser.browser import Browser
+from browser_use.browser.context import BrowserContext
+from browser_use.browser.views import BrowserStateHistory
+from browser_use.controller.service import Controller
+from browser_use.telemetry.views import (
+    AgentEndTelemetryEvent,
+    AgentRunTelemetryEvent,
+    AgentStepErrorTelemetryEvent,
+)
+from browser_use.utils import time_execution_async
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import (
+    BaseMessage,
+)
+from src.utils.agent_state import AgentState
+
+from .custom_massage_manager import CustomMassageManager
+from .custom_views import CustomAgentOutput, CustomAgentStepInfo
+
+logger = logging.getLogger(__name__)
+
+
+class CustomAgent(Agent):
+    def __init__(
+            self,
+            task: str,
+            llm: BaseChatModel,
+            add_infos: str = "",
+            browser: Browser | None = None,
+            browser_context: BrowserContext | None = None,
+            controller: Controller = Controller(),
+            use_vision: bool = True,
+            save_conversation_path: Optional[str] = None,
+            max_failures: int = 5,
+            retry_delay: int = 10,
+            system_prompt_class: Type[SystemPrompt] = SystemPrompt,
+            max_input_tokens: int = 128000,
+            validate_output: bool = False,
+            include_attributes: list[str] = [
+                "title",
+                "type",
+                "name",
+                "role",
+                "tabindex",
+                "aria-label",
+                "placeholder",
+                "value",
+                "alt",
+                "aria-expanded",
+            ],
+            max_error_length: int = 400,
+            max_actions_per_step: int = 10,
+            tool_call_in_content: bool = True,
+            agent_state: AgentState = None,
+    ):
+        super().__init__(
+            task=task,
+            llm=llm,
+            browser=browser,
+            browser_context=browser_context,
+            controller=controller,
+            use_vision=use_vision,
+            save_conversation_path=save_conversation_path,
+            max_failures=max_failures,
+            retry_delay=retry_delay,
+            system_prompt_class=system_prompt_class,
+            max_input_tokens=max_input_tokens,
+            validate_output=validate_output,
+            include_attributes=include_attributes,
+            max_error_length=max_error_length,
+            max_actions_per_step=max_actions_per_step,
+            tool_call_in_content=tool_call_in_content,
+        )
+        self.add_infos = add_infos
+        self.agent_state = agent_state
+        self.message_manager = CustomMassageManager(
+            llm=self.llm,
+            task=self.task,
+            action_descriptions=self.controller.registry.get_prompt_description(),
+            system_prompt_class=self.system_prompt_class,
+            max_input_tokens=self.max_input_tokens,
+            include_attributes=self.include_attributes,
+            max_error_length=self.max_error_length,
+            max_actions_per_step=self.max_actions_per_step,
+            tool_call_in_content=tool_call_in_content,
+        )
+
+    def _setup_action_models(self) -> None:
+        """Setup dynamic action models from controller's registry"""
+        # Get the dynamic action model from controller's registry
+        self.ActionModel = self.controller.registry.create_action_model()
+        # Create output model with the dynamic actions
+        self.AgentOutput = CustomAgentOutput.type_with_custom_actions(self.ActionModel)
+
+    def _log_response(self, response: CustomAgentOutput) -> None:
+        """Log the model's response"""
+        if "Success" in response.current_state.prev_action_evaluation:
+            emoji = "✅"
+        elif "Failed" in response.current_state.prev_action_evaluation:
+            emoji = "❌"
+        else:
+            emoji = "🤷"
+
+        logger.info(f"{emoji} Eval: {response.current_state.prev_action_evaluation}")
+        logger.info(f"🧠 New Memory: {response.current_state.important_contents}")
+        logger.info(f"⏳ Task Progress: {response.current_state.completed_contents}")
+        logger.info(f"🤔 Thought: {response.current_state.thought}")
+        logger.info(f"🎯 Summary: {response.current_state.summary}")
+        for i, action in enumerate(response.action):
+            logger.info(
+                f"🛠️  Action {i + 1}/{len(response.action)}: {action.model_dump_json(exclude_unset=True)}"
+            )
+
+    def update_step_info(
+            self, model_output: CustomAgentOutput, step_info: CustomAgentStepInfo = None
+    ):
+        """
+        update step info
+        """
+        if step_info is None:
+            return
+
+        step_info.step_number += 1
+        important_contents = model_output.current_state.important_contents
+        if (
+                important_contents
+                and "None" not in important_contents
+                and important_contents not in step_info.memory
+        ):
+            step_info.memory += important_contents + "\n"
+
+        completed_contents = model_output.current_state.completed_contents
+        if completed_contents and "None" not in completed_contents:
+            step_info.task_progress = completed_contents
+
+    @time_execution_async("--get_next_action")
+    async def get_next_action(self, input_messages: list[BaseMessage]) -> AgentOutput:
+        """Get next action from LLM based on current state"""
+        try:
+            structured_llm = self.llm.with_structured_output(self.AgentOutput, include_raw=True)
+            response: dict[str, Any] = await structured_llm.ainvoke(input_messages)  # type: ignore
+
+            parsed: AgentOutput = response['parsed']
+            # cut the number of actions to max_actions_per_step
+            parsed.action = parsed.action[: self.max_actions_per_step]
+            self._log_response(parsed)
+            self.n_steps += 1
+
+            return parsed
+        except Exception as e:
+            # If something goes wrong, try to invoke the LLM again without structured output,
+            # and Manually parse the response. Temporarily solution for DeepSeek
+            ret = self.llm.invoke(input_messages)
+            if isinstance(ret.content, list):
+                parsed_json = json.loads(ret.content[0].replace("```json", "").replace("```", ""))
+            else:
+                parsed_json = json.loads(ret.content.replace("```json", "").replace("```", ""))
+            parsed: AgentOutput = self.AgentOutput(**parsed_json)
+            if parsed is None:
+                raise ValueError(f'Could not parse response.')
+
+            # cut the number of actions to max_actions_per_step
+            parsed.action = parsed.action[: self.max_actions_per_step]
+            self._log_response(parsed)
+            self.n_steps += 1
+
+            return parsed
+
+    @time_execution_async("--step")
+    async def step(self, step_info: Optional[CustomAgentStepInfo] = None) -> None:
+        """Execute one step of the task"""
+        logger.info(f"\n📍 Step {self.n_steps}")
+        state = None
+        model_output = None
+        result: list[ActionResult] = []
+
+        try:
+            state = await self.browser_context.get_state(use_vision=self.use_vision)
+            self.message_manager.add_state_message(state, self._last_result, step_info)
+            input_messages = self.message_manager.get_messages()
+            model_output = await self.get_next_action(input_messages)
+            self.update_step_info(model_output, step_info)
+            logger.info(f"🧠 All Memory: {step_info.memory}")
+            self._save_conversation(input_messages, model_output)
+            self.message_manager._remove_last_state_message()  # we dont want the whole state in the chat history
+            self.message_manager.add_model_output(model_output)
+
+            result: list[ActionResult] = await self.controller.multi_act(
+                model_output.action, self.browser_context
+            )
+            self._last_result = result
+
+            if len(result) > 0 and result[-1].is_done:
+                logger.info(f"📄 Result: {result[-1].extracted_content}")
+
+            self.consecutive_failures = 0
+
+        except Exception as e:
+            result = self._handle_step_error(e)
+            self._last_result = result
+
+        finally:
+            if not result:
+                return
+            for r in result:
+                if r.error:
+                    self.telemetry.capture(
+                        AgentStepErrorTelemetryEvent(
+                            agent_id=self.agent_id,
+                            error=r.error,
+                        )
+                    )
+            if state:
+                self._make_history_item(model_output, state, result)
+    def create_history_gif(
+            self,
+            output_path: str = 'agent_history.gif',
+            duration: int = 3000,
+            show_goals: bool = True,
+            show_task: bool = True,
+            show_logo: bool = False,
+            font_size: int = 40,
+            title_font_size: int = 56,
+            goal_font_size: int = 44,
+            margin: int = 40,
+            line_spacing: float = 1.5,
+    ) -> None:
+        """Create a GIF from the agent's history with overlaid task and goal text."""
+        if not self.history.history:
+            logger.warning('No history to create GIF from')
+            return
+
+        images = []
+        # if history is empty or first screenshot is None, we can't create a gif
+        if not self.history.history or not self.history.history[0].state.screenshot:
+            logger.warning('No history or first screenshot to create GIF from')
+            return
+
+        # Try to load nicer fonts
+        try:
+            # Try different font options in order of preference
+            font_options = ['Helvetica', 'Arial', 'DejaVuSans', 'Verdana']
+            font_loaded = False
+
+            for font_name in font_options:
+                try:
+                    import platform
+                    if platform.system() == "Windows":
+                        # Need to specify the abs font path on Windows
+                        font_name = os.path.join(os.getenv("WIN_FONT_DIR", "C:\\Windows\\Fonts"), font_name + ".ttf")
+                    regular_font = ImageFont.truetype(font_name, font_size)
+                    title_font = ImageFont.truetype(font_name, title_font_size)
+                    goal_font = ImageFont.truetype(font_name, goal_font_size)
+                    font_loaded = True
+                    break
+                except OSError:
+                    continue
+
+            if not font_loaded:
+                raise OSError('No preferred fonts found')
+
+        except OSError:
+            regular_font = ImageFont.load_default()
+            title_font = ImageFont.load_default()
+
+            goal_font = regular_font
+
+        # Load logo if requested
+        logo = None
+        if show_logo:
+            try:
+                logo = Image.open('./static/browser-use.png')
+                # Resize logo to be small (e.g., 40px height)
+                logo_height = 150
+                aspect_ratio = logo.width / logo.height
+                logo_width = int(logo_height * aspect_ratio)
+                logo = logo.resize((logo_width, logo_height), Image.Resampling.LANCZOS)
+            except Exception as e:
+                logger.warning(f'Could not load logo: {e}')
+
+        # Create task frame if requested
+        if show_task and self.task:
+            task_frame = self._create_task_frame(
+                self.task,
+                self.history.history[0].state.screenshot,
+                title_font,
+                regular_font,
+                logo,
+                line_spacing,
+            )
+            images.append(task_frame)
+
+        # Process each history item
+        for i, item in enumerate(self.history.history, 1):
+            if not item.state.screenshot:
+                continue
+
+            # Convert base64 screenshot to PIL Image
+            img_data = base64.b64decode(item.state.screenshot)
+            image = Image.open(io.BytesIO(img_data))
+
+            if show_goals and item.model_output:
+                image = self._add_overlay_to_image(
+                    image=image,
+                    step_number=i,
+                    goal_text=item.model_output.current_state.thought,
+                    regular_font=regular_font,
+                    title_font=title_font,
+                    margin=margin,
+                    logo=logo,
+                )
+
+            images.append(image)
+
+        if images:
+            # Save the GIF
+            images[0].save(
+                output_path,
+                save_all=True,
+                append_images=images[1:],
+                duration=duration,
+                loop=0,
+                optimize=False,
+            )
+            logger.info(f'Created GIF at {output_path}')
+        else:
+            logger.warning('No images found in history to create GIF')
+
+    async def run(self, max_steps: int = 100) -> AgentHistoryList:
+        """Execute the task with maximum number of steps"""
+        try:
+            logger.info(f"🚀 Starting task: {self.task}")
+
+            self.telemetry.capture(
+                AgentRunTelemetryEvent(
+                    agent_id=self.agent_id,
+                    task=self.task,
+                )
+            )
+
+            step_info = CustomAgentStepInfo(
+                task=self.task,
+                add_infos=self.add_infos,
+                step_number=1,
+                max_steps=max_steps,
+                memory="",
+                task_progress="",
+            )
+
+            for step in range(max_steps):
+                # 1) Check if stop requested
+                if self.agent_state and self.agent_state.is_stop_requested():
+                    logger.info("🛑 Stop requested by user")
+                    self._create_stop_history_item()
+                    break
+
+                # 2) Store last valid state before step
+                if self.browser_context and self.agent_state:
+                    state = await self.browser_context.get_state(use_vision=self.use_vision)
+                    self.agent_state.set_last_valid_state(state)
+
+                if self._too_many_failures():
+                    break
+
+                # 3) Do the step
+                await self.step(step_info)
+
+                if self.history.is_done():
+                    if (
+                            self.validate_output and step < max_steps - 1
+                    ):  # if last step, we dont need to validate
+                        if not await self._validate_output():
+                            continue
+
+                    logger.info("✅ Task completed successfully")
+                    break
+            else:
+                logger.info("❌ Failed to complete task in maximum steps")
+
+            return self.history
+
+        finally:
+            self.telemetry.capture(
+                AgentEndTelemetryEvent(
+                    agent_id=self.agent_id,
+                    task=self.task,
+                    success=self.history.is_done(),
+                    steps=len(self.history.history),
+                )
+            )
+            if not self.injected_browser_context:
+                await self.browser_context.close()
+
+            if not self.injected_browser and self.browser:
+                await self.browser.close()
+
+            if self.generate_gif:
+                self.create_history_gif()
+
+    def _create_stop_history_item(self):
+        """Create a history item for when the agent is stopped."""
+        try:
+            # Attempt to retrieve the last valid state from agent_state
+            state = None
+            if self.agent_state:
+                last_state = self.agent_state.get_last_valid_state()
+                if last_state:
+                    # Convert to BrowserStateHistory
+                    state = BrowserStateHistory(
+                        url=getattr(last_state, 'url', ""),
+                        title=getattr(last_state, 'title', ""),
+                        tabs=getattr(last_state, 'tabs', []),
+                        interacted_element=[None],
+                        screenshot=getattr(last_state, 'screenshot', None)
+                    )
+                else:
+                    state = self._create_empty_state()
+            else:
+                state = self._create_empty_state()
+
+            # Create a final item in the agent history indicating done
+            stop_history = AgentHistory(
+                model_output=None,
+                state=state,
+                result=[ActionResult(extracted_content=None, error=None, is_done=True)]
+            )
+            self.history.history.append(stop_history)
+
+        except Exception as e:
+            logger.error(f"Error creating stop history item: {e}")
+            # Create empty state as fallback
+            state = self._create_empty_state()
+            stop_history = AgentHistory(
+                model_output=None,
+                state=state,
+                result=[ActionResult(extracted_content=None, error=None, is_done=True)]
+            )
+            self.history.history.append(stop_history)
+
+    def _convert_to_browser_state_history(self, browser_state):
+        return BrowserStateHistory(
+            url=getattr(browser_state, 'url', ""),
+            title=getattr(browser_state, 'title', ""),
+            tabs=getattr(browser_state, 'tabs', []),
+            interacted_element=[None],
+            screenshot=getattr(browser_state, 'screenshot', None)
+        )
+
+    def _create_empty_state(self):
+        return BrowserStateHistory(
+            url="",
+            title="",
+            tabs=[],
+            interacted_element=[None],
+            screenshot=None
+        )

src/agent/custom_massage_manager.py

@@ -0,0 +1,121 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2025/1/2
+# @Author  : wenshao
+# @ProjectName: browser-use-webui
+# @FileName: custom_massage_manager.py
+
+from __future__ import annotations
+
+import logging
+from typing import List, Optional, Type
+
+from browser_use.agent.message_manager.service import MessageManager
+from browser_use.agent.message_manager.views import MessageHistory
+from browser_use.agent.prompts import SystemPrompt
+from browser_use.agent.views import ActionResult, AgentStepInfo
+from browser_use.browser.views import BrowserState
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import (
+    HumanMessage,
+    AIMessage
+)
+
+from .custom_prompts import CustomAgentMessagePrompt
+
+logger = logging.getLogger(__name__)
+
+
+class CustomMassageManager(MessageManager):
+    def __init__(
+            self,
+            llm: BaseChatModel,
+            task: str,
+            action_descriptions: str,
+            system_prompt_class: Type[SystemPrompt],
+            max_input_tokens: int = 128000,
+            estimated_tokens_per_character: int = 3,
+            image_tokens: int = 800,
+            include_attributes: list[str] = [],
+            max_error_length: int = 400,
+            max_actions_per_step: int = 10,
+            tool_call_in_content: bool = False,
+    ):
+        super().__init__(
+            llm=llm,
+            task=task,
+            action_descriptions=action_descriptions,
+            system_prompt_class=system_prompt_class,
+            max_input_tokens=max_input_tokens,
+            estimated_tokens_per_character=estimated_tokens_per_character,
+            image_tokens=image_tokens,
+            include_attributes=include_attributes,
+            max_error_length=max_error_length,
+            max_actions_per_step=max_actions_per_step,
+            tool_call_in_content=tool_call_in_content,
+        )
+
+        # Custom: Move Task info to state_message
+        self.history = MessageHistory()
+        self._add_message_with_tokens(self.system_prompt)
+        tool_calls = [
+            {
+                'name': 'CustomAgentOutput',
+                'args': {
+                    'current_state': {
+                        'prev_action_evaluation': 'Unknown - No previous actions to evaluate.',
+                        'important_contents': '',
+                        'completed_contents': '',
+                        'thought': 'Now Google is open. Need to type OpenAI to search.',
+                        'summary': 'Type OpenAI to search.',
+                    },
+                    'action': [],
+                },
+                'id': '',
+                'type': 'tool_call',
+            }
+        ]
+        if self.tool_call_in_content:
+            # openai throws error if tool_calls are not responded -> move to content
+            example_tool_call = AIMessage(
+                content=f'{tool_calls}',
+                tool_calls=[],
+            )
+        else:
+            example_tool_call = AIMessage(
+                content=f'',
+                tool_calls=tool_calls,
+            )
+
+        self._add_message_with_tokens(example_tool_call)
+
+    def add_state_message(
+            self,
+            state: BrowserState,
+            result: Optional[List[ActionResult]] = None,
+            step_info: Optional[AgentStepInfo] = None,
+    ) -> None:
+        """Add browser state as human message"""
+
+        # if keep in memory, add to directly to history and add state without result
+        if result:
+            for r in result:
+                if r.include_in_memory:
+                    if r.extracted_content:
+                        msg = HumanMessage(content=str(r.extracted_content))
+                        self._add_message_with_tokens(msg)
+                    if r.error:
+                        msg = HumanMessage(
+                            content=str(r.error)[-self.max_error_length:]
+                        )
+                        self._add_message_with_tokens(msg)
+                    result = None  # if result in history, we dont want to add it again
+
+        # otherwise add state message and result to next message (which will not stay in memory)
+        state_message = CustomAgentMessagePrompt(
+            state,
+            result,
+            include_attributes=self.include_attributes,
+            max_error_length=self.max_error_length,
+            step_info=step_info,
+        ).get_user_message()
+        self._add_message_with_tokens(state_message)

src/agent/custom_prompts.py

@@ -0,0 +1,205 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2025/1/2
+# @Author  : wenshao
+# @ProjectName: browser-use-webui
+# @FileName: custom_prompts.py
+
+from typing import List, Optional
+
+from browser_use.agent.prompts import SystemPrompt
+from browser_use.agent.views import ActionResult
+from browser_use.browser.views import BrowserState
+from langchain_core.messages import HumanMessage, SystemMessage
+
+from .custom_views import CustomAgentStepInfo
+
+
+class CustomSystemPrompt(SystemPrompt):
+    def important_rules(self) -> str:
+        """
+        Returns the important rules for the agent.
+        """
+        text = """
+    1. RESPONSE FORMAT: You must ALWAYS respond with valid JSON in this exact format:
+       {
+         "current_state": {
+           "prev_action_evaluation": "Success|Failed|Unknown - Analyze the current elements and the image to check if the previous goals/actions are successful like intended by the task. Ignore the action result. The website is the ground truth. Also mention if something unexpected happened like new suggestions in an input field. Shortly state why/why not. Note that the result you output must be consistent with the reasoning you output afterwards. If you consider it to be 'Failed,' you should reflect on this during your thought.",
+           "important_contents": "Output important contents closely related to user\'s instruction or task on the current page. If there is, please output the contents. If not, please output empty string ''.",
+           "completed_contents": "Update the input Task Progress. Completed contents is a general summary of the current contents that have been completed. Just summarize the contents that have been actually completed based on the current page and the history operations. Please list each completed item individually, such as: 1. Input username. 2. Input Password. 3. Click confirm button",
+           "thought": "Think about the requirements that have been completed in previous operations and the requirements that need to be completed in the next one operation. If the output of prev_action_evaluation is 'Failed', please reflect and output your reflection here. If you think you have entered the wrong page, consider to go back to the previous page in next action.",
+           "summary": "Please generate a brief natural language description for the operation in next actions based on your Thought."
+         },
+         "action": [
+           {
+             "action_name": {
+               // action-specific parameters
+             }
+           },
+           // ... more actions in sequence
+         ]
+       }
+
+    2. ACTIONS: You can specify multiple actions to be executed in sequence. 
+
+       Common action sequences:
+       - Form filling: [
+           {"input_text": {"index": 1, "text": "username"}},
+           {"input_text": {"index": 2, "text": "password"}},
+           {"click_element": {"index": 3}}
+         ]
+       - Navigation and extraction: [
+           {"open_new_tab": {}},
+           {"go_to_url": {"url": "https://example.com"}},
+           {"extract_page_content": {}}
+         ]
+
+
+    3. ELEMENT INTERACTION:
+       - Only use indexes that exist in the provided element list
+       - Each element has a unique index number (e.g., "33[:]<button>")
+       - Elements marked with "_[:]" are non-interactive (for context only)
+
+    4. NAVIGATION & ERROR HANDLING:
+       - If no suitable elements exist, use other functions to complete the task
+       - If stuck, try alternative approaches
+       - Handle popups/cookies by accepting or closing them
+       - Use scroll to find elements you are looking for
+
+    5. TASK COMPLETION:
+       - If you think all the requirements of user\'s instruction have been completed and no further operation is required, output the done action to terminate the operation process.
+       - Don't hallucinate actions.
+       - If the task requires specific information - make sure to include everything in the done function. This is what the user will see.
+       - If you are running out of steps (current step), think about speeding it up, and ALWAYS use the done action as the last action.
+
+    6. VISUAL CONTEXT:
+       - When an image is provided, use it to understand the page layout
+       - Bounding boxes with labels correspond to element indexes
+       - Each bounding box and its label have the same color
+       - Most often the label is inside the bounding box, on the top right
+       - Visual context helps verify element locations and relationships
+       - sometimes labels overlap, so use the context to verify the correct element
+
+    7. Form filling:
+       - If you fill an input field and your action sequence is interrupted, most often a list with suggestions poped up under the field and you need to first select the right element from the suggestion list.
+
+    8. ACTION SEQUENCING:
+       - Actions are executed in the order they appear in the list 
+       - Each action should logically follow from the previous one
+       - If the page changes after an action, the sequence is interrupted and you get the new state.
+       - If content only disappears the sequence continues.
+       - Only provide the action sequence until you think the page will change.
+       - Try to be efficient, e.g. fill forms at once, or chain actions where nothing changes on the page like saving, extracting, checkboxes...
+       - only use multiple actions if it makes sense. 
+    """
+        text += f"   - use maximum {self.max_actions_per_step} actions per sequence"
+        return text
+
+    def input_format(self) -> str:
+        return """
+    INPUT STRUCTURE:
+    1. Task: The user\'s instructions you need to complete.
+    2. Hints(Optional): Some hints to help you complete the user\'s instructions.
+    3. Memory: Important contents are recorded during historical operations for use in subsequent operations.
+    4. Task Progress: Up to the current page, the content you have completed can be understood as the progress of the task.
+    5. Current URL: The webpage you're currently on
+    6. Available Tabs: List of open browser tabs
+    7. Interactive Elements: List in the format:
+       index[:]<element_type>element_text</element_type>
+       - index: Numeric identifier for interaction
+       - element_type: HTML element type (button, input, etc.)
+       - element_text: Visible text or element description
+
+    Example:
+    33[:]<button>Submit Form</button>
+    _[:] Non-interactive text
+
+
+    Notes:
+    - Only elements with numeric indexes are interactive
+    - _[:] elements provide context but cannot be interacted with
+    """
+
+    def get_system_message(self) -> SystemMessage:
+        """
+        Get the system prompt for the agent.
+
+        Returns:
+            str: Formatted system prompt
+        """
+        time_str = self.current_date.strftime("%Y-%m-%d %H:%M")
+
+        AGENT_PROMPT = f"""You are a precise browser automation agent that interacts with websites through structured commands. Your role is to:
+    1. Analyze the provided webpage elements and structure
+    2. Plan a sequence of actions to accomplish the given task
+    3. Respond with valid JSON containing your action sequence and state assessment
+
+    Current date and time: {time_str}
+
+    {self.input_format()}
+
+    {self.important_rules()}
+
+    Functions:
+    {self.default_action_description}
+
+    Remember: Your responses must be valid JSON matching the specified format. Each action in the sequence must be valid."""
+        return SystemMessage(content=AGENT_PROMPT)
+
+
+class CustomAgentMessagePrompt:
+    def __init__(
+            self,
+            state: BrowserState,
+            result: Optional[List[ActionResult]] = None,
+            include_attributes: list[str] = [],
+            max_error_length: int = 400,
+            step_info: Optional[CustomAgentStepInfo] = None,
+    ):
+        self.state = state
+        self.result = result
+        self.max_error_length = max_error_length
+        self.include_attributes = include_attributes
+        self.step_info = step_info
+
+    def get_user_message(self) -> HumanMessage:
+        state_description = f"""
+    1. Task: {self.step_info.task}
+    2. Hints(Optional): 
+    {self.step_info.add_infos}
+    3. Memory: 
+    {self.step_info.memory}
+    4. Task Progress: 
+    {self.step_info.task_progress}
+    5. Current url: {self.state.url}
+    6. Available tabs:
+    {self.state.tabs}
+    7. Interactive elements:
+    {self.state.element_tree.clickable_elements_to_string(include_attributes=self.include_attributes)}
+            """
+
+        if self.result:
+            for i, result in enumerate(self.result):
+                if result.extracted_content:
+                    state_description += f"\nResult of action {i + 1}/{len(self.result)}: {result.extracted_content}"
+                if result.error:
+                    # only use last 300 characters of error
+                    error = result.error[-self.max_error_length:]
+                    state_description += (
+                        f"\nError of action {i + 1}/{len(self.result)}: ...{error}"
+                    )
+
+        if self.state.screenshot:
+            # Format message for vision model
+            return HumanMessage(
+                content=[
+                    {"type": "text", "text": state_description},
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": f"data:image/png;base64,{self.state.screenshot}"
+                        },
+                    },
+                ]
+            )
+
+        return HumanMessage(content=state_description)

src/agent/custom_views.py

@@ -0,0 +1,59 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2025/1/2
+# @Author  : wenshao
+# @ProjectName: browser-use-webui
+# @FileName: custom_views.py
+
+from dataclasses import dataclass
+from typing import Type
+
+from browser_use.agent.views import AgentOutput
+from browser_use.controller.registry.views import ActionModel
+from pydantic import BaseModel, ConfigDict, Field, create_model
+
+
+@dataclass
+class CustomAgentStepInfo:
+    step_number: int
+    max_steps: int
+    task: str
+    add_infos: str
+    memory: str
+    task_progress: str
+
+
+class CustomAgentBrain(BaseModel):
+    """Current state of the agent"""
+
+    prev_action_evaluation: str
+    important_contents: str
+    completed_contents: str
+    thought: str
+    summary: str
+
+
+class CustomAgentOutput(AgentOutput):
+    """Output model for agent
+
+    @dev note: this model is extended with custom actions in AgentService. You can also use some fields that are not in this model as provided by the linter, as long as they are registered in the DynamicActions model.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    current_state: CustomAgentBrain
+    action: list[ActionModel]
+
+    @staticmethod
+    def type_with_custom_actions(
+        custom_actions: Type[ActionModel],
+    ) -> Type["CustomAgentOutput"]:
+        """Extend actions with custom actions"""
+        return create_model(
+            "AgentOutput",
+            __base__=CustomAgentOutput,
+            action=(
+                list[custom_actions],
+                Field(...),
+            ),  # Properly annotated field with no default
+            __module__=CustomAgentOutput.__module__,
+        )