update to bu=0.1.40

src/agent/custom_agent.py

@@ -2,21 +2,31 @@ import json
 import logging
 import pdb
 import traceback
-from typing import Optional, Type, List, Dict, Any, Callable
+from typing import Any, Awaitable, Callable, Dict, Generic, List, Optional, Type, TypeVar
 from PIL import Image, ImageDraw, ImageFont
 import os
 import base64
 import io
+import asyncio
+import time
 import platform
 from browser_use.agent.prompts import SystemPrompt, AgentMessagePrompt
 from browser_use.agent.service import Agent
+from browser_use.agent.message_manager.utils import convert_input_messages, extract_json_from_model_output, \
+    save_conversation
 from browser_use.agent.views import (
     ActionResult,
-    ActionModel,
+    AgentError,
+    AgentHistory,
     AgentHistoryList,
     AgentOutput,
-    AgentHistory,
+    AgentSettings,
+    AgentState,
+    AgentStepInfo,
+    StepMetadata,
+    ToolCallingMethod,
 )
+from browser_use.agent.gif import create_history_gif
 from browser_use.browser.browser import Browser
 from browser_use.browser.context import BrowserContext
 from browser_use.browser.views import BrowserStateHistory
@@ -33,26 +43,58 @@ from langchain_core.messages import (
     HumanMessage,
     AIMessage
 )
+from browser_use.browser.views import BrowserState, BrowserStateHistory
 from browser_use.agent.prompts import PlannerPrompt
 
 from json_repair import repair_json
 from src.utils.agent_state import AgentState
 
-from .custom_message_manager import CustomMessageManager
-from .custom_views import CustomAgentOutput, CustomAgentStepInfo
+from .custom_message_manager import CustomMessageManager, CustomMessageManagerSettings
+from .custom_views import CustomAgentOutput, CustomAgentStepInfo, CustomAgentState
 
 logger = logging.getLogger(__name__)
 
 
+def _log_response(response: CustomAgentOutput) -> None:
+    """Log the model's response"""
+    if "Success" in response.current_state.evaluation_previous_goal:
+        emoji = "✅"
+    elif "Failed" in response.current_state.evaluation_previous_goal:
+        emoji = "❌"
+    else:
+        emoji = "🤷"
+
+    logger.info(f"{emoji} Eval: {response.current_state.evaluation_previous_goal}")
+    logger.info(f"🧠 New Memory: {response.current_state.important_contents}")
+    logger.info(f"🤔 Thought: {response.current_state.thought}")
+    logger.info(f"🎯 Next Goal: {response.current_state.next_goal}")
+    for i, action in enumerate(response.action):
+        logger.info(
+            f"🛠️  Action {i + 1}/{len(response.action)}: {action.model_dump_json(exclude_unset=True)}"
+        )
+
+
+Context = TypeVar('Context')
+
+
 class CustomAgent(Agent):
     def __init__(
             self,
             task: str,
             llm: BaseChatModel,
             add_infos: str = "",
+            # Optional parameters
             browser: Browser | None = None,
             browser_context: BrowserContext | None = None,
-            controller: Controller = Controller(),
+            controller: Controller[Context] = Controller(),
+            # Initial agent run parameters
+            sensitive_data: Optional[Dict[str, str]] = None,
+            initial_actions: Optional[List[Dict[str, Dict[str, Any]]]] = None,
+            # Cloud Callbacks
+            register_new_step_callback: Callable[['BrowserState', 'AgentOutput', int], Awaitable[None]] | None = None,
+            register_done_callback: Callable[['AgentHistoryList'], Awaitable[None]] | None = None,
+            register_external_agent_status_raise_error_callback: Callable[[], Awaitable[bool]] | None = None,
+            # Agent settings
             use_vision: bool = True,
             use_vision_for_planner: bool = False,
             save_conversation_path: Optional[str] = None,
@@ -64,53 +106,40 @@ class CustomAgent(Agent):
             max_input_tokens: int = 128000,
             validate_output: bool = False,
             message_context: Optional[str] = None,
-            generate_gif: bool | str = True,
-            sensitive_data: Optional[Dict[str, str]] = None,
+            generate_gif: bool | str = False,
             available_file_paths: Optional[list[str]] = None,
             include_attributes: list[str] = [
                 'title',
                 'type',
                 'name',
                 'role',
-                'tabindex',
                 'aria-label',
                 'placeholder',
                 'value',
                 'alt',
                 'aria-expanded',
+                'data-date-format',
             ],
-            max_error_length: int = 400,
             max_actions_per_step: int = 10,
-            tool_call_in_content: bool = True,
-            initial_actions: Optional[List[Dict[str, Dict[str, Any]]]] = None,
-            # Cloud Callbacks
-            register_new_step_callback: Callable[['BrowserState', 'AgentOutput', int], None] | None = None,
-            register_done_callback: Callable[['AgentHistoryList'], None] | None = None,
-            tool_calling_method: Optional[str] = 'auto',
+            tool_calling_method: Optional[ToolCallingMethod] = 'auto',
             page_extraction_llm: Optional[BaseChatModel] = None,
             planner_llm: Optional[BaseChatModel] = None,
             planner_interval: int = 1,  # Run planner every N steps
+            # Inject state
+            injected_agent_state: Optional[AgentState] = None,
+            context: Context | None = None,
     ):
-
-        # Load sensitive data from environment variables
-        env_sensitive_data = {}
-        for key, value in os.environ.items():
-            if key.startswith('SENSITIVE_'):
-                env_key = key.replace('SENSITIVE_', '', 1).lower()
-                env_sensitive_data[env_key] = value
-    
-        # Merge environment variables with provided sensitive_data
-        if sensitive_data is None:
-            sensitive_data = {}
-        sensitive_data = {**env_sensitive_data, **sensitive_data}  # Provided data takes precedence
-
-
-        super().__init__(
+        super(CustomAgent, self).__init__(
             task=task,
             llm=llm,
             browser=browser,
             browser_context=browser_context,
             controller=controller,
+            sensitive_data=sensitive_data,
+            initial_actions=initial_actions,
+            register_new_step_callback=register_new_step_callback,
+            register_done_callback=register_done_callback,
+            register_external_agent_status_raise_error_callback=register_external_agent_status_raise_error_callback,
             use_vision=use_vision,
             use_vision_for_planner=use_vision_for_planner,
             save_conversation_path=save_conversation_path,
@@ -122,47 +151,34 @@ class CustomAgent(Agent):
             validate_output=validate_output,
             message_context=message_context,
             generate_gif=generate_gif,
-            sensitive_data=sensitive_data,
             available_file_paths=available_file_paths,
             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,
-            initial_actions=initial_actions,
-            register_new_step_callback=register_new_step_callback,
-            register_done_callback=register_done_callback,
             tool_calling_method=tool_calling_method,
+            page_extraction_llm=page_extraction_llm,
             planner_llm=planner_llm,
-            planner_interval=planner_interval
+            planner_interval=planner_interval,
+            injected_agent_state=injected_agent_state,
+            context=context,
         )
-        if self.model_name in ["deepseek-reasoner"] or "deepseek-r1" in self.model_name:
-            # deepseek-reasoner does not support function calling
-            self.use_deepseek_r1 = True
-            # deepseek-reasoner only support 64000 context
-            self.max_input_tokens = 64000
-        else:
-            self.use_deepseek_r1 = False
-
-        # record last actions
-        self._last_actions = None
-        # record extract content
-        self.extracted_content = ""
-        # custom new info
+        self.state = injected_agent_state or CustomAgentState()
         self.add_infos = add_infos
 
-        self.agent_prompt_class = agent_prompt_class
-        self.message_manager = CustomMessageManager(
-            llm=self.llm,
-            task=self.task,
-            action_descriptions=self.controller.registry.get_prompt_description(),
-            system_prompt_class=self.system_prompt_class,
-            agent_prompt_class=agent_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,
-            message_context=self.message_context,
-            sensitive_data=self.sensitive_data
+        self._message_manager = CustomMessageManager(
+            task=task,
+            system_message=self.settings.system_prompt_class(
+                self.available_actions,
+                max_actions_per_step=self.settings.max_actions_per_step,
+            ).get_system_message(),
+            settings=CustomMessageManagerSettings(
+                max_input_tokens=self.settings.max_input_tokens,
+                include_attributes=self.settings.include_attributes,
+                message_context=self.settings.message_context,
+                sensitive_data=sensitive_data,
+                available_file_paths=self.settings.available_file_paths,
+                agent_prompt_class=agent_prompt_class
+            ),
+            state=self.state.message_manager_state,
         )
 
     def _setup_action_models(self) -> None:
@@ -172,26 +188,6 @@ class CustomAgent(Agent):
         # 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: \n{response.current_state.task_progress}")
-        logger.info(f"📋 Future Plans: \n{response.current_state.future_plans}")
-        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
     ):
@@ -210,14 +206,6 @@ class CustomAgent(Agent):
         ):
             step_info.memory += important_contents + "\n"
 
