Spaces:
Runtime error
Runtime error
| # -*- coding: utf-8 -*- | |
| # @Time : 2025/1/2 | |
| # @Author : wenshao | |
| # @ProjectName: browser-use-webui | |
| # @FileName: custom_massage_manager.py | |
| 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 ( | |
| HumanMessage, | |
| AIMessage | |
| ) | |
| 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, | |
| ): | |
| 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, | |
| ) | |
| # Custom: Move Task info to state_message | |
| self.history = MessageHistory() | |
| self._add_message_with_tokens(self.system_prompt) | |
| tool_calls = [ | |
| { | |
| 'name': 'CustomAgentOutput', | |
| 'args': { | |
| 'current_state': { | |
| 'prev_action_evaluation': 'Unknown - No previous actions to evaluate.', | |
| 'important_contents': '', | |
| 'completed_contents': '', | |
| 'thought': 'Now Google is open. Need to type OpenAI to search.', | |
| 'summary': 'Type OpenAI to search.', | |
| }, | |
| 'action': [], | |
| }, | |
| 'id': '', | |
| 'type': 'tool_call', | |
| } | |
| ] | |
| if self.tool_call_in_content: | |
| # openai throws error if tool_calls are not responded -> move to content | |
| example_tool_call = AIMessage( | |
| content=f'{tool_calls}', | |
| tool_calls=[], | |
| ) | |
| else: | |
| example_tool_call = AIMessage( | |
| content=f'', | |
| tool_calls=tool_calls, | |
| ) | |
| self._add_message_with_tokens(example_tool_call) | |
| def add_state_message( | |
| self, | |
| state: BrowserState, | |
| result: Optional[List[ActionResult]] = None, | |
| step_info: Optional[AgentStepInfo] = None, | |
| ) -> None: | |
| """Add browser state as human message""" | |
| # if keep in memory, add to directly to history and add state without result | |
| if result: | |
| for r in result: | |
| if r.include_in_memory: | |
| if r.extracted_content: | |
| msg = HumanMessage(content=str(r.extracted_content)) | |
| self._add_message_with_tokens(msg) | |
| if r.error: | |
| msg = HumanMessage( | |
| content=str(r.error)[-self.max_error_length:] | |
| ) | |
| self._add_message_with_tokens(msg) | |
| result = None # if result in history, we dont want to add it again | |
| # otherwise add state message and result to next message (which will not stay in memory) | |
| 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) | |