File size: 1,140 Bytes
4fd8631 | 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 | from dataclasses import dataclass
from typing import Type
from browser_use.agent.views import AgentOutput
from browser_use.controller.registry.views import ActionModel
from pydantic import BaseModel, ConfigDict, Field, create_model
@dataclass
class CustomAgentStepInfo:
step_number: int
max_steps: int
task: str
add_infos: str
memory: str
task_progress: str
future_plans: str
class CustomAgentBrain(BaseModel):
prev_action_evaluation: str
important_contents: str
task_progress: str
future_plans: str
thought: str
summary: str
class CustomAgentOutput(AgentOutput):
model_config = ConfigDict(arbitrary_types_allowed=True)
current_state: CustomAgentBrain
action: list[ActionModel]
@staticmethod
def type_with_custom_actions(
custom_actions: Type[ActionModel],
) -> Type["CustomAgentOutput"]:
return create_model(
"AgentOutput",
__base__=CustomAgentOutput,
action=(
list[custom_actions],
Field(...),
),
__module__=CustomAgentOutput.__module__,
)
|