-        task_progress = model_output.current_state.task_progress
-        if task_progress and "None" not in task_progress:
-            step_info.task_progress = task_progress
-
-        future_plans = model_output.current_state.future_plans
-        if future_plans and "None" not in future_plans:
-            step_info.future_plans = future_plans
-
         logger.info(f"🧠 All Memory: \n{step_info.memory}")
 
     @time_execution_async("--get_next_action")
@@ -246,27 +234,26 @@ class CustomAgent(Agent):
             logger.debug(ai_message.content)
             raise ValueError('Could not parse response.')
 
-        # Limit actions to maximum allowed per step
-        parsed.action = parsed.action[: self.max_actions_per_step]
-        self._log_response(parsed)
-        self.n_steps += 1
-
+        # cut the number of actions to max_actions_per_step if needed
+        if len(parsed.action) > self.settings.max_actions_per_step:
+            parsed.action = parsed.action[: self.settings.max_actions_per_step]
+        _log_response(parsed)
         return parsed
 
     async def _run_planner(self) -> Optional[str]:
         """Run the planner to analyze state and suggest next steps"""
         # Skip planning if no planner_llm is set
-        if not self.planner_llm:
+        if not self.settings.planner_llm:
             return None
 
         # Create planner message history using full message history
         planner_messages = [
-            PlannerPrompt(self.action_descriptions).get_system_message(),
+            PlannerPrompt(self.controller.registry.get_prompt_description()).get_system_message(),
             *self.message_manager.get_messages()[1:],  # Use full message history except the first
         ]
 
-        if not self.use_vision_for_planner and self.use_vision:
-            last_state_message = planner_messages[-1]
+        if not self.settings.use_vision_for_planner and self.settings.use_vision:
+            last_state_message: HumanMessage = planner_messages[-1]
             # remove image from last state message
             new_msg = ''
             if isinstance(last_state_message.content, list):
@@ -281,16 +268,17 @@ class CustomAgent(Agent):
             planner_messages[-1] = HumanMessage(content=new_msg)
 
         # Get planner output
-        response = await self.planner_llm.ainvoke(planner_messages)
-        plan = response.content
-        last_state_message = planner_messages[-1]
-        # remove image from last state message
-        if isinstance(last_state_message.content, list):
-            for msg in last_state_message.content:
-                if msg['type'] == 'text':
-                    msg['text'] += f"\nPlanning Agent outputs plans:\n {plan}\n"
-        else:
-            last_state_message.content += f"\nPlanning Agent outputs plans:\n {plan}\n "
+        response = await self.settings.planner_llm.ainvoke(planner_messages)
+        plan = str(response.content)
+        last_state_message = self.message_manager.get_messages()[-1]
+        if isinstance(last_state_message, HumanMessage):
+            # remove image from last state message
+            if isinstance(last_state_message.content, list):
+                for msg in last_state_message.content:
+                    if msg['type'] == 'text':
+                        msg['text'] += f"\nPlanning Agent outputs plans:\n {plan}\n"
+            else:
+                last_state_message.content += f"\nPlanning Agent outputs plans:\n {plan}\n "
 
         try:
             plan_json = json.loads(plan.replace("```json", "").replace("```", ""))
@@ -306,86 +294,91 @@ class CustomAgent(Agent):
         except Exception as e:
             logger.debug(f'Error parsing planning analysis: {e}')
             logger.info(f'📋 Plans: {plan}')
+        return plan
 
     @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}")
+        logger.info(f"\n📍 Step {self.state.n_steps}")
         state = None
         model_output = None
         result: list[ActionResult] = []
-        actions: list[ActionModel] = []
+        step_start_time = time.time()
+        tokens = 0
 
         try:
             state = await self.browser_context.get_state()
-            self._check_if_stopped_or_paused()
+            await self._raise_if_stopped_or_paused()
 
-            self.message_manager.add_state_message(state, self._last_actions, self._last_result, step_info,
-                                                   self.use_vision)
+            self.message_manager.add_state_message(state, self.state.last_action, self.state.last_result, step_info,
+                                                   self.settings.use_vision)
 
             # Run planner at specified intervals if planner is configured
-            if self.planner_llm and self.n_steps % self.planning_interval == 0:
+            if self.settings.planner_llm and self.state.n_steps % self.settings.planner_interval == 0:
                 await self._run_planner()
             input_messages = self.message_manager.get_messages()
-            self._check_if_stopped_or_paused()
+            tokens = self._message_manager.state.history.current_tokens
+
             try:
                 model_output = await self.get_next_action(input_messages)
-                if self.register_new_step_callback:
-                    self.register_new_step_callback(state, model_output, self.n_steps)
                 self.update_step_info(model_output, step_info)
-                self._save_conversation(input_messages, model_output)
+                self.state.n_steps += 1
+
+                if self.register_new_step_callback:
+                    await self.register_new_step_callback(state, model_output, self.state.n_steps)
+
+                if self.settings.save_conversation_path:
+                    target = self.settings.save_conversation_path + f'_{self.state.n_steps}.txt'
+                    save_conversation(input_messages, model_output, target,
+                                      self.settings.save_conversation_path_encoding)
+
                 if self.model_name != "deepseek-reasoner":
                     # remove prev message
                     self.message_manager._remove_state_message_by_index(-1)
-                self._check_if_stopped_or_paused()
+                await self._raise_if_stopped_or_paused()
             except Exception as e:
                 # model call failed, remove last state message from history
                 self.message_manager._remove_state_message_by_index(-1)
                 raise e
 
-            actions: list[ActionModel] = model_output.action
-            result: list[ActionResult] = await self.controller.multi_act(
-                actions,
-                self.browser_context,
-                page_extraction_llm=self.page_extraction_llm,
-                sensitive_data=self.sensitive_data,
-                check_break_if_paused=lambda: self._check_if_stopped_or_paused(),
-                available_file_paths=self.available_file_paths,
-            )
-            if len(result) != len(actions):
-                # I think something changes, such information should let LLM know
-                for ri in range(len(result), len(actions)):
-                    result.append(ActionResult(extracted_content=None,
-                                               include_in_memory=True,
-                                               error=f"{actions[ri].model_dump_json(exclude_unset=True)} is Failed to execute. \
-                                                    Something new appeared after action {actions[len(result) - 1].model_dump_json(exclude_unset=True)}",
-                                               is_done=False))
+            result: list[ActionResult] = await self.multi_act(model_output.action)
             for ret_ in result:
                 if ret_.extracted_content and "Extracted page" in ret_.extracted_content:
                     # record every extracted page
-                    self.extracted_content += ret_.extracted_content
-            self._last_result = result
-            self._last_actions = actions
+                    self.state.extracted_content += ret_.extracted_content
+            self.state.last_result = result
+            self.state.last_action = model_output.action
             if len(result) > 0 and result[-1].is_done:
-                if not self.extracted_content:
-                    self.extracted_content = step_info.memory
-                result[-1].extracted_content = self.extracted_content
+                if not self.state.extracted_content:
+                    self.state.extracted_content = step_info.memory
+                result[-1].extracted_content = self.state.extracted_content
                 logger.info(f"📄 Result: {result[-1].extracted_content}")
 
-            self.consecutive_failures = 0
+            self.state.consecutive_failures = 0
+
+        except InterruptedError:
+            logger.debug('Agent paused')
+            self.state.last_result = [
+                ActionResult(
+                    error='The agent was paused - now continuing actions might need to be repeated',
+                    include_in_memory=True
+                )
+            ]
+            return
 
         except Exception as e:
             result = await self._handle_step_error(e)
-            self._last_result = result
+            self.state.last_result = result
 
         finally:
