跳到内容

短期记忆API参考

  • SummarizationNode

    一个LangGraph节点,当消息超出令牌限制时,会对其进行总结,并用一个总结消息替换它们。

  • SummarizationResult

    消息总结的结果。

  • RunningSummary

    用于存储先前总结信息的对象。

函数

  • summarize_messages

    当消息超出令牌限制时,对其进行总结,并用一个总结消息替换它们。

SummarizationNode

基类:RunnableCallable

一个LangGraph节点,当消息超出令牌限制时,会对其进行总结,并用一个总结消息替换它们。

方法

  • __init__

    一个LangGraph节点,当消息超出令牌限制时,会对其进行总结,并用一个总结消息替换它们。

__init__

__init__(
    *,
    model: LanguageModelLike,
    max_tokens: int,
    max_tokens_before_summary: int | None = None,
    max_summary_tokens: int = 256,
    token_counter: TokenCounter = count_tokens_approximately,
    initial_summary_prompt: ChatPromptTemplate = DEFAULT_INITIAL_SUMMARY_PROMPT,
    existing_summary_prompt: ChatPromptTemplate = DEFAULT_EXISTING_SUMMARY_PROMPT,
    final_prompt: ChatPromptTemplate = DEFAULT_FINAL_SUMMARY_PROMPT,
    input_messages_key: str = "messages",
    output_messages_key: str = "summarized_messages",
    name: str = "summarization",
) -> None

一个LangGraph节点,当消息超出令牌限制时,会对其进行总结,并用一个总结消息替换它们。

从最旧到最新处理消息:一旦累积消息令牌数达到 max_tokens_before_summary,则 max_tokens_before_summary 内的所有消息(如果存在系统消息则排除)都将被总结,并替换为一个新的总结消息。最终的消息列表是 [总结消息] + 剩余消息。

参数

  • model (LanguageModelLike) –

    用于生成总结的语言模型。

  • max_tokens (int) –

    最终输出中返回的最大令牌数。仅在总结后强制执行。

  • max_tokens_before_summary (int | None, 默认值:None ) –

    在触发总结之前要累积的最大令牌数。如果未提供,则默认为 max_tokens 的相同值。这允许在需要时将更多令牌放入总结LLM中。

    注意

    如果 max_tokens_before_summary 内的最后一条消息是带有工具调用的AI消息,那么所有后续的、相应的工具消息也将被总结。

    注意

    如果要总结的令牌数大于 max_tokens,则只总结其中最后 max_tokens 个。这样做是为了防止超出总结LLM的上下文窗口(假定上限为 max_tokens)。

  • max_summary_tokens (int, 默认值:256 ) –

    为总结预留的最大令牌数。

    注意

    此参数不会传递给生成总结的LLM以限制总结的长度。它仅用于正确估算最大允许令牌预算。如果需要强制执行,则需要将 model.bind(max_tokens=max_summary_tokens) 作为 model 参数传递给此函数。

  • token_counter (TokenCounter, 默认值:count_tokens_approximately ) –

    用于计算消息中令牌的函数。默认为近似计数。要获得更准确的计数,可以使用 model.get_num_tokens_from_messages

  • initial_summary_prompt (ChatPromptTemplate, 默认值:DEFAULT_INITIAL_SUMMARY_PROMPT ) –

    用于生成第一个总结的提示模板。

  • existing_summary_prompt (ChatPromptTemplate, 默认值:DEFAULT_EXISTING_SUMMARY_PROMPT ) –

    用于更新现有(运行中)总结的提示模板。

  • final_prompt (ChatPromptTemplate, 默认值:DEFAULT_FINAL_SUMMARY_PROMPT ) –

    在返回之前将总结与剩余消息结合起来的提示模板。

  • input_messages_key (str, 默认值:'messages' ) –

    输入图状态中包含要总结的消息列表的键。

  • output_messages_key (str, 默认值:'summarized_messages' ) –

    状态更新中包含已更新消息列表的键。

    警告

    默认情况下,output_messages_key 不同于 input_messages_key。这样做是为了将总结后的消息从图状态中的主消息列表(即 input_messages_key)中解耦。只有当您希望覆盖主消息列表(即 input_messages_key)时,才应将它们设置为相同。

  • name (str, 默认值:'summarization' ) –

    总结节点的名称。

