InternLM/lagent

👋 加入我们的 𝕏 (Twitter)、Discord 和微信社区。

安装

从源码安装：

git clone https://github.com/InternLM/lagent.git
cd lagent
pip install -e .

使用方法

Lagent 的设计灵感来源于 PyTorch 的设计理念。我们期望通过类比神经网络层，使得工作流程更加清晰直观，用户只需专注于以 Pythonic 的方式创建层并定义层与层之间的消息传递。以下是一个帮助你快速构建多智能体（Multi-agent）应用程序的简易教程。

模型即智能体

智能体使用 AgentMessage 进行通信。

from typing import Dict, List
from lagent.agents import Agent
from lagent.schema import AgentMessage
from lagent.llms import VllmModel, INTERNLM2_META

llm = VllmModel(
    path='Qwen/Qwen2-7B-Instruct',
    meta_template=INTERNLM2_META,
    tp=1,
    top_k=1,
    temperature=1.0,
    stop_words=[''],
    max_new_tokens=1024,
)
system_prompt = '你的回答只能从“典”、“孝”、“急”三个字中选一个。'
agent = Agent(llm, system_prompt)

user_msg = AgentMessage(sender='user', content='今天天气情况')
bot_msg = agent(user_msg)
print(bot_msg)

content='急' sender='Agent' formatted=None extra_info=None type=None receiver=None stream_state=

内存作为状态

在每次前向传播中，输入和输出消息都会被添加到 Agent 的内存中。此过程是在 __call__ 而非 forward 中执行的。请参考以下伪代码：

    def __call__(self, *message):
        message = pre_hooks(message)
        add_memory(message)
        message = self.forward(*message)
        add_memory(message)
        message = post_hooks(message)
        return message

可以通过以下两种方式查看内存：

memory: List[AgentMessage] = agent.memory.get_memory()
print(memory)
print('-' * 120)
dumped_memory: Dict[str, List[dict]] = agent.state_dict()
print(dumped_memory['memory'])

[AgentMessage(content='今天天气情况', sender='user', formatted=None, extra_info=None, type=None, receiver=None, stream_state=), AgentMessage(content='急', sender='Agent', formatted=None, extra_info=None, type=None, receiver=None, stream_state=)]
------------------------------------------------------------------------------------------------------------------------
[{'content': '今天天气情况', 'sender': 'user', 'formatted': None, 'extra_info': None, 'type': None, 'receiver': None, 'stream_state': }, {'content': '急', 'sender': 'Agent', 'formatted': None, 'extra_info': None, 'type': None, 'receiver': None, 'stream_state': }]

清除当前会话的内存（默认 session_id=0）：

agent.reset()

自定义消息聚合（Message Aggregation）

底层会调用 DefaultAggregator 来组装并将 AgentMessage 转换为 OpenAI 消息格式。

    def forward(self, *message: AgentMessage, session_id=0, **kwargs) -> Union[AgentMessage, str]:
        formatted_messages = self.aggregator.aggregate(
            self.memory.get(session_id),
            self.name,
            self.output_format,
            self.template,
        )
        llm_response = self.llm.chat(formatted_messages, **kwargs)
        ...

实现一个简单的聚合器，使其能够支持少样本（few-shots）：

from typing import List, Union
from lagent.memory import Memory
from lagent.prompts import StrParser
from lagent.agents.aggregator import DefaultAggregator

class FewshotAggregator(DefaultAggregator):
    def __init__(self, few_shot: List[dict] = None):
        self.few_shot = few_shot or []

    def aggregate(self,
                  messages: Memory,
                  name: str,
                  parser: StrParser = None,
                  system_instruction: Union[str, dict, List[dict]] = None) -> List[dict]:
        _message = []
        if system_instruction:
            _message.extend(
                self.aggregate_system_intruction(system_instruction))
        _message.extend(self.few_shot)
        messages = messages.get_memory()
        for message in messages:
            if message.sender == name:
                _message.append(
                    ...