LangChain —— Message —— How to trim messages

一、概述

 所有模型都有 有限的 上下文窗口，这意味着它们可以作为输入的 token 数量是有限的。如果你有很长的消息，或者一个 chain 或 agent 累积了很长的 历史消息，你需要管理你传递给模型的消息的长度。
 trim_messages util 提供了一些基本策略，用于将消息列表修剪为特定的 token 长度。

二、获取最后的 max_tokens 令牌

 为了获取消息列表中的最后一个 max_tokens，我们可以设置 strategy=“last”。请注意，对于我们的 token_counter，我们可以将其传入到一个函数 (下面将详细介绍) 或一个语言模型 (因为语言模型有一个消息令牌计数方法) 中。当调整消息以适应特定模型的上下文窗口时，将其传入到模型是有意义的：

# pip install -U langchain-openai
from langchain_core.messages import (
    AIMessage,
    HumanMessage,
    SystemMessage,
    trim_messages,
)
from langchain_openai import ChatOpenAI

messages = [
    SystemMessage("you're a good assistant, you always respond with a joke."),
    HumanMessage("i wonder why it's called langchain"),
    AIMessage(
        'Well, I guess they thought "WordRope" and "SentenceString" just didn\'t have the same ring to it!'
    ),
    HumanMessage("and who is harrison chasing anyways"),
    AIMessage(
        "Hmmm let me think.\n\nWhy, he's probably chasing after the last cup of coffee in the office!"
    ),
    HumanMessage("what do you call a speechless parrot"),
]

trim_messages(
    messages,
    max_tokens=45,
    strategy="last",
    token_counter=ChatOpenAI(model="gpt-4o"),
)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

 如果我们想始终保留初始系统消息，我们可以指定 include_system=True：
 如果我们想允许拆分消息的内容，我们可以指定 allow_partial=True：
 如果我们需要确保我们的第一条消息 (不包括 SystemMessage) 始终是特定类型的，我们可以指定 start_on：

trim_messages(
    messages,
    max_tokens=60,
    strategy="last",
    token_counter=ChatOpenAI(model="gpt-4o"),
    include_system=True,
    start_on="human",
)
1
2
3
4
5
6
7
8

三、获取第一个 max_tokens 令牌

 我们可以通过指定 strategy=“first” 来执行获取第一个 max_tokens 的翻转操作：

trim_messages(
    messages,
    max_tokens=45,
    strategy="first",
    token_counter=ChatOpenAI(model="gpt-4o"),
)
1
2
3
4
5
6

四、编写自定义令牌计数器

 我们可以编写一个自定义令牌计数器函数，该函数接收消息列表并返回一个整数。

from typing import List
# pip install tiktoken
import tiktoken
from langchain_core.messages import BaseMessage, ToolMessage

def str_token_counter(text: str) -> int:
    enc = tiktoken.get_encoding("o200k_base")
    return len(enc.encode(text))

def tiktoken_counter(messages: List[BaseMessage]) -> int:
    """Approximately reproduce https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb

    For simplicity only supports str Message.contents.
    """
    num_tokens = 3  # every reply is primed with <|start|>assistant<|message|>
    tokens_per_message = 3
    tokens_per_name = 1
    for msg in messages:
        if isinstance(msg, HumanMessage):
            role = "user"
        elif isinstance(msg, AIMessage):
            role = "assistant"
        elif isinstance(msg, ToolMessage):
            role = "tool"
        elif isinstance(msg, SystemMessage):
            role = "system"
        else:
            raise ValueError(f"Unsupported messages type {msg.__class__}")
        num_tokens += (
            tokens_per_message
            + str_token_counter(role)
            + str_token_counter(msg.content)
        )
        if msg.name:
            num_tokens += tokens_per_name + str_token_counter(msg.name)
    return num_tokens


trim_messages(
    messages,
    max_tokens=45,
    strategy="last",
    token_counter=tiktoken_counter,
)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44

• 定义 str_token_counter 函数
  • 该函数接受一个字符串 text 并返回该字符串的令牌数量。
  • 使用 tiktoken.get_encoding(“o200k_base”) 获取编码器，然后使用 enc.encode(text) 将文本编码为令牌，并返回令牌的长度。