返回

  • None

    LangGraph状态更新,格式如下

    {
        "output_messages_key": <list of updated messages ready to be input to the LLM after summarization, including a message with a summary (if any)>,
        "context": {"running_summary": <RunningSummary object>}
    }
    

示例
from typing import Any, TypedDict
from langchain_openai import ChatOpenAI
from langchain_core.messages import AnyMessage
from langgraph.graph import StateGraph, START, MessagesState
from langgraph.checkpoint.memory import InMemorySaver
from langmem.short_term import SummarizationNode, RunningSummary

model = ChatOpenAI(model="gpt-4o")
summarization_model = model.bind(max_tokens=128)


class State(MessagesState):
    context: dict[str, Any]


class LLMInputState(TypedDict):
    summarized_messages: list[AnyMessage]
    context: dict[str, Any]


summarization_node = SummarizationNode(
    model=summarization_model,
    max_tokens=256,
    max_tokens_before_summary=256,
    max_summary_tokens=128,
)


def call_model(state: LLMInputState):
    response = model.invoke(state["summarized_messages"])
    return {"messages": [response]}


checkpointer = InMemorySaver()
workflow = StateGraph(State)
workflow.add_node(call_model)
workflow.add_node("summarize", summarization_node)
workflow.add_edge(START, "summarize")
workflow.add_edge("summarize", "call_model")
graph = workflow.compile(checkpointer=checkpointer)

config = {"configurable": {"thread_id": "1"}}
graph.invoke({"messages": "hi, my name is bob"}, config)
graph.invoke({"messages": "write a short poem about cats"}, config)
graph.invoke({"messages": "now do the same but for dogs"}, config)
graph.invoke({"messages": "what's my name?"}, config)

SummarizationResult dataclass

消息总结的结果。

属性

  • messages (list[AnyMessage]) –

    总结后准备输入LLM的更新消息列表,包括带有总结的消息(如果有)。

  • running_summary (RunningSummary | None) –

    关于先前总结的信息(总结文本和先前总结的消息ID)。

messages instance-attribute

messages: list[AnyMessage]

总结后准备输入LLM的更新消息列表,包括带有总结的消息(如果有)。

running_summary class-attribute instance-attribute

running_summary: RunningSummary | None = None

关于先前总结的信息(总结文本和先前总结的消息ID)。如果未执行总结(没有足够的消息进行总结),则可以为None。

RunningSummary dataclass

用于存储先前总结信息的对象。

在后续调用 `summarize_messages` 时使用,以避免总结相同的消息。

属性

summary instance-attribute

summary: str

消息的最新总结,每次执行总结时更新。

summarized_message_ids instance-attribute

summarized_message_ids: set[str]

先前已总结的所有消息的ID。

last_summarized_message_id instance-attribute

last_summarized_message_id: str | None

最后一条被总结消息的ID。

summarize_messages

summarize_messages(
    messages: list[AnyMessage],
    *,
    running_summary: RunningSummary | None,
    model: LanguageModelLike,
    max_tokens: int,
    max_tokens_before_summary: int | None = None,
    max_summary_tokens: int = 256,
    token_counter: TokenCounter = count_tokens_approximately,
    initial_summary_prompt: ChatPromptTemplate = DEFAULT_INITIAL_SUMMARY_PROMPT,
    existing_summary_prompt: ChatPromptTemplate = DEFAULT_EXISTING_SUMMARY_PROMPT,
    final_prompt: ChatPromptTemplate = DEFAULT_FINAL_SUMMARY_PROMPT,
) -> SummarizationResult

当消息超出令牌限制时,对其进行总结,并用一个总结消息替换它们。

此函数从最旧到最新处理消息:一旦累积消息令牌数达到 max_tokens_before_summary,则 max_tokens_before_summary 内的所有消息(如果存在系统消息则排除)都将被总结,并替换为一个新的总结消息。最终的消息列表是 [总结消息] + 剩余消息。

