| from __future__ import annotations | |
| import logging | |
| from typing import List, Optional, Type | |
| from browser_use.agent.message_manager.service import MessageManager | |
| from browser_use.agent.message_manager.views import MessageHistory | |
| from browser_use.agent.prompts import SystemPrompt | |
| from browser_use.agent.views import ActionResult, AgentStepInfo | |
| from browser_use.browser.views import BrowserState | |
| from langchain_core.language_models import BaseChatModel | |
| from langchain_core.messages import ( | |
| AIMessage, | |
| BaseMessage, | |
| HumanMessage, | |
| ) | |
| from langchain_openai import ChatOpenAI | |
| from ..utils.llm import DeepSeekR1ChatOpenAI | |
| from .custom_prompts import CustomAgentMessagePrompt | |
| logger = logging.getLogger(__name__) | |
| class CustomMassageManager(MessageManager): | |
| def __init__( | |
| self, | |
| llm: BaseChatModel, | |
| task: str, | |
| action_descriptions: str, | |
| system_prompt_class: Type[SystemPrompt], | |
| max_input_tokens: int = 128000, | |
| estimated_tokens_per_character: int = 3, | |
| image_tokens: int = 800, | |
| include_attributes: list[str] = [], | |
| max_error_length: int = 400, | |
| max_actions_per_step: int = 10, | |
| tool_call_in_content: bool = False, | |
| use_function_calling: bool = True | |
| ): | |
| super().__init__( | |
| llm=llm, | |
| task=task, | |
| action_descriptions=action_descriptions, | |
| system_prompt_class=system_prompt_class, | |
| max_input_tokens=max_input_tokens, | |
| estimated_tokens_per_character=estimated_tokens_per_character, | |
| image_tokens=image_tokens, | |
| include_attributes=include_attributes, | |
| max_error_length=max_error_length, | |
| max_actions_per_step=max_actions_per_step, | |
| tool_call_in_content=tool_call_in_content, | |
| ) | |
| self.use_function_calling = use_function_calling | |
| self.history = MessageHistory() | |
| self._add_message_with_tokens(self.system_prompt) | |
| if self.use_function_calling: | |
| tool_calls = [ | |
| { | |
| 'name': 'CustomAgentOutput', | |
| 'args': { | |
| 'current_state': { | |
| 'prev_action_evaluation': 'Unknown', | |
| 'important_contents': '', | |
| 'completed_contents': '', | |
| 'thought': 'Ready to proceed with the task.', | |
| 'summary': 'Initializing task.', | |
| }, | |
| 'action': [], | |
| }, | |
| 'id': '', | |
| 'type': 'tool_call', | |
| } | |
| ] | |
| example_tool_call = AIMessage( | |
| content='', | |
| tool_calls=tool_calls if not self.tool_call_in_content else [] | |
| ) | |
| self._add_message_with_tokens(example_tool_call) | |
| def cut_messages(self): | |
| diff = self.history.total_tokens - self.max_input_tokens | |
| while diff > 0 and len(self.history.messages) > 1: | |
| self.history.remove_message(1) | |
| diff = self.history.total_tokens - self.max_input_tokens | |
| def add_state_message( | |
| self, | |
| state: BrowserState, | |
| result: Optional[List[ActionResult]] = None, | |
| step_info: Optional[AgentStepInfo] = None, | |
| ) -> None: | |
| state_message = CustomAgentMessagePrompt( | |
| state, | |
| result, | |
| include_attributes=self.include_attributes, | |
| max_error_length=self.max_error_length, | |
| step_info=step_info, | |
| ).get_user_message() | |
| self._add_message_with_tokens(state_message) | |
| def _count_text_tokens(self, text: str) -> int: | |
| try: | |
| if isinstance(self.llm, (ChatOpenAI, DeepSeekR1ChatOpenAI)): | |
| return self.llm.get_num_tokens(text) | |
| except Exception: | |
| pass | |
| return len(text) // self.ESTIMATED_TOKENS_PER_CHARACTER | |