InternLM/lagent

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

安装

从源码安装：

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

使用方法

Lagent 的设计灵感来源于 PyTorch 的设计哲学。我们希望通过类比神经网络层，使工作流变得更加清晰直观。因此，用户只需专注于创建“层”（Layer）并以 Pythonic 的方式定义它们之间的消息传递即可。这是一个帮助你快速入门构建多 Agent 应用的简易教程。

以模型作为 Agent

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=

以记忆（Memory）作为状态

在每次前向传播（forward pass）中，输入和输出消息都会被添加到 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()

自定义消息聚合（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-shot（少样本）的简单聚合器：

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(
                    dict(role='assistant', content=message.content))
            else:
                _message.append(
                    dict(role='user', content=message.content))
        return _message