File size: 4,034 Bytes
acdb322 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | 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
|