• 定义 tiktoken_counter 函数：
  • 该函数接受一个 BaseMessage 类型的消息列表 messages 并返回总的令牌数量。
  • 由于每个回复都以 <|start|>assistance<|message|> 开头，所以每个消息列表初始都默认有 3 个 token，每个消息有一个基本的令牌数 tokens_per_message，每个 name 属性预设的固定令牌数 tokens_per_name，假设其值为 1。
  • 函数通过迭代消息列表，并根据消息的角色 (如 user、assistant、tool、system) 计算令牌数量。
  • 根据消息的 tokens_per_message、role、content，计算总的令牌数。
  • 如果遇到不支持的消息类型，会引发 ValueError。
  • 对于每个消息，如果消息对象 msg 有 name 属性 (即 msg.name 不为 None 或空)，那么就要计算该 name 属性所包含的令牌数量，并将其加入到总令牌数 num_tokens 中。
• 调用 trim_messages 函数 (假设定义在其他地方)：
  • messages：要处理的消息列表。
  • max_tokens=45：最大允许的令牌数。
  • strategy=“last”：修剪策略 (假设修剪最后的消息)。
  • token_counter=tiktoken_counter：用于计算令牌数的函数。

五、连成链

 trim_message可以以命令式（如上所述）或声明式的方式使用，从而便于与链中的其他组件组合。

llm = ChatOpenAI(model="gpt-4o")

# Notice we don't pass in messages. This creates
# a RunnableLambda that takes messages as input
trimmer = trim_messages(
    max_tokens=45,
    strategy="last",
    token_counter=llm,
    include_system=True,
)

chain = trimmer | llm
chain.invoke(messages)
1
2
3
4
5
6
7
8
9
10
11
12
13

 查看 LangSmith 跟踪，我们可以看到，在消息传递到模型之前，它们首先被修剪。
 如果只看 trimer，我们可以看到它是一个Runnable对象，可以像所有Runnables一样被调用：

trimmer.invoke(messages)
1

六、使用 ChatMessageHistory

 在处理聊天历史记录时，修剪消息特别有用，因为聊天历史记录可能会变得任意长：

from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory

chat_history = InMemoryChatMessageHistory(messages=messages[:-1])

def dummy_get_session_history(session_id):
    if session_id != "1":
        return InMemoryChatMessageHistory()
    return chat_history

llm = ChatOpenAI(model="gpt-4o")

trimmer = trim_messages(
    max_tokens=45,
    strategy="last",
    token_counter=llm,
    include_system=True,
)

chain = trimmer | llm
chain_with_history = RunnableWithMessageHistory(chain, dummy_get_session_history)
chain_with_history.invoke(
    [HumanMessage("what do you call a speechless parrot")],
    config={"configurable": {"session_id": "1"}},
)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

• 第四行，创建一个 InMemoryChatMessageHistory 对象 chat_history，并初始化它的消息历史为 messages 列表 (除了最后一个消息)。
• 第六行，定义一个函数 dummy_get_session_history，根据 session_id 返回相应的聊天历史记录：
  • 如果 session_id 不等于 “1”，则返回一个新的空的 InMemoryChatMessageHistory 对象。
  • 如果 session_id 等于 “1”，则返回之前定义的 chat_history。
• 倒数第五行，创建一个 RunnableWithMessageHistory 对象 chain_with_history，将 chain 和 dummy_get_session_history 结合在一起，使其能够处理带有历史记录的消息。
• 倒数第四行，调用 chain_with_history 对象的 invoke 方法，传入一个包含 HumanMessage(“what do you call a speechless parrot”) 的列表，和配置 {“configurable”: {“session_id”: “1”}}：
  • 该方法将根据 session_id 为 “1” 调用 dummy_get_session_history 返回相应的历史记录。
  • 将消息传递给链 chain，先修剪再由 llm 处理。

 查看 LangSmith 跟踪，我们可以看到我们检索了所有消息，但在将消息传递给模型之前，它们被修剪成只有系统消息和最后一条人类消息。