参数

  • messages (list[AnyMessage]) –

    要处理的消息列表。

  • running_summary (RunningSummary | None) –

    可选的运行中总结对象,包含先前总结的信息。如果提供: - 将仅处理先前总结的消息 - 如果没有生成新的总结,则运行中总结将添加到返回的消息中 - 如果需要生成新的总结,则通过合并运行中总结中的现有总结值来生成

  • model (LanguageModelLike) –

    用于生成总结的语言模型。

  • max_tokens (int) –

    最终输出中返回的最大令牌数。仅在总结后强制执行。此值也将用作提供给总结LLM的最大令牌数。

  • max_tokens_before_summary (int | None, 默认值:None ) –

    在触发总结之前要累积的最大令牌数。如果未提供,则默认为 max_tokens 的相同值。这允许在需要时将更多令牌放入总结LLM中。

    注意

    如果 max_tokens_before_summary 内的最后一条消息是带有工具调用的AI消息,那么所有后续的、相应的工具消息也将被总结。

    注意

    如果要总结的令牌数大于 max_tokens,则只总结其中最后 max_tokens 个。这样做是为了防止超出总结LLM的上下文窗口(假定上限为 max_tokens)。

  • max_summary_tokens (int, 默认值:256 ) –

    为总结预留的最大令牌数。

    注意

    此参数不会传递给生成总结的LLM以限制总结的长度。它仅用于正确估算最大允许令牌预算。如果需要强制执行,则需要将 model.bind(max_tokens=max_summary_tokens) 作为 model 参数传递给此函数。

  • token_counter (TokenCounter, 默认值:count_tokens_approximately ) –

    用于计算消息中令牌的函数。默认为近似计数。要获得更准确的计数,可以使用 model.get_num_tokens_from_messages

  • initial_summary_prompt (ChatPromptTemplate, 默认值:DEFAULT_INITIAL_SUMMARY_PROMPT ) –

    用于生成第一个总结的提示模板。

  • existing_summary_prompt (ChatPromptTemplate, 默认值:DEFAULT_EXISTING_SUMMARY_PROMPT ) –

    用于更新现有(运行中)总结的提示模板。

  • final_prompt (ChatPromptTemplate, 默认值:DEFAULT_FINAL_SUMMARY_PROMPT ) –

    在返回之前将总结与剩余消息结合起来的提示模板。

返回

  • SummarizationResult

    一个SummarizationResult对象,包含更新后的消息和运行中的总结。 - messages:准备输入到LLM的更新消息列表 - running_summary:RunningSummary对象 - summary:最新总结的文本 - summarized_message_ids:先前已总结的消息ID集合 - last_summarized_message_id:最后一条被总结消息的ID

示例
from langgraph.graph import StateGraph, START, MessagesState
from langgraph.checkpoint.memory import InMemorySaver
from langmem.short_term import summarize_messages, RunningSummary
from langchain_openai import ChatOpenAI

model = ChatOpenAI(model="gpt-4o")
summarization_model = model.bind(max_tokens=128)


class SummaryState(MessagesState):
    summary: RunningSummary | None


def call_model(state):
    summarization_result = summarize_messages(
        state["messages"],
        running_summary=state.get("summary"),
        model=summarization_model,
        max_tokens=256,
        max_tokens_before_summary=256,
        max_summary_tokens=128,
    )
    response = model.invoke(summarization_result.messages)
    state_update = {"messages": [response]}
    if summarization_result.running_summary:
        state_update["summary"] = summarization_result.running_summary
    return state_update


checkpointer = InMemorySaver()
workflow = StateGraph(SummaryState)
workflow.add_node(call_model)
workflow.add_edge(START, "call_model")
graph = workflow.compile(checkpointer=checkpointer)

config = {"configurable": {"thread_id": "1"}}
graph.invoke({"messages": "hi, my name is bob"}, config)
graph.invoke({"messages": "write a short poem about cats"}, config)
graph.invoke({"messages": "now do the same but for dogs"}, config)
graph.invoke({"messages": "what's my name?"}, config)

评论