+            step_end_time = time.time()
             actions = [a.model_dump(exclude_unset=True) for a in model_output.action] if model_output else []
             self.telemetry.capture(
                 AgentStepTelemetryEvent(
-                    agent_id=self.agent_id,
-                    step=self.n_steps,
+                    agent_id=self.state.agent_id,
+                    step=self.state.n_steps,
                     actions=actions,
-                    consecutive_failures=self.consecutive_failures,
+                    consecutive_failures=self.state.consecutive_failures,
                     step_error=[r.error for r in result if r.error] if result else ['No result'],
                 )
             )
@@ -393,7 +386,13 @@ class CustomAgent(Agent):
                 return
 
             if state:
-                self._make_history_item(model_output, state, result)
+                metadata = StepMetadata(
+                    step_number=self.state.n_steps,
+                    step_start_time=step_start_time,
+                    step_end_time=step_end_time,
+                    input_tokens=tokens,
+                )
+                self._make_history_item(model_output, state, result, metadata)
 
     async def run(self, max_steps: int = 100) -> AgentHistoryList:
         """Execute the task with maximum number of steps"""
@@ -402,15 +401,8 @@ class CustomAgent(Agent):
 
             # Execute initial actions if provided
             if self.initial_actions:
-                result = await self.controller.multi_act(
-                    self.initial_actions,
-                    self.browser_context,
-                    check_for_new_elements=False,
-                    page_extraction_llm=self.page_extraction_llm,
-                    check_break_if_paused=lambda: self._check_if_stopped_or_paused(),
-                    available_file_paths=self.available_file_paths,
-                )
-                self._last_result = result
+                result = await self.multi_act(self.initial_actions, check_for_new_elements=False)
+                self.state.last_result = result
 
             step_info = CustomAgentStepInfo(
                 task=self.task,
@@ -418,43 +410,53 @@ class CustomAgent(Agent):
                 step_number=1,
                 max_steps=max_steps,
                 memory="",
-                task_progress="",
-                future_plans=""
             )
 
             for step in range(max_steps):
-                if self._too_many_failures():
+                # Check if we should stop due to too many failures
+                if self.state.consecutive_failures >= self.settings.max_failures:
+                    logger.error(f'❌ Stopping due to {self.settings.max_failures} consecutive failures')
                     break
 
-                # 3) Do the step
+                # Check control flags before each step
+                if self.state.stopped:
+                    logger.info('Agent stopped')
+                    break
+
+                while self.state.paused:
+                    await asyncio.sleep(0.2)  # Small delay to prevent CPU spinning
+                    if self.state.stopped:  # Allow stopping while paused
+                        break
+
                 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 self.state.history.is_done():
+                    if self.settings.validate_output and step < max_steps - 1:
                         if not await self._validate_output():
                             continue
 
-                    logger.info("✅ Task completed successfully")
+                    await self.log_completion()
                     break
             else:
                 logger.info("❌ Failed to complete task in maximum steps")
-                if not self.extracted_content:
-                    self.history.history[-1].result[-1].extracted_content = step_info.memory
+                if not self.state.extracted_content:
+                    self.state.history.history[-1].result[-1].extracted_content = step_info.memory
                 else:
-                    self.history.history[-1].result[-1].extracted_content = self.extracted_content
+                    self.state.history.history[-1].result[-1].extracted_content = self.state.extracted_content
 
-            return self.history
+            return self.state.history
 
         finally:
             self.telemetry.capture(
                 AgentEndTelemetryEvent(
-                    agent_id=self.agent_id,
-                    success=self.history.is_done(),
-                    steps=self.n_steps,
-                    max_steps_reached=self.n_steps >= max_steps,
-                    errors=self.history.errors(),
+                    agent_id=self.state.agent_id,
+                    is_done=self.state.history.is_done(),
+                    success=self.state.history.is_successful(),
+                    steps=self.state.n_steps,
+                    max_steps_reached=self.state.n_steps >= max_steps,
+                    errors=self.state.history.errors(),
+                    total_input_tokens=self.state.history.total_input_tokens(),
+                    total_duration_seconds=self.state.history.total_duration_seconds(),
                 )
             )
 
@@ -464,122 +466,9 @@ class CustomAgent(Agent):
             if not self.injected_browser and self.browser:
                 await self.browser.close()
 
-            if self.generate_gif:
+            if self.settings.generate_gif:
                 output_path: str = 'agent_history.gif'
-                if isinstance(self.generate_gif, str):
-                    output_path = self.generate_gif
+                if isinstance(self.settings.generate_gif, str):
+                    output_path = self.settings.generate_gif
 
-                self.create_history_gif(output_path=output_path)
-
-    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:
-                    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')
+                create_history_gif(task=self.task, history=self.state.history, output_path=output_path)

src/agent/custom_message_manager.py

@@ -8,14 +8,17 @@ from browser_use.agent.message_manager.views import MessageHistory
 from browser_use.agent.prompts import SystemPrompt, AgentMessagePrompt
 from browser_use.agent.views import ActionResult, AgentStepInfo, ActionModel
 from browser_use.browser.views import BrowserState
+from browser_use.agent.message_manager.service import MessageManagerSettings
+from browser_use.agent.views import ActionResult, AgentOutput, AgentStepInfo, MessageManagerState
 from langchain_core.language_models import BaseChatModel
 from langchain_anthropic import ChatAnthropic
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import (
-	AIMessage,
-	BaseMessage,
-	HumanMessage,
-    ToolMessage
+    AIMessage,
+    BaseMessage,
+    HumanMessage,
+    ToolMessage,
+    SystemMessage
 )
 from langchain_openai import ChatOpenAI
 from ..utils.llm import DeepSeekR1ChatOpenAI
@@ -24,55 +27,55 @@ from .custom_prompts import CustomAgentMessagePrompt
 logger = logging.getLogger(__name__)
 
 
+class CustomMessageManagerSettings(MessageManagerSettings):
+    agent_prompt_class: Type[AgentMessagePrompt] = AgentMessagePrompt
+
+
 class CustomMessageManager(MessageManager):
     def __init__(
             self,
-            llm: BaseChatModel,
             task: str,
-            action_descriptions: str,
-            system_prompt_class: Type[SystemPrompt],
-            agent_prompt_class: Type[AgentMessagePrompt],
-            max_input_tokens: int = 128000,
-            estimated_characters_per_token: int = 3,
-            image_tokens: int = 800,
-            include_attributes: list[str] = [],
-            max_error_length: int = 400,
-            max_actions_per_step: int = 10,
-            message_context: Optional[str] = None,
-            sensitive_data: Optional[Dict[str, str]] = None,
+            system_message: SystemMessage,
+            settings: MessageManagerSettings = MessageManagerSettings(),
+            state: MessageManagerState = MessageManagerState(),
     ):
         super().__init__(
-            llm=llm,
             task=task,
-            action_descriptions=action_descriptions,
-            system_prompt_class=system_prompt_class,
-            max_input_tokens=max_input_tokens,
-            estimated_characters_per_token=estimated_characters_per_token,
-            image_tokens=image_tokens,
-            include_attributes=include_attributes,
-            max_error_length=max_error_length,
-            max_actions_per_step=max_actions_per_step,
-            message_context=message_context,
-            sensitive_data=sensitive_data
+            system_message=system_message,
+            settings=settings,
+            state=state
         )
-        self.agent_prompt_class = agent_prompt_class
-        # Custom: Move Task info to state_message
-        self.history = MessageHistory()
+
+    def _init_messages(self) -> None:
+        """Initialize the message history with system message, context, task, and other initial messages"""
         self._add_message_with_tokens(self.system_prompt)
-        
-        if self.message_context:
-            context_message = HumanMessage(content=self.message_context)
+        self.context_content = ""
+
+        if self.settings.message_context:
+            self.context_content += 'Context for the task' + self.settings.message_context
+
+        if self.settings.sensitive_data:
+            info = f'Here are placeholders for sensitive data: {list(self.settings.sensitive_data.keys())}'
+            info += 'To use them, write <secret>the placeholder name</secret>'
+            self.context_content += info
+
+        if self.settings.available_file_paths:
+            filepaths_msg = f'Here are file paths you can use: {self.settings.available_file_paths}'
+            self.context_content += filepaths_msg
+
+        if self.context_content:
+            context_message = HumanMessage(content=self.context_content)
             self._add_message_with_tokens(context_message)
 
     def cut_messages(self):
         """Get current message list, potentially trimmed to max tokens"""
