Merge pull request #413 from vvincent1234/feat/update_bu

Feat/update bu

Author: warmshao

.env.example

@@ -25,7 +25,7 @@ MOONSHOT_ENDPOINT=https://api.moonshot.cn/v1
 MOONSHOT_API_KEY=
 
 # Set to false to disable anonymized telemetry
-ANONYMIZED_TELEMETRY=true
+ANONYMIZED_TELEMETRY=false
 
 # LogLevel: Set to debug to enable verbose logging, set to result to get results only. Available: result | debug | info
 BROWSER_USE_LOGGING_LEVEL=info

requirements.txt

@@ -1,4 +1,4 @@
-browser-use==0.1.37
+browser-use==0.1.40
 pyperclip==1.9.0
 gradio==5.10.0
 json-repair

src/agent/custom_agent.py

@@ -8,14 +8,17 @@
 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 @@
 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 @@ def add_state_message(
     ) -> 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_message_manager.py

@@ -6,113 +6,20 @@
 from browser_use.browser.views import BrowserState
 from langchain_core.messages import HumanMessage, SystemMessage
 from datetime import datetime
+import importlib
 
 from .custom_views import CustomAgentStepInfo
 
 
 class CustomSystemPrompt(SystemPrompt):
-    def important_rules(self) -> str:
-        """
-        Returns the important rules for the agent.
-        """
-        text = r"""
-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 on the current page. If there is, please output the contents. If not, please output empty string ''.",
-       "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 actions 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.",
-       "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 prev_action_evaluation is 'Failed', please reflect and output your reflection here.",
-       "summary": "Please generate a brief natural language description for the operation in next actions based on your Thought."
-     },
-     "action": [
-       * actions in sequences, please refer to **Common action sequences**. Each output action MUST be formated as: \{action_name\: action_params\}* 
-     ]
-   }
-
-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: [
-       {"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.
-   - Note that you must verify if you've truly fulfilled the user's request by examining the actual page content, not just by looking at the actions you output but also whether the action is executed successfully. Pay particular attention when errors occur during action execution.
-
-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. 
-
-9. Extraction:
-    - If your task is to find information or do research - call extract_content on the specific pages to get and store the information.
-
-"""
-        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. Current URL: The webpage you're currently on
-5. Available Tabs: List of open browser tabs
-6. 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 inside [] are interactive
-- [] elements provide context but cannot be interacted with
-    """
+    def _load_prompt_template(self) -> None:
+        """Load the prompt template from the markdown file."""
+        try:
+            # This works both in development and when installed as a package
+            with importlib.resources.files('src.agent').joinpath('custom_system_prompt.md').open('r') as f:
+                self.prompt_template = f.read()
+        except Exception as e:
+            raise RuntimeError(f'Failed to load system prompt template: {e}')
 
 
 class CustomAgentMessagePrompt(AgentMessagePrompt):
@@ -122,13 +29,11 @@ def __init__(
             actions: Optional[List[ActionModel]] = None,
             result: Optional[List[ActionResult]] = None,
             include_attributes: list[str] = [],
-            max_error_length: int = 400,
             step_info: Optional[CustomAgentStepInfo] = None,
     ):
         super(CustomAgentMessagePrompt, self).__init__(state=state,
                                                        result=result,
                                                        include_attributes=include_attributes,
-                                                       max_error_length=max_error_length,
                                                        step_info=step_info
                                                        )
         self.actions = actions
@@ -179,19 +84,19 @@ def get_user_message(self, use_vision: bool = True) -> HumanMessage:
 
         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"
+                if result.error:
+                    # only use last 300 characters of error
+                    error = result.error.split('\n')[-1]
+                    state_description += (
+                        f"Error of previous action {i + 1}/{len(self.result)}: ...{error}\n"
+                    )
                 if result.include_in_memory:
                     if result.extracted_content:
                         state_description += f"Result of previous action {i + 1}/{len(self.result)}: {result.extracted_content}\n"
-                    if result.error:
-                        # only use last 300 characters of error
-                        error = result.error[-self.max_error_length:]
-                        state_description += (
-                            f"Error of previous action {i + 1}/{len(self.result)}: ...{error}\n"
-                        )
 
         if self.state.screenshot and use_vision == True:
             # Format message for vision model