| 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 |
|
|
