Files
browser-use/browser_use/controller/registry/views.py
2024-11-22 13:08:37 +01:00

50 lines
1.3 KiB
Python

from typing import Callable, Dict, Type
from pydantic import BaseModel, ConfigDict
class RegisteredAction(BaseModel):
"""Model for a registered action"""
name: str
description: str
function: Callable
param_model: Type[BaseModel]
requires_browser: bool = False
model_config = ConfigDict(arbitrary_types_allowed=True)
def prompt_description(self) -> str:
"""Get a description of the action for the prompt"""
skip_keys = ['title']
s = f'{self.description}: \n'
s += '{' + str(self.name) + ': '
s += str(
{
k: {sub_k: sub_v for sub_k, sub_v in v.items() if sub_k not in skip_keys}
for k, v in self.param_model.schema()['properties'].items()
}
)
s += '}'
return s
class ActionModel(BaseModel):
"""Base model for dynamically created action models"""
# this will have all the registered actions, e.g.
# click_element = param_model = ClickElementParams
# done = param_model = None
#
model_config = ConfigDict(arbitrary_types_allowed=True)
class ActionRegistry(BaseModel):
"""Model representing the action registry"""
actions: Dict[str, RegisteredAction] = {}
def get_prompt_description(self) -> str:
"""Get a description of all actions for the prompt"""
return '\n'.join([action.prompt_description() for action in self.actions.values()])