-        diff = self.history.total_tokens - self.max_input_tokens
-        min_message_len = 2 if self.message_context is not None else 1
-        
-        while diff > 0 and len(self.history.messages) > min_message_len:
-            self.history.remove_message(min_message_len)  # always remove the oldest message
-            diff = self.history.total_tokens - self.max_input_tokens
-        
+        diff = self.state.history.current_tokens - self.settings.max_input_tokens
+        min_message_len = 2 if self.context_content is not None else 1
+
+        while diff > 0 and len(self.state.history.messages) > min_message_len:
+            self.state.history.remove_message(min_message_len)  # always remove the oldest message
+            diff = self.state.history.current_tokens - self.settings.max_input_tokens
+
     def add_state_message(
             self,
             state: BrowserState,
@@ -83,38 +86,23 @@ class CustomMessageManager(MessageManager):
     ) -> None:
         """Add browser state as human message"""
         # otherwise add state message and result to next message (which will not stay in memory)
-        state_message = self.agent_prompt_class(
+        state_message = self.settings.agent_prompt_class(
             state,
             actions,
             result,
-            include_attributes=self.include_attributes,
-            max_error_length=self.max_error_length,
+            include_attributes=self.settings.include_attributes,
             step_info=step_info,
         ).get_user_message(use_vision)
         self._add_message_with_tokens(state_message)
-    
-    def _count_text_tokens(self, text: str) -> int:
-        if isinstance(self.llm, (ChatOpenAI, ChatAnthropic, DeepSeekR1ChatOpenAI)):
-            try:
-                tokens = self.llm.get_num_tokens(text)
-            except Exception:
-                tokens = (
-					len(text) // self.estimated_characters_per_token
-				)  # Rough estimate if no tokenizer available
-        else:
-            tokens = (
-				len(text) // self.estimated_characters_per_token
-			)  # Rough estimate if no tokenizer available
-        return tokens
 
     def _remove_state_message_by_index(self, remove_ind=-1) -> None:
         """Remove last state message from history"""
-        i = len(self.history.messages) - 1
+        i = len(self.state.history.messages) - 1
         remove_cnt = 0
         while i >= 0:
-            if isinstance(self.history.messages[i].message, HumanMessage): 
+            if isinstance(self.state.history.messages[i].message, HumanMessage):
                 remove_cnt += 1
             if remove_cnt == abs(remove_ind):
-                self.history.remove_message(i)
+                self.state.history.messages.pop(i)
                 break
             i -= 1

src/agent/custom_prompts.py

@@ -20,9 +20,7 @@ class CustomSystemPrompt(SystemPrompt):
    {
      "current_state": {
        "evaluation_previous_goal": "Success|Failed|Unknown - Analyze the current elements and the image to check if the previous goals/actions are successful like intended by the task. Mention if something unexpected happened. Shortly state why/why not.",
-       "important_contents": "Output important contents closely related to user\'s instruction on the current page. If there is, please output the contents. If not, please output ''.",
-       "task_progress": "Task Progress is a general summary of the current contents that have been completed. Just summarize the contents that have been actually completed based on the content at current step and the history operations. Please list each completed item individually, such as: 1. Input username. 2. Input Password. 3. Click confirm button. Please return string type not a list.",
-       "future_plans": "Based on the user's request and the current state, outline the remaining steps needed to complete the task. This should be a concise list of sub-goals yet to be performed, such as: 1. Select a date. 2. Choose a specific time slot. 3. Confirm booking. Please return string type not a list.",
+       "important_contents": "Output important contents closely related to user's instruction on the current page. If there is, please output the contents. If not, please output ''.",
        "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 your output of evaluation_previous_goal is 'Failed', please reflect and output your reflection here.",
        "next_goal": "Please generate a brief natural language description for the goal of your next actions based on your thought."
      },
@@ -167,7 +165,7 @@ class CustomAgentMessagePrompt(AgentMessagePrompt):
 
         if self.actions and self.result:
             state_description += "\n **Previous Actions** \n"
-            state_description += f'Previous step: {self.step_info.step_number-1}/{self.step_info.max_steps} \n'
+            state_description += f'Previous step: {self.step_info.step_number - 1}/{self.step_info.max_steps} \n'
             for i, result in enumerate(self.result):
                 action = self.actions[i]
                 state_description += f"Previous action {i + 1}/{len(self.result)}: {action.model_dump_json(exclude_unset=True)}\n"

src/agent/custom_views.py

@@ -1,7 +1,8 @@
 from dataclasses import dataclass
-from typing import Type
+from typing import Any, Dict, List, Literal, Optional, Type
+import uuid
 
-from browser_use.agent.views import AgentOutput
+from browser_use.agent.views import AgentOutput, AgentState, ActionResult, AgentHistoryList, MessageManagerState
 from browser_use.controller.registry.views import ActionModel
 from pydantic import BaseModel, ConfigDict, Field, create_model
 
@@ -13,8 +14,6 @@ class CustomAgentStepInfo:
     task: str
     add_infos: str
     memory: str
-    task_progress: str
-    future_plans: str
 
 
 class CustomAgentBrain(BaseModel):
@@ -22,8 +21,6 @@ class CustomAgentBrain(BaseModel):
 
     evaluation_previous_goal: str
     important_contents: str
-    task_progress: str
-    future_plans: str
     thought: str
     next_goal: str
 
@@ -38,7 +35,7 @@ class CustomAgentOutput(AgentOutput):
 
     @staticmethod
     def type_with_custom_actions(
-        custom_actions: Type[ActionModel],
+            custom_actions: Type[ActionModel],
     ) -> Type["CustomAgentOutput"]:
         """Extend actions with custom actions"""
         model_ = create_model(
@@ -52,3 +49,19 @@ class CustomAgentOutput(AgentOutput):
         )
         model_.__doc__ = 'AgentOutput model with custom actions'
         return model_
+
+
+class CustomAgentState(BaseModel):
+    agent_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    n_steps: int = 1
+    consecutive_failures: int = 0
+    last_result: Optional[List['ActionResult']] = None
+    history: AgentHistoryList = Field(default_factory=lambda: AgentHistoryList(history=[]))
+    last_plan: Optional[str] = None
+    paused: bool = False
+    stopped: bool = False
+
+    message_manager_state: MessageManagerState = Field(default_factory=MessageManagerState)
+
+    last_action: Optional[List['ActionModel']] = None
+    extracted_content: str = ''

src/browser/custom_browser.py

@@ -18,10 +18,11 @@ from .custom_context import CustomBrowserContext
 
 logger = logging.getLogger(__name__)
 
+
 class CustomBrowser(Browser):
 
     async def new_context(
-        self,
-        config: BrowserContextConfig = BrowserContextConfig()
+            self,
+            config: BrowserContextConfig = BrowserContextConfig()
     ) -> CustomBrowserContext:
         return CustomBrowserContext(config=config, browser=self)

src/browser/custom_context.py

@@ -12,8 +12,8 @@ logger = logging.getLogger(__name__)
 
 class CustomBrowserContext(BrowserContext):
     def __init__(
-        self,
-        browser: "Browser",
-        config: BrowserContextConfig = BrowserContextConfig()
+            self,
+            browser: "Browser",
+            config: BrowserContextConfig = BrowserContextConfig()
     ):
-        super(CustomBrowserContext, self).__init__(browser=browser, config=config)
+        super(CustomBrowserContext, self).__init__(browser=browser, config=config)

src/utils/deep_research.py

@@ -310,11 +310,12 @@ Provide your output as a JSON formatted list. Each item in the list must adhere
             await browser_context.close()
         logger.info("Browser closed.")
 
+
 async def generate_final_report(task, history_infos, save_dir, llm, error_msg=None):
     """Generate report from collected information with error handling"""
     try:
         logger.info("\nAttempting to generate final report from collected data...")
-        
+
         writer_system_prompt = """
         You are a **Deep Researcher** and a professional report writer tasked with creating polished, high-quality reports that fully meet the user's needs, based on the user's instructions and the relevant information provided. You will write the report using Markdown format, ensuring it is both informative and visually appealing.
 
@@ -366,9 +367,9 @@ async def generate_final_report(task, history_infos, save_dir, llm, error_msg=No
         # Add error notification to the report
         if error_msg:
             report_content = f"## ⚠️ Research Incomplete - Partial Results\n" \
-                            f"**The research process was interrupted by an error:** {error_msg}\n\n" \
-                            f"{report_content}"
-            
+                             f"**The research process was interrupted by an error:** {error_msg}\n\n" \
+                             f"{report_content}"
+
         report_file_path = os.path.join(save_dir, "final_report.md")
         with open(report_file_path, "w", encoding="utf-8") as f:
             f.write(report_content)

src/utils/llm.py

@@ -40,22 +40,23 @@ from typing import (
     cast,
 )
 
+
 class DeepSeekR1ChatOpenAI(ChatOpenAI):
-    
+
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         self.client = OpenAI(
             base_url=kwargs.get("base_url"),
             api_key=kwargs.get("api_key")
-        ) 
-        
+        )
+
     async def ainvoke(
-        self,
-        input: LanguageModelInput,
-        config: Optional[RunnableConfig] = None,
-        *,
-        stop: Optional[list[str]] = None,
-        **kwargs: Any,
+            self,
+            input: LanguageModelInput,
+            config: Optional[RunnableConfig] = None,
+            *,
+            stop: Optional[list[str]] = None,
+            **kwargs: Any,
     ) -> AIMessage:
         message_history = []
         for input_ in input:
@@ -65,7 +66,7 @@ class DeepSeekR1ChatOpenAI(ChatOpenAI):
                 message_history.append({"role": "assistant", "content": input_.content})
             else:
                 message_history.append({"role": "user", "content": input_.content})
-        
+
         response = self.client.chat.completions.create(
             model=self.model_name,
             messages=message_history
@@ -74,14 +75,14 @@ class DeepSeekR1ChatOpenAI(ChatOpenAI):
         reasoning_content = response.choices[0].message.reasoning_content
         content = response.choices[0].message.content
         return AIMessage(content=content, reasoning_content=reasoning_content)
-    
+
     def invoke(
-        self,
-        input: LanguageModelInput,
-        config: Optional[RunnableConfig] = None,
-        *,
-        stop: Optional[list[str]] = None,
-        **kwargs: Any,
+            self,
+            input: LanguageModelInput,
+            config: Optional[RunnableConfig] = None,
+            *,
+            stop: Optional[list[str]] = None,
+            **kwargs: Any,
     ) -> AIMessage:
         message_history = []
         for input_ in input:
@@ -91,7 +92,7 @@ class DeepSeekR1ChatOpenAI(ChatOpenAI):
                 message_history.append({"role": "assistant", "content": input_.content})
             else:
                 message_history.append({"role": "user", "content": input_.content})
-        
+
         response = self.client.chat.completions.create(
             model=self.model_name,
             messages=message_history
@@ -100,16 +101,17 @@ class DeepSeekR1ChatOpenAI(ChatOpenAI):
         reasoning_content = response.choices[0].message.reasoning_content
         content = response.choices[0].message.content
         return AIMessage(content=content, reasoning_content=reasoning_content)
-    
+
+
 class DeepSeekR1ChatOllama(ChatOllama):
-        
+
     async def ainvoke(
-        self,
-        input: LanguageModelInput,
-        config: Optional[RunnableConfig] = None,
-        *,
-        stop: Optional[list[str]] = None,
-        **kwargs: Any,
+            self,
+            input: LanguageModelInput,
+            config: Optional[RunnableConfig] = None,
+            *,
+            stop: Optional[list[str]] = None,
+            **kwargs: Any,
     ) -> AIMessage:
         org_ai_message = await super().ainvoke(input=input)
         org_content = org_ai_message.content
@@ -118,14 +120,14 @@ class DeepSeekR1ChatOllama(ChatOllama):
         if "**JSON Response:**" in content:
             content = content.split("**JSON Response:**")[-1]
         return AIMessage(content=content, reasoning_content=reasoning_content)
-    
+
     def invoke(
-        self,
-        input: LanguageModelInput,
-        config: Optional[RunnableConfig] = None,
-        *,
-        stop: Optional[list[str]] = None,
-        **kwargs: Any,
+            self,
+            input: LanguageModelInput,
+            config: Optional[RunnableConfig] = None,
+            *,
+            stop: Optional[list[str]] = None,
+            **kwargs: Any,
     ) -> AIMessage:
         org_ai_message = super().invoke(input=input)
         org_content = org_ai_message.content
@@ -133,4 +135,4 @@ class DeepSeekR1ChatOllama(ChatOllama):
         content = org_content.split("</think>")[1]
         if "**JSON Response:**" in content:
             content = content.split("**JSON Response:**")[-1]
-        return AIMessage(content=content, reasoning_content=reasoning_content)
+        return AIMessage(content=content, reasoning_content=reasoning_content)

src/utils/utils.py

@@ -24,6 +24,7 @@ PROVIDER_DISPLAY_NAMES = {
     "moonshot": "MoonShot"
 }
 
+
 def get_llm_model(provider: str, **kwargs):
     """
     获取LLM 模型
@@ -161,19 +162,23 @@ def get_llm_model(provider: str, **kwargs):
     else:
         raise ValueError(f"Unsupported provider: {provider}")
 
+
 # Predefined model names for common providers
 model_names = {
     "anthropic": ["claude-3-5-sonnet-20241022", "claude-3-5-sonnet-20240620", "claude-3-opus-20240229"],
     "openai": ["gpt-4o", "gpt-4", "gpt-3.5-turbo", "o3-mini"],
     "deepseek": ["deepseek-chat", "deepseek-reasoner"],
-    "google": ["gemini-2.0-flash", "gemini-2.0-flash-thinking-exp", "gemini-1.5-flash-latest", "gemini-1.5-flash-8b-latest", "gemini-2.0-flash-thinking-exp-01-21", "gemini-2.0-pro-exp-02-05"],
-    "ollama": ["qwen2.5:7b", "qwen2.5:14b", "qwen2.5:32b", "qwen2.5-coder:14b", "qwen2.5-coder:32b", "llama2:7b", "deepseek-r1:14b", "deepseek-r1:32b"],
+    "google": ["gemini-2.0-flash", "gemini-2.0-flash-thinking-exp", "gemini-1.5-flash-latest",
+               "gemini-1.5-flash-8b-latest", "gemini-2.0-flash-thinking-exp-01-21", "gemini-2.0-pro-exp-02-05"],
+    "ollama": ["qwen2.5:7b", "qwen2.5:14b", "qwen2.5:32b", "qwen2.5-coder:14b", "qwen2.5-coder:32b", "llama2:7b",
+               "deepseek-r1:14b", "deepseek-r1:32b"],
     "azure_openai": ["gpt-4o", "gpt-4", "gpt-3.5-turbo"],
-    "mistral": ["pixtral-large-latest", "mistral-large-latest", "mistral-small-latest", "ministral-8b-latest"],
+    "mistral": ["mixtral-large-latest", "mistral-large-latest", "mistral-small-latest", "ministral-8b-latest"],
     "alibaba": ["qwen-plus", "qwen-max", "qwen-turbo", "qwen-long"],
     "moonshot": ["moonshot-v1-32k-vision-preview", "moonshot-v1-8k-vision-preview"],
 }
 
+
 # Callback to update the model name dropdown based on the selected provider
 def update_model_dropdown(llm_provider, api_key=None, base_url=None):
     """
@@ -191,6 +196,7 @@ def update_model_dropdown(llm_provider, api_key=None, base_url=None):
     else:
         return gr.Dropdown(choices=[], value="", interactive=True, allow_custom_value=True)
 
+
 def handle_api_key_error(provider: str, env_var: str):
     """
     Handles the missing API key error by raising a gr.Error with a clear message.
@@ -201,6 +207,7 @@ def handle_api_key_error(provider: str, env_var: str):
         f"`{env_var}` environment variable or provide it in the UI."
     )
 
+
 def encode_image(img_path):
     if not img_path:
         return None
@@ -212,7 +219,7 @@ def encode_image(img_path):
 def get_latest_files(directory: str, file_types: list = ['.webm', '.zip']) -> Dict[str, Optional[str]]:
     """Get the latest recording and trace files"""
     latest_files: Dict[str, Optional[str]] = {ext: None for ext in file_types}
-    
+
     if not os.path.exists(directory):
         os.makedirs(directory, exist_ok=True)
         return latest_files
@@ -227,8 +234,10 @@ def get_latest_files(directory: str, file_types: list = ['.webm', '.zip']) -> Di
                     latest_files[file_type] = str(latest)
         except Exception as e:
             print(f"Error getting latest {file_type} file: {e}")
-            
+
     return latest_files
+
+
 async def capture_screenshot(browser_context):
     """Capture and encode a screenshot"""
     # Extract the Playwright browser instance

tests/test_browser_use.py

@@ -37,7 +37,7 @@ async def test_browser_use_org():
     #     model_name="deepseek-chat",
     #     temperature=0.8
     # )
-    
+
     llm = utils.get_llm_model(
         provider="ollama", model_name="deepseek-r1:14b", temperature=0.5
     )
@@ -51,7 +51,7 @@ async def test_browser_use_org():
             chrome_path = None
     else:
         chrome_path = None
-        
+
     tool_calling_method = "json_schema"  # setting to json_schema when using ollma
 
     browser = Browser(
@@ -63,14 +63,14 @@ async def test_browser_use_org():
         )
     )
     async with await browser.new_context(
-        config=BrowserContextConfig(
-            trace_path="./tmp/traces",
-            save_recording_path="./tmp/record_videos",
-            no_viewport=False,
-            browser_window_size=BrowserContextWindowSize(
-                width=window_w, height=window_h
-            ),
-        )
+            config=BrowserContextConfig(
+                trace_path="./tmp/traces",
+                save_recording_path="./tmp/record_videos",
+                no_viewport=False,
+                browser_window_size=BrowserContextWindowSize(
+                    width=window_w, height=window_h
+                ),
+            )
     ) as browser_context:
         agent = Agent(
             task="go to google.com and type 'OpenAI' click search and give me the first url",
@@ -108,8 +108,8 @@ async def test_browser_use_custom():
     from src.browser.custom_context import BrowserContextConfig
     from src.controller.custom_controller import CustomController
 
-    window_w, window_h = 1920, 1080
-    
+    window_w, window_h = 1280, 1100
+
     # llm = utils.get_llm_model(
     #     provider="openai",
     #     model_name="gpt-4o",
@@ -138,7 +138,7 @@ async def test_browser_use_custom():
     #     model_name="deepseek-reasoner",
     #     temperature=0.8
     # )
-    
+
     # llm = utils.get_llm_model(
     #     provider="deepseek",
     #     model_name="deepseek-chat",
@@ -148,7 +148,7 @@ async def test_browser_use_custom():
     # llm = utils.get_llm_model(
     #     provider="ollama", model_name="qwen2.5:7b", temperature=0.5
     # )
-    
+
     # llm = utils.get_llm_model(
     #     provider="ollama", model_name="deepseek-r1:14b", temperature=0.5
     # )
@@ -157,7 +157,7 @@ async def test_browser_use_custom():
     use_own_browser = True
     disable_security = True
     use_vision = False  # Set to False when using DeepSeek
-    
+
     max_actions_per_step = 1
     playwright = None
     browser = None
@@ -209,16 +209,6 @@ async def test_browser_use_custom():
         print("Final Result:")
         pprint(history.final_result(), indent=4)
 
-        print("\nErrors:")
-        pprint(history.errors(), indent=4)
-
-        # e.g. xPaths the model clicked on
-        print("\nModel Outputs:")
-        pprint(history.model_actions(), indent=4)
-
-        print("\nThoughts:")
-        pprint(history.model_thoughts(), indent=4)
-        # close browser
     except Exception:
         import traceback
 
@@ -233,7 +223,8 @@ async def test_browser_use_custom():
             await playwright.stop()
         if browser:
             await browser.close()
-            
+
+
 async def test_browser_use_parallel():
     from browser_use.browser.context import BrowserContextWindowSize
     from browser_use.browser.browser import BrowserConfig
@@ -246,7 +237,7 @@ async def test_browser_use_parallel():
     from src.controller.custom_controller import CustomController
 
     window_w, window_h = 1920, 1080
-    
+
     # llm = utils.get_llm_model(
     #     provider="openai",
     #     model_name="gpt-4o",
@@ -275,7 +266,7 @@ async def test_browser_use_parallel():
     #     model_name="deepseek-reasoner",
     #     temperature=0.8
     # )
-    
+
     # llm = utils.get_llm_model(
     #     provider="deepseek",
     #     model_name="deepseek-chat",
@@ -285,7 +276,7 @@ async def test_browser_use_parallel():
     # llm = utils.get_llm_model(
     #     provider="ollama", model_name="qwen2.5:7b", temperature=0.5
     # )
-    
+
     # llm = utils.get_llm_model(
     #     provider="ollama", model_name="deepseek-r1:14b", temperature=0.5
     # )
@@ -294,12 +285,12 @@ async def test_browser_use_parallel():
     use_own_browser = True
     disable_security = True
     use_vision = True  # Set to False when using DeepSeek
-    
+
     max_actions_per_step = 1
     playwright = None
     browser = None
     browser_context = None
-    
+
     browser = Browser(
         config=BrowserConfig(
             disable_security=True,
@@ -310,7 +301,7 @@ async def test_browser_use_parallel():
 
     try:
         agents = [
-		Agent(task=task, llm=llm, browser=browser)
+            Agent(task=task, llm=llm, browser=browser)
             for task in [
                 'Search Google for weather in Tokyo',
                 'Check Reddit front page title',
@@ -355,6 +346,7 @@ async def test_browser_use_parallel():
         if browser:
             await browser.close()
 
+
 if __name__ == "__main__":
     # asyncio.run(test_browser_use_org())
     # asyncio.run(test_browser_use_parallel())

webui.py

@@ -32,10 +32,10 @@ from src.agent.custom_prompts import CustomSystemPrompt, CustomAgentMessagePromp
 from src.browser.custom_context import BrowserContextConfig, CustomBrowserContext
 from src.controller.custom_controller import CustomController
 from gradio.themes import Citrus, Default, Glass, Monochrome, Ocean, Origin, Soft, Base
-from src.utils.default_config_settings import default_config, load_config_from_file, save_config_to_file, save_current_config, update_ui_from_config
+from src.utils.default_config_settings import default_config, load_config_from_file, save_config_to_file, \
+    save_current_config, update_ui_from_config
 from src.utils.utils import update_model_dropdown, get_latest_files, capture_screenshot
 
-
 # Global variables for persistence
 _global_browser = None
 _global_browser_context = None
@@ -44,6 +44,7 @@ _global_agent = None
 # Create the global agent state instance
 _global_agent_state = AgentState()
 
+
 def resolve_sensitive_env_variables(text):
     """
     Replace environment variable placeholders ($SENSITIVE_*) with their values.
@@ -51,12 +52,12 @@ def resolve_sensitive_env_variables(text):
     """
     if not text:
         return text
-        
+
     import re
-    
+
     # Find all $SENSITIVE_* patterns
     env_vars = re.findall(r'\$SENSITIVE_[A-Za-z0-9_]*', text)
-    
+
     result = text
     for var in env_vars:
         # Remove the $ prefix to get the actual environment variable name
@@ -65,9 +66,10 @@ def resolve_sensitive_env_variables(text):
         if env_value is not None:
             # Replace $SENSITIVE_VAR_NAME with its value
             result = result.replace(var, env_value)
-        
+
     return result
 
+
 async def stop_agent():
     """Request the agent to stop and update UI with enhanced feedback"""
     global _global_agent_state, _global_browser_context, _global_browser, _global_agent
@@ -82,9 +84,9 @@ async def stop_agent():
 
         # Return UI updates
         return (
-            message,                                        # errors_output
+            message,  # errors_output
             gr.update(value="Stopping...", interactive=False),  # stop_button
-            gr.update(interactive=False),                      # run_button
+            gr.update(interactive=False),  # run_button
         )
     except Exception as e:
         error_msg = f"Error during stop: {str(e)}"
@@ -94,7 +96,8 @@ async def stop_agent():
             gr.update(value="Stop", interactive=True),
             gr.update(interactive=True)
         )
-        
+
+
 async def stop_research_agent():
     """Request the agent to stop and update UI with enhanced feedback"""
     global _global_agent_state, _global_browser_context, _global_browser
@@ -108,9 +111,9 @@ async def stop_research_agent():
         logger.info(f"🛑 {message}")
 
         # Return UI updates
-        return (                                   # errors_output
+        return (  # errors_output
             gr.update(value="Stopping...", interactive=False),  # stop_button
-            gr.update(interactive=False),                      # run_button
+            gr.update(interactive=False),  # run_button
         )
     except Exception as e:
         error_msg = f"Error during stop: {str(e)}"
@@ -120,6 +123,7 @@ async def stop_research_agent():
             gr.update(interactive=True)
         )
 
+
 async def run_browser_agent(
         agent_type,
         llm_provider,
@@ -238,7 +242,7 @@ async def run_browser_agent(
             trace_file,
             history_file,
             gr.update(value="Stop", interactive=True),  # Re-enable stop button
-            gr.update(interactive=True)    # Re-enable run button
+            gr.update(interactive=True)  # Re-enable run button
         )
 
     except gr.Error:
@@ -249,15 +253,15 @@ async def run_browser_agent(
         traceback.print_exc()
         errors = str(e) + "\n" + traceback.format_exc()
         return (
-            '',                                         # final_result
-            errors,                                     # errors
-            '',                                         # model_actions
-            '',                                         # model_thoughts
-            None,                                       # latest_video
-            None,                                       # history_file
-            None,                                       # trace_file
+            '',  # final_result
+            errors,  # errors
+            '',  # model_actions
+            '',  # model_thoughts
+            None,  # latest_video
+            None,  # history_file
+            None,  # trace_file
             gr.update(value="Stop", interactive=True),  # Re-enable stop button
-            gr.update(interactive=True)    # Re-enable run button
+            gr.update(interactive=True)  # Re-enable run button
         )
 
 
@@ -281,7 +285,7 @@ async def run_org_agent(
 ):
     try:
         global _global_browser, _global_browser_context, _global_agent_state, _global_agent
-        
+
         # Clear any previous stop request
         _global_agent_state.clear_stop()
 
@@ -298,9 +302,8 @@ async def run_org_agent(
                 extra_chromium_args += [f"--user-data-dir={chrome_user_data}"]
         else:
             chrome_path = None
-            
-        if _global_browser is None:
 
+        if _global_browser is None:
             _global_browser = Browser(
                 config=BrowserConfig(
                     headless=headless,
@@ -363,6 +366,7 @@ async def run_org_agent(
                 await _global_browser.close()
                 _global_browser = None
 
+
 async def run_custom_agent(
         llm,
         use_own_browser,
@@ -405,7 +409,7 @@ async def run_custom_agent(
         controller = CustomController()
 
         # Initialize global browser if needed
-        #if chrome_cdp not empty string nor None
+        # if chrome_cdp not empty string nor None
         if (_global_browser is None) or (cdp_url and cdp_url != "" and cdp_url != None):
             _global_browser = CustomBrowser(
                 config=BrowserConfig(
@@ -417,7 +421,7 @@ async def run_custom_agent(
                 )
             )
 
-        if _global_browser_context is None  or (chrome_cdp and cdp_url != "" and cdp_url != None):
+        if _global_browser_context is None or (chrome_cdp and cdp_url != "" and cdp_url != None):
             _global_browser_context = await _global_browser.new_context(
                 config=BrowserContextConfig(
                     trace_path=save_trace_path if save_trace_path else None,
@@ -429,7 +433,6 @@ async def run_custom_agent(
                 )
             )
 
-
         # Create and run agent
         if _global_agent is None:
             _global_agent = CustomAgent(
@@ -455,7 +458,7 @@ async def run_custom_agent(
         model_actions = history.model_actions()
         model_thoughts = history.model_thoughts()
 
-        trace_file = get_latest_files(save_trace_path)        
+        trace_file = get_latest_files(save_trace_path)
 
         return final_result, errors, model_actions, model_thoughts, trace_file.get('.zip'), history_file
     except Exception as e:
@@ -475,31 +478,32 @@ async def run_custom_agent(
                 await _global_browser.close()
                 _global_browser = None
 
+
 async def run_with_stream(
-    agent_type,
-    llm_provider,
-    llm_model_name,
-    llm_num_ctx,
-    llm_temperature,
-    llm_base_url,
-    llm_api_key,
-    use_own_browser,
-    keep_browser_open,
-    headless,
-    disable_security,
-    window_w,
-    window_h,
-    save_recording_path,
-    save_agent_history_path,
-    save_trace_path,
-    enable_recording,
-    task,
-    add_infos,
-    max_steps,
-    use_vision,
-    max_actions_per_step,
-    tool_calling_method,
-    chrome_cdp
+        agent_type,
+        llm_provider,
+        llm_model_name,
+        llm_num_ctx,
+        llm_temperature,
+        llm_base_url,
+        llm_api_key,
+        use_own_browser,
+        keep_browser_open,
+        headless,
+        disable_security,
+        window_w,
+        window_h,
+        save_recording_path,
+        save_agent_history_path,
+        save_trace_path,
+        enable_recording,
+        task,
+        add_infos,
+        max_steps,
+        use_vision,
+        max_actions_per_step,
+        tool_calling_method,
+        chrome_cdp
 ):
     global _global_agent_state
     stream_vw = 80
@@ -572,7 +576,6 @@ async def run_with_stream(
             final_result = errors = model_actions = model_thoughts = ""
             latest_videos = trace = history_file = None
 
-
             # Periodically update the stream while the agent task is running
             while not agent_task.done():
                 try:
@@ -651,9 +654,10 @@ async def run_with_stream(
                 None,
                 None,
                 gr.update(value="Stop", interactive=True),  # Re-enable stop button
-                gr.update(interactive=True)    # Re-enable run button
+                gr.update(interactive=True)  # Re-enable run button
             ]
 
+
 # Define the theme map globally
 theme_map = {
     "Default": Default(),
@@ -666,6 +670,7 @@ theme_map = {
     "Base": Base()
 }
 
+
 async def close_global_browser():
     global _global_browser, _global_browser_context
 
@@ -676,33 +681,36 @@ async def close_global_browser():
     if _global_browser:
         await _global_browser.close()
         _global_browser = None
-        
-async def run_deep_search(research_task, max_search_iteration_input, max_query_per_iter_input, llm_provider, llm_model_name, llm_num_ctx, llm_temperature, llm_base_url, llm_api_key, use_vision, use_own_browser, headless, chrome_cdp):
+
+
+async def run_deep_search(research_task, max_search_iteration_input, max_query_per_iter_input, llm_provider,
+                          llm_model_name, llm_num_ctx, llm_temperature, llm_base_url, llm_api_key, use_vision,
+                          use_own_browser, headless, chrome_cdp):
     from src.utils.deep_research import deep_research
     global _global_agent_state
 
     # Clear any previous stop request
     _global_agent_state.clear_stop()
-    
+
     llm = utils.get_llm_model(
-            provider=llm_provider,
-            model_name=llm_model_name,
-            num_ctx=llm_num_ctx,
-            temperature=llm_temperature,
-            base_url=llm_base_url,
-            api_key=llm_api_key,
-        )
+        provider=llm_provider,
+        model_name=llm_model_name,
+        num_ctx=llm_num_ctx,
+        temperature=llm_temperature,
+        base_url=llm_base_url,
+        api_key=llm_api_key,
+    )
     markdown_content, file_path = await deep_research(research_task, llm, _global_agent_state,
-                                                        max_search_iterations=max_search_iteration_input,
-                                                        max_query_num=max_query_per_iter_input,
-                                                        use_vision=use_vision,
-                                                        headless=headless,
-                                                        use_own_browser=use_own_browser,
-                                                        chrome_cdp=chrome_cdp
-                                                        )
-    
-    return markdown_content, file_path, gr.update(value="Stop", interactive=True),  gr.update(interactive=True) 
-    
+                                                      max_search_iterations=max_search_iteration_input,
+                                                      max_query_num=max_query_per_iter_input,
+                                                      use_vision=use_vision,
+                                                      headless=headless,
+                                                      use_own_browser=use_own_browser,
+                                                      chrome_cdp=chrome_cdp
+                                                      )
+
+    return markdown_content, file_path, gr.update(value="Stop", interactive=True), gr.update(interactive=True)
+
 
 def create_ui(config, theme_name="Ocean"):
     css = """
@@ -779,7 +787,7 @@ def create_ui(config, theme_name="Ocean"):
             with gr.TabItem("🔧 LLM Configuration", id=2):
                 with gr.Group():
                     llm_provider = gr.Dropdown(
-                        choices=[provider for provider,model in utils.model_names.items()],
+                        choices=[provider for provider, model in utils.model_names.items()],
                         label="LLM Provider",
                         value=config['llm_provider'],
                         info="Select your preferred language model provider"
@@ -790,11 +798,11 @@ def create_ui(config, theme_name="Ocean"):
                         value=config['llm_model_name'],
                         interactive=True,
                         allow_custom_value=True,  # Allow users to input custom model names
-                        info="Select a model from the dropdown or type a custom model name"
+                        info="Select a model in the dropdown options or directly type a custom model name"
                     )
                     llm_num_ctx = gr.Slider(
-                        minimum=2**8,
-                        maximum=2**16,
+                        minimum=2 ** 8,
+                        maximum=2 ** 16,
                         value=config['llm_num_ctx'],
                         step=1,
                         label="Max Context Length",
@@ -874,7 +882,6 @@ def create_ui(config, theme_name="Ocean"):
                             info="Browser window height",
                         )
 
-
                     save_recording_path = gr.Textbox(
                         label="Recording Path",
                         placeholder="e.g. ./tmp/record_videos",
@@ -933,28 +940,29 @@ def create_ui(config, theme_name="Ocean"):
                 with gr.Row():
                     run_button = gr.Button("▶️ Run Agent", variant="primary", scale=2)
                     stop_button = gr.Button("⏹️ Stop", variant="stop", scale=1)
-                    
+
                 with gr.Row():
                     browser_view = gr.HTML(
                         value="<h1 style='width:80vw; height:50vh'>Waiting for browser session...</h1>",
                         label="Live Browser View",
-                )
-            
+                    )
+
             with gr.TabItem("🧐 Deep Research", id=5):
-                research_task_input = gr.Textbox(label="Research Task", lines=5, value="Compose a report on the use of Reinforcement Learning for training Large Language Models, encompassing its origins, current advancements, and future prospects, substantiated with examples of relevant models and techniques. The report should reflect original insights and analysis, moving beyond mere summarization of existing literature.")
+                research_task_input = gr.Textbox(label="Research Task", lines=5,
+                                                 value="Compose a report on the use of Reinforcement Learning for training Large Language Models, encompassing its origins, current advancements, and future prospects, substantiated with examples of relevant models and techniques. The report should reflect original insights and analysis, moving beyond mere summarization of existing literature.")
                 with gr.Row():
-                    max_search_iteration_input = gr.Number(label="Max Search Iteration", value=3, precision=0) # precision=0 确保是整数
-                    max_query_per_iter_input = gr.Number(label="Max Query per Iteration", value=1, precision=0) # precision=0 确保是整数
+                    max_search_iteration_input = gr.Number(label="Max Search Iteration", value=3,
+                                                           precision=0)  # precision=0 确保是整数
+                    max_query_per_iter_input = gr.Number(label="Max Query per Iteration", value=1,
+                                                         precision=0)  # precision=0 确保是整数
                 with gr.Row():
                     research_button = gr.Button("▶️ Run Deep Research", variant="primary", scale=2)
-                    stop_research_button = gr.Button("⏹️ Stop", variant="stop", scale=1)
+                    stop_research_button = gr.Button("⏹ Stop", variant="stop", scale=1)
                 markdown_output_display = gr.Markdown(label="Research Report")
                 markdown_download = gr.File(label="Download Research Report")
 
-
             with gr.TabItem("📊 Results", id=6):
                 with gr.Group():
-
                     recording_display = gr.Video(label="Latest Recording")
 
                     gr.Markdown("### Results")
@@ -991,31 +999,35 @@ def create_ui(config, theme_name="Ocean"):
                 # Run button click handler
                 run_button.click(
                     fn=run_with_stream,
-                        inputs=[
-                            agent_type, llm_provider, llm_model_name, llm_num_ctx, llm_temperature, llm_base_url, llm_api_key,
-                            use_own_browser, keep_browser_open, headless, disable_security, window_w, window_h,
-                            save_recording_path, save_agent_history_path, save_trace_path,  # Include the new path
-                            enable_recording, task, add_infos, max_steps, use_vision, max_actions_per_step, tool_calling_method, chrome_cdp
-                        ],
+                    inputs=[
+                        agent_type, llm_provider, llm_model_name, llm_num_ctx, llm_temperature, llm_base_url,
+                        llm_api_key,
+                        use_own_browser, keep_browser_open, headless, disable_security, window_w, window_h,
+                        save_recording_path, save_agent_history_path, save_trace_path,  # Include the new path
+                        enable_recording, task, add_infos, max_steps, use_vision, max_actions_per_step,
+                        tool_calling_method, chrome_cdp
+                    ],
                     outputs=[
-                        browser_view,           # Browser view
-                        final_result_output,    # Final result
-                        errors_output,          # Errors
-                        model_actions_output,   # Model actions
+                        browser_view,  # Browser view
+                        final_result_output,  # Final result
+                        errors_output,  # Errors
+                        model_actions_output,  # Model actions
                         model_thoughts_output,  # Model thoughts
-                        recording_display,      # Latest recording
-                        trace_file,             # Trace file
-                        agent_history_file,     # Agent history file
-                        stop_button,            # Stop button
-                        run_button              # Run button
+                        recording_display,  # Latest recording
+                        trace_file,  # Trace file
+                        agent_history_file,  # Agent history file
+                        stop_button,  # Stop button
+                        run_button  # Run button
                     ],
                 )
-                
+
                 # Run Deep Research
                 research_button.click(
-                        fn=run_deep_search,
-                        inputs=[research_task_input, max_search_iteration_input, max_query_per_iter_input, llm_provider, llm_model_name, llm_num_ctx, llm_temperature, llm_base_url, llm_api_key, use_vision, use_own_browser, headless, chrome_cdp],
-                        outputs=[markdown_output_display, markdown_download, stop_research_button, research_button]
+                    fn=run_deep_search,
+                    inputs=[research_task_input, max_search_iteration_input, max_query_per_iter_input, llm_provider,
+                            llm_model_name, llm_num_ctx, llm_temperature, llm_base_url, llm_api_key, use_vision,
+                            use_own_browser, headless, chrome_cdp],
+                    outputs=[markdown_output_display, markdown_download, stop_research_button, research_button]
                 )
                 # Bind the stop button click event after errors_output is defined
                 stop_research_button.click(
@@ -1030,7 +1042,8 @@ def create_ui(config, theme_name="Ocean"):
                         return []
 
                     # Get all video files
-                    recordings = glob.glob(os.path.join(save_recording_path, "*.[mM][pP]4")) + glob.glob(os.path.join(save_recording_path, "*.[wW][eE][bB][mM]"))
+                    recordings = glob.glob(os.path.join(save_recording_path, "*.[mM][pP]4")) + glob.glob(
+                        os.path.join(save_recording_path, "*.[wW][eE][bB][mM]"))
 
                     # Sort recordings by creation time (oldest first)
                     recordings.sort(key=os.path.getctime)
@@ -1057,7 +1070,7 @@ def create_ui(config, theme_name="Ocean"):
                     inputs=save_recording_path,
                     outputs=recordings_gallery
                 )
-            
+
             with gr.TabItem("📁 Configuration", id=8):
                 with gr.Group():
                     config_file_input = gr.File(
@@ -1095,11 +1108,10 @@ def create_ui(config, theme_name="Ocean"):
                         use_own_browser, keep_browser_open, headless, disable_security,
                         enable_recording, window_w, window_h, save_recording_path, save_trace_path,
                         save_agent_history_path, task,
-                    ],  
+                    ],
                     outputs=[config_status]
                 )
 
-
         # Attach the callback to the LLM provider dropdown
         llm_provider.change(
             lambda provider, api_key, base_url: update_model_dropdown(provider, api_key, base_url),
@@ -1119,6 +1131,7 @@ def create_ui(config, theme_name="Ocean"):
 
     return demo
 
+
 def main():
     parser = argparse.ArgumentParser(description="Gradio UI for Browser Agent")
     parser.add_argument("--ip", type=str, default="127.0.0.1", help="IP address to bind to")
@@ -1132,5 +1145,6 @@ def main():
     demo = create_ui(config_dict, theme_name=args.theme)
     demo.launch(server_name=args.ip, server_port=args.port)
 
+
 if __name__ == '__main__':
     main()