File size: 1,938 Bytes
db4810d |
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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
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]
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)
def get_index(self) -> int | None:
"""Get the index of the action"""
# {'clicked_element': {'index':5}}
params = self.model_dump(exclude_unset=True).values()
if not params:
return None
for param in params:
if param is not None and 'index' in param:
return param['index']
return None
def set_index(self, index: int):
"""Overwrite the index of the action"""
# Get the action name and params
action_data = self.model_dump(exclude_unset=True)
action_name = next(iter(action_data.keys()))
action_params = getattr(self, action_name)
# Update the index directly on the model
if hasattr(action_params, 'index'):
action_params.index = index
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()])
|