CCCasEEE/internlm_lagent

1from copy import deepcopy2 3from lagent.schema import ActionReturn, ActionStatusCode, FunctionCall4from .hook import Hook5 6 7class ActionPreprocessor(Hook):8    """The ActionPreprocessor is a hook that preprocesses the action message9    and postprocesses the action return message.10 11    """12 13    def before_action(self, executor, message, session_id):14        assert isinstance(message.formatted, FunctionCall) or (15            isinstance(message.formatted, dict) and 'name' in message.content16            and 'parameters' in message.formatted) or (17                'action' in message.formatted18                and 'parameters' in message.formatted['action']19                and 'name' in message.formatted['action'])20        if isinstance(message.formatted, dict):21            name = message.formatted.get('name',22                                         message.formatted['action']['name'])23            parameters = message.formatted.get(24                'parameters', message.formatted['action']['parameters'])25        else:26            name = message.formatted.name27            parameters = message.formatted.parameters28        message.content = dict(name=name, parameters=parameters)29        return message30 31    def after_action(self, executor, message, session_id):32        action_return = message.content33        if isinstance(action_return, ActionReturn):34            if action_return.state == ActionStatusCode.SUCCESS:35                response = action_return.format_result()36            else:37                response = action_return.errmsg38        else:39            response = action_return40        message.content = response41        return message42 43 44class InternLMActionProcessor(ActionPreprocessor):45 46    def __init__(self, code_parameter: str = 'command'):47        self.code_parameter = code_parameter48 49    def before_action(self, executor, message, session_id):50        message = deepcopy(message)51        assert isinstance(message.formatted, dict) and set(52            message.formatted).issuperset(53                {'tool_type', 'thought', 'action', 'status'})54        if isinstance(message.formatted['action'], str):55            # encapsulate code interpreter arguments56            action_name = next(iter(executor.actions))57            parameters = {self.code_parameter: message.formatted['action']}58            if action_name in ['AsyncIPythonInterpreter']:59                parameters['session_id'] = session_id60            message.formatted['action'] = dict(61                name=action_name, parameters=parameters)62        return super().before_action(executor, message, session_id)63