src/agent/custom_prompts.py

import pdb
from typing import List, Optional

from browser_use.agent.prompts import SystemPrompt
from browser_use.agent.views import ActionResult
from browser_use.browser.views import BrowserState
from langchain_core.messages import HumanMessage, SystemMessage

from .custom_views import CustomAgentStepInfo

class CustomSystemPrompt(SystemPrompt):
    def important_rules(self) -> str:
        text = """
    1. RESPONSE FORMAT: You must ALWAYS respond with valid JSON in this exact format:
       {
         "current_state": {
           "prev_action_evaluation": "Success|Failed|Unknown - Analyze the current elements and the image to check if the previous goals/actions are successful like intended by the task. Ignore the action result. The website is the ground truth. Also mention if something unexpected happened like new suggestions in an input field. Shortly state why/why not. Note that the result you output must be consistent with the reasoning you output afterwards. If you consider it to be 'Failed,' you should reflect on this during your thought.",
           "important_contents": "Output important contents closely related to user\'s instruction or task on the current page. If there is, please output the contents. If not, please output empty string ''.",
           "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": [
           {
             "action_name": {
               // action-specific parameters
             }
           }
         ]
       }

    2. ACTIONS: Specify multiple actions to be executed in sequence if needed.

    3. ELEMENT INTERACTION:
       - Only use indexes that exist in the provided element list.
       - Ensure each action is valid based on the current page state.

    4. NAVIGATION & ERROR HANDLING:
       - Handle popups/cookies by accepting or closing them.
       - If no suitable elements exist, use alternative methods to proceed.

    5. TASK COMPLETION:
       - Ensure the task is fully completed before returning a 'done' action.
       - Always validate your output with the actual page content.

    6. VISUAL CONTEXT:
       - Use visual context when provided to verify layout and relationships.

    7. ACTION SEQUENCING:
       - Plan and execute actions efficiently, minimizing unnecessary steps.
    """
        text += f"   - Use a maximum of {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. Memory: Important contents recorded during historical operations.
    3. Current URL: The webpage currently being viewed.
    4. Interactive Elements: List in the format:
       index[:]<element_type>element_text</element_type>
    """

    def get_system_message(self) -> SystemMessage:
        time_str = self.current_date.strftime("%Y-%m-%d %H:%M")

        AGENT_PROMPT = f"""You are a precise browser automation agent. Your role is to:
    1. Analyze the provided webpage elements and structure.
    2. Plan a sequence of actions to accomplish the task.
    3. Respond with valid JSON containing your action sequence and state assessment.

    Current date and time: {time_str}

    {self.input_format()}

    {self.important_rules()}
    """
        return SystemMessage(content=AGENT_PROMPT)

class CustomAgentMessagePrompt:
    def __init__(
            self,
            state: BrowserState,
            result: Optional[List[ActionResult]] = None,
            include_attributes: list[str] = [],
            max_error_length: int = 400,
            step_info: Optional[CustomAgentStepInfo] = None,
    ):
        self.state = state
        self.result = result
        self.max_error_length = max_error_length
        self.include_attributes = include_attributes
        self.step_info = step_info

    def get_user_message(self) -> HumanMessage:
        step_info_description = f'Current step: {self.step_info.step_number + 1}/{self.step_info.max_steps}' if self.step_info else ''

        elements_text = self.state.element_tree.clickable_elements_to_string(include_attributes=self.include_attributes)
        if not elements_text:
            elements_text = 'empty page'

        state_description = f"""
{step_info_description}
1. Task: {self.step_info.task if self.step_info else 'No task provided.'}
2. Memory: {self.step_info.memory if self.step_info else 'No memory available.'}
3. Current URL: {self.state.url}
4. Interactive elements: {elements_text}
        """

        if self.result:
            for i, result in enumerate(self.result):
                if result.extracted_content:
                    state_description += f"\nResult of action {i + 1}/{len(self.result)}: {result.extracted_content}"
                if result.error:
                    error = result.error[-self.max_error_length:]
                    state_description += f"\nError of action {i + 1}/{len(self.result)}: ...{error}"

        return HumanMessage(content=state_description)