Files
browser-use/browser_use/agent/message_manager/views.py
2025-07-01 12:32:54 +02:00

98 lines
2.6 KiB
Python

from __future__ import annotations
from typing import TYPE_CHECKING
from pydantic import BaseModel, ConfigDict, Field
from browser_use.llm.messages import (
BaseMessage,
UserMessage,
)
if TYPE_CHECKING:
pass
class HistoryItem(BaseModel):
"""Represents a single agent history item with its data and string representation"""
step_number: int | None = None
evaluation_previous_goal: str | None = None
memory: str | None = None
next_goal: str | None = None
action_results: str | None = None
error: str | None = None
system_message: str | None = None
model_config = ConfigDict(arbitrary_types_allowed=True)
def model_post_init(self, __context) -> None:
"""Validate that error and system_message are not both provided"""
if self.error is not None and self.system_message is not None:
raise ValueError('Cannot have both error and system_message at the same time')
def to_string(self) -> str:
"""Get string representation of the history item"""
step_str = f'step_{self.step_number}' if self.step_number is not None else 'step_unknown'
if self.error:
return f"""<{step_str}>
{self.error}
</{step_str}>"""
elif self.system_message:
return f"""<sys>
{self.system_message}
</sys>"""
else:
content_parts = [
f'Evaluation of Previous Step: {self.evaluation_previous_goal}',
f'Memory: {self.memory}',
f'Next Goal: {self.next_goal}',
]
if self.action_results:
content_parts.append(self.action_results)
content = '\n'.join(content_parts)
return f"""<{step_str}>
{content}
</{step_str}>"""
class MessageHistory(BaseModel):
"""History of messages"""
messages: list[BaseMessage] = Field(default_factory=list)
model_config = ConfigDict(arbitrary_types_allowed=True)
def add_message(self, message: BaseMessage, position: int | None = None) -> None:
"""Add message to history"""
if position is None:
self.messages.append(message)
else:
self.messages.insert(position, message)
def get_messages(self) -> list[BaseMessage]:
"""Get all messages"""
return self.messages
def remove_last_state_message(self) -> None:
"""Remove last state message from history"""
if len(self.messages) > 2 and isinstance(self.messages[-1], UserMessage):
self.messages.pop()
class MessageManagerState(BaseModel):
"""Holds the state for MessageManager"""
history: MessageHistory = Field(default_factory=MessageHistory)
tool_id: int = 1
agent_history_items: list[HistoryItem] = Field(
default_factory=lambda: [HistoryItem(step_number=0, system_message='Agent initialized')]
)
read_state_description: str = ''
model_config = ConfigDict(arbitrary_types_allowed=True)