| import json | |
| import warnings | |
| from copy import deepcopy | |
| from typing import Callable, Dict, List, Union | |
| from lagent.actions import ActionExecutor, AsyncActionExecutor, AsyncIPythonInterpreter, IPythonInteractive | |
| from lagent.agents.agent import Agent, AsyncAgent | |
| from lagent.agents.aggregator import InternLMToolAggregator | |
| from lagent.hooks import InternLMActionProcessor | |
| from lagent.llms import BaseLLM | |
| from lagent.memory import Memory | |
| from lagent.prompts.parsers import InterpreterParser, MixedToolParser, PluginParser, ToolStatusCode | |
| from lagent.schema import AgentMessage | |
| from lagent.utils import create_object | |
| API_PREFIX = ( | |
| "This is the subfunction for tool '{tool_name}', you can use this tool. " | |
| 'The description of this function is: \n{description}') | |
| META_CN = ('当开启工具以及代码时，根据需求选择合适的工具进行调用') | |
| INTERPRETER_CN = ('你现在已经能够在一个有状态的 Jupyter 笔记本环境中运行 Python 代码。' | |
| '当你向 python 发送含有 Python 代码的消息时，它将在该环境中执行。' | |
| '这个工具适用于多种场景，如数据分析或处理（包括数据操作、统计分析、图表绘制），' | |
| '复杂的计算问题（解决数学和物理难题），编程示例（理解编程概念或特性），' | |
| '文本处理和分析（比如文本解析和自然语言处理），' | |
| '机器学习和数据科学（用于展示模型训练和数据可视化），' | |
| '以及文件操作和数据导入（处理CSV、JSON等格式的文件）。') | |
| PLUGIN_CN = ('你可以使用如下工具：' | |
| '\n{prompt}\n' | |
| '如果你已经获得足够信息，请直接给出答案. 避免不必要的工具调用! ' | |
| '同时注意你可以使用的工具，不要随意捏造！') | |
| def get_plugin_prompt(actions, api_desc_template=API_PREFIX): | |
| plugin_descriptions = [] | |
| for action in actions if isinstance(actions, list) else [actions]: | |
| action = create_object(action) | |
| action_desc = deepcopy(action.description) | |
| if action.is_toolkit: | |
| for api in action_desc['api_list']: | |
| api['name'] = f"{action.name}.{api['name']}" | |
| api['description'] = api_desc_template.format( | |
| tool_name=action.name, description=api['description']) | |
| api['parameters'] = [ | |
| param for param in api['parameters'] | |
| if param['name'] in api['required'] | |
| ] | |
| plugin_descriptions.append(api) | |
| else: | |
| action_desc['description'] = api_desc_template.format( | |
| tool_name=action.name, description=action_desc['description']) | |
| action_desc['parameters'] = [ | |
| param for param in action_desc['parameters'] | |
| if param['name'] in action_desc['required'] | |
| ] | |
| plugin_descriptions.append(action_desc) | |
| return json.dumps(plugin_descriptions, ensure_ascii=False, indent=4) | |
| class AgentForInternLM(Agent): | |
| _INTERNAL_AGENT_CLS = Agent | |
| def __init__( | |
| self, | |
| llm: Union[BaseLLM, Dict], | |
| plugins: Union[dict, List[dict]] = None, | |
| interpreter: dict = None, | |
| template: Union[str, dict, List[dict]] = None, | |
| memory: Dict = dict(type=Memory), | |
| output_format: Dict = dict( | |
| type=MixedToolParser, | |
| template=META_CN, | |
| parsers=[ | |
| dict(type=PluginParser, template=PLUGIN_CN), | |
| dict(type=InterpreterParser, template=INTERPRETER_CN), | |
| ]), | |
| aggregator: Dict = dict(type=InternLMToolAggregator), | |
| action_hooks: List = [dict(type=InternLMActionProcessor)], | |
| finish_condition: Callable[ | |
| [AgentMessage], | |
| bool] = lambda m: m.formatted['status'] == ToolStatusCode.NO_TOOL, | |
| max_turn: int = 4, | |
| **kwargs, | |
| ): | |
| agent = dict( | |
| type=self._INTERNAL_AGENT_CLS, | |
| llm=llm, | |
| template=template, | |
| output_format=output_format, | |
| memory=memory, | |
| aggregator=aggregator, | |
| hooks=kwargs.pop('hooks', None), | |
| ) | |
| self.agent = create_object(agent) | |
| self.plugin_executor = plugins and ActionExecutor( | |
| plugins, hooks=action_hooks) | |
| self.interpreter_executor = interpreter and ActionExecutor( | |
| interpreter, hooks=action_hooks) | |
| if not (self.plugin_executor or self.interpreter_executor): | |
| warnings.warn( | |
| 'Neither plugin nor interpreter executor is initialized. ' | |
| 'An exception will be thrown when the agent call a tool.') | |
| self.finish_condition = finish_condition | |
| self.max_turn = max_turn | |
| super().__init__(**kwargs) | |
| def forward(self, message: AgentMessage, session_id=0, **kwargs): | |
| if isinstance(message, str): | |
| message = AgentMessage(sender='user', content=message) | |
| for _ in range(self.max_turn): | |
| message = self.agent(message, session_id=session_id, **kwargs) | |
| assert isinstance(message.formatted, dict) | |
| if self.finish_condition(message): | |
| return message | |
| if message.formatted['tool_type']: | |
| tool_type = message.formatted["tool_type"] | |
| executor = getattr(self, f'{tool_type}_executor', None) | |
| if not executor: | |
| raise RuntimeError(f'No available {tool_type} executor') | |
| message = executor(message, session_id=session_id) | |
| return message | |
| def get_steps(self, session_id=0): | |
| steps, tool_type = [], None | |
| for msg in self.agent.memory.get_memory(session_id): | |
| if msg.sender == self.agent.name: | |
| steps.append( | |
| dict(role='thought', content=msg.formatted['thought'])) | |
| if msg.formatted['tool_type']: | |
| tool_type = msg.formatted['tool_type'] | |
| steps.append( | |
| dict( | |
| role='tool', | |
| content=msg.formatted['action'], | |
| name=tool_type)) | |
| elif msg.sender != 'user': | |
| feedback = dict(role='environment', content=msg.content) | |
| if tool_type: | |
| feedback['name'] = tool_type | |
| steps.append(feedback) | |
| return steps | |
| class MathCoder(AgentForInternLM): | |
| def __init__( | |
| self, | |
| llm: Union[BaseLLM, Dict], | |
| interpreter: dict = dict( | |
| type=IPythonInteractive, timeout=20, max_out_len=8192), | |
| template: Union[str, dict, List[dict]] = None, | |
| memory: Dict = dict(type=Memory), | |
| output_format: Dict = dict( | |
| type=InterpreterParser, | |
| template= | |
| ('Integrate step-by-step reasoning and Python code to solve math problems ' | |
| 'using the following guidelines:\n' | |
| '- Analyze the question and write jupyter code to solve the problem;\n' | |
| r"- Present the final result in LaTeX using a '\boxed{{}}' without any " | |
| 'units. \n')), | |
| aggregator: Dict = dict(type=InternLMToolAggregator), | |
| action_hooks: List = [dict(type=InternLMActionProcessor)], | |
| finish_condition: Callable[ | |
| [AgentMessage], | |
| bool] = lambda m: m.formatted['status'] == ToolStatusCode.NO_TOOL, | |
| max_turn: int = 6, | |
| **kwargs, | |
| ): | |
| kwargs.pop('plugins', None) | |
| super().__init__( | |
| llm=llm, | |
| interpreter=interpreter, | |
| template=template, | |
| memory=memory, | |
| output_format=output_format, | |
| aggregator=aggregator, | |
| action_hooks=action_hooks, | |
| finish_condition=finish_condition, | |
| max_turn=max_turn, | |
| **kwargs) | |
| class AsyncAgentForInternLM(AsyncAgent): | |
| _INTERNAL_AGENT_CLS = AsyncAgent | |
| def __init__( | |
| self, | |
| llm: Union[BaseLLM, Dict], | |
| plugins: Union[dict, List[dict]] = None, | |
| interpreter: dict = None, | |
| template: Union[str, dict, List[dict]] = None, | |
| memory: Dict = dict(type=Memory), | |
| output_format: Dict = dict( | |
| type=MixedToolParser, | |
| template=META_CN, | |
| parsers=[ | |
| dict(type=PluginParser, template=PLUGIN_CN), | |
| dict(type=InterpreterParser, template=INTERPRETER_CN), | |
| ]), | |
| aggregator: Dict = dict(type=InternLMToolAggregator), | |
| action_hooks: List = [dict(type=InternLMActionProcessor)], | |
| finish_condition: Callable[ | |
| [AgentMessage], | |
| bool] = lambda m: m.formatted['status'] == ToolStatusCode.NO_TOOL, | |
| max_turn: int = 4, | |
| **kwargs, | |
| ): | |
| agent = dict( | |
| type=self._INTERNAL_AGENT_CLS, | |
| llm=llm, | |
| template=template, | |
| output_format=output_format, | |
| memory=memory, | |
| aggregator=aggregator, | |
| hooks=kwargs.pop('hooks', None), | |
| ) | |
| self.agent = create_object(agent) | |
| self.plugin_executor = plugins and AsyncActionExecutor( | |
| plugins, hooks=action_hooks) | |
| self.interpreter_executor = interpreter and AsyncActionExecutor( | |
| interpreter, hooks=action_hooks) | |
| if not (self.plugin_executor or self.interpreter_executor): | |
| warnings.warn( | |
| 'Neither plugin nor interpreter executor is initialized. ' | |
| 'An exception will be thrown when the agent call a tool.') | |
| self.finish_condition = finish_condition | |
| self.max_turn = max_turn | |
| super().__init__(**kwargs) | |
| async def forward(self, message: AgentMessage, session_id=0, **kwargs): | |
| if isinstance(message, str): | |
| message = AgentMessage(sender='user', content=message) | |
| for _ in range(self.max_turn): | |
| message = await self.agent( | |
| message, session_id=session_id, **kwargs) | |
| assert isinstance(message.formatted, dict) | |
| if self.finish_condition(message): | |
| return message | |
| if message.formatted['tool_type']: | |
| tool_type = message.formatted["tool_type"] | |
| executor = getattr(self, f'{tool_type}_executor', None) | |
| if not executor: | |
| raise RuntimeError(f'No available {tool_type} executor') | |
| message = await executor(message, session_id=session_id) | |
| return message | |
| def get_steps(self, session_id=0): | |
| steps, tool_type = [], None | |
| for msg in self.agent.memory.get_memory(session_id): | |
| if msg.sender == self.agent.name: | |
| steps.append( | |
| dict(role='thought', content=msg.formatted['thought'])) | |
| if msg.formatted['tool_type']: | |
| tool_type = msg.formatted['tool_type'] | |
| steps.append( | |
| dict( | |
| role='tool', | |
| content=msg.formatted['action'], | |
| name=tool_type)) | |
| elif msg.sender != 'user': | |
| feedback = dict(role='environment', content=msg.content) | |
| if tool_type: | |
| feedback['name'] = tool_type | |
| steps.append(feedback) | |
| return steps | |
| class AsyncMathCoder(AsyncAgentForInternLM): | |
| def __init__( | |
| self, | |
| llm: Union[BaseLLM, Dict], | |
| interpreter: dict = dict(type=AsyncIPythonInterpreter), | |
| template: Union[str, dict, List[dict]] = None, | |
| memory: Dict = dict(type=Memory), | |
| output_format: Dict = dict( | |
| type=InterpreterParser, | |
| template= | |
| ('Integrate step-by-step reasoning and Python code to solve math problems ' | |
| 'using the following guidelines:\n' | |
| '- Analyze the question and write jupyter code to solve the problem;\n' | |
| r"- Present the final result in LaTeX using a '\boxed{{}}' without any " | |
| 'units. \n')), | |
| aggregator: Dict = dict(type=InternLMToolAggregator), | |
| action_hooks: List = [dict(type=InternLMActionProcessor)], | |
| finish_condition: Callable[ | |
| [AgentMessage], | |
| bool] = lambda m: m.formatted['status'] == ToolStatusCode.NO_TOOL, | |
| max_turn: int = 6, | |
| **kwargs, | |
| ): | |
| kwargs.pop('plugins', None) | |
| super().__init__( | |
| llm=llm, | |
| interpreter=interpreter, | |
| template=template, | |
| memory=memory, | |
| output_format=output_format, | |
| aggregator=aggregator, | |
| action_hooks=action_hooks, | |
| finish_condition=finish_condition, | |
| max_turn=max_turn, | |
| **kwargs) | |
| async def forward(self, message: AgentMessage, session_id=0, **kwargs): | |
| try: | |
| return await super().forward(message, session_id, **kwargs) | |
| finally: | |
| interpreter = next( | |
| iter(self.interpreter_executor.actions.values())) | |
| if interpreter.name == 'AsyncIPythonInterpreter': | |
| await interpreter.close_session(session_id) | |