mirror of
https://github.com/browser-use/browser-use
synced 2026-05-06 17:52:15 +02:00
65 lines
1.7 KiB
Python
65 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, Literal
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field
|
|
|
|
from browser_use.llm.messages import (
|
|
BaseMessage,
|
|
UserMessage,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
pass
|
|
|
|
|
|
SupportedMessageTypes = Literal['init', 'memory']
|
|
|
|
|
|
class MessageMetadata(BaseModel):
|
|
"""Metadata for a message"""
|
|
|
|
message_type: SupportedMessageTypes | None = None
|
|
|
|
|
|
class ManagedMessage(BaseModel):
|
|
"""A message with its metadata"""
|
|
|
|
message: BaseMessage
|
|
metadata: MessageMetadata = Field(default_factory=MessageMetadata)
|
|
|
|
|
|
class MessageHistory(BaseModel):
|
|
"""History of messages with metadata"""
|
|
|
|
messages: list[ManagedMessage] = Field(default_factory=list)
|
|
|
|
model_config = ConfigDict(arbitrary_types_allowed=True)
|
|
|
|
def add_message(self, message: BaseMessage, metadata: MessageMetadata, position: int | None = None) -> None:
|
|
"""Add message with metadata to history"""
|
|
if position is None:
|
|
self.messages.append(ManagedMessage(message=message, metadata=metadata))
|
|
else:
|
|
self.messages.insert(position, ManagedMessage(message=message, metadata=metadata))
|
|
|
|
def get_messages(self) -> list[BaseMessage]:
|
|
"""Get all messages"""
|
|
return [m.message for m in 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].message, UserMessage):
|
|
self.messages.pop()
|
|
|
|
|
|
class MessageManagerState(BaseModel):
|
|
"""Holds the state for MessageManager"""
|
|
|
|
history: MessageHistory = Field(default_factory=MessageHistory)
|
|
tool_id: int = 1
|
|
agent_history_description: str = '<s>Agent initialized</s>\n'
|
|
read_state_description: str = ''
|
|
|
|
model_config = ConfigDict(arbitrary_types_allowed=True)
|