| from typing import Dict, List, Optional, Union | |
| from lagent.agents.aggregator.default_aggregator import DefaultAggregator | |
| from lagent.memory.base_memory import Memory | |
| from lagent.prompts.parsers.tool_parser import MixedToolParser, ToolParser, ToolStatusCode | |
| class InternLMToolAggregator(DefaultAggregator): | |
| def __init__(self, | |
| environment_role='environment', | |
| environment_begin='', | |
| environment_end='', | |
| user_names: Optional[List[str]] = None, | |
| few_shot: Optional[List[List[dict]]] = None): | |
| self.environment_role = environment_role | |
| self.environment_begin = environment_begin | |
| self.environment_end = environment_end | |
| self.user_names = user_names or ['user'] | |
| self.few_shot = few_shot or [] | |
| def aggregate(self, | |
| messages: Memory, | |
| name: str, | |
| parser: Union[ToolParser, MixedToolParser], | |
| system_instruction: str = None) -> List[Dict[str, str]]: | |
| _message = [] | |
| messages = messages.get_memory() | |
| if system_instruction: | |
| _message.extend( | |
| self.aggregate_system_intruction(system_instruction)) | |
| tool_instruction = parser.format_instruction() | |
| if tool_instruction: | |
| if isinstance(tool_instruction, str): | |
| tool_instruction = dict( | |
| role='system', content=tool_instruction) | |
| if parser.tool_type: | |
| tool_instruction['name'] = parser.tool_type | |
| if isinstance(tool_instruction, dict): | |
| tool_instruction = [tool_instruction] | |
| _message.extend(tool_instruction) | |
| for shot in self.few_shot: | |
| i = 0 | |
| while i < len(shot): | |
| msg = shot[i] | |
| if msg['role'] in ['assistant', 'user', 'system']: | |
| _message.append(msg) | |
| elif msg['role'] == self.environment_role: | |
| if not msg['content'].startswith(self.environment_begin): | |
| msg['content'] = self.environment_begin + msg['content'] | |
| if not msg['content'].endswith(self.environment_end): | |
| msg['content'] += self.environment_end | |
| _message.append(msg) | |
| elif msg['role'] in ['thought', 'language']: | |
| if i < len(shot) - 1 and shot[i + 1]['role'] == 'tool': | |
| _message.append( | |
| dict( | |
| role='assistant', | |
| content=parser.format_response( | |
| dict( | |
| tool_type=shot[i + 1]['name'], | |
| thought=msg['content'], | |
| action=shot[i + 1]['content'], | |
| status=None)))) | |
| i += 1 | |
| else: | |
| _message.append( | |
| dict( | |
| role='assistant', | |
| content=parser.format_response( | |
| dict( | |
| tool_type=None, | |
| thought=msg['content'], | |
| action=None, | |
| status=None)))) | |
| else: | |
| raise KeyError(f'Unkown role: {msg["role"]}') | |
| i += 1 | |
| tool_type = None | |
| for message in messages: | |
| if message.sender == name: | |
| if isinstance(message.formatted, dict): | |
| parsed = message.formatted | |
| if parsed['status'] == ToolStatusCode.PARSING_ERROR: | |
| continue | |
| _message.append( | |
| dict( | |
| role='assistant', | |
| content=parser.format_response(parsed))) | |
| tool_type = parsed['tool_type'] | |
| else: | |
| _message.append( | |
| dict(role='assistant', content=str(message.content))) | |
| elif message.sender in self.user_names: | |
| _message.append(dict(role='user', content=message.content)) | |
| else: | |
| msg = dict( | |
| role=self.environment_role, | |
| content=self.environment_begin + str(message.content) + | |
| self.environment_end) | |
| if tool_type: | |
| msg['name'] = tool_type | |
| _message.append(msg) | |
| return _message | |