LangChain-13 Memory 通过 ConversationBufferMemory 记录会话的上下文 并在对话中自动检索

问题背景

平常我们的任务都是一次性会话，大模型是不会记录你的上下文的。如果我们要进行持续会话，或者希望大模型有记忆的话，我们需要对上下文进行管理。但这个过程是比较复杂的，LangChain 提供了一系列的工具来帮助我们简化过程。

安装依赖

pip install --upgrade --quiet  langchain-core langchain langchain-openai

编写代码

from operator import itemgetter

from langchain.memory import ConversationBufferMemory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from langchain_openai import ChatOpenAI

model = ChatOpenAI()
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "You are a helpful chatbot"),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{input}"),
    ]
)

memory = ConversationBufferMemory(return_messages=True)
memory.load_memory_variables({})

chain = (
    RunnablePassthrough.assign(
        history=RunnableLambda(memory.load_memory_variables) | itemgetter("history")
    )
    | prompt
    | model
)

inputs1 = {"input": "hi im wzk"}
message1 = chain.invoke(inputs1)
print(f"message1: {message1}")

# no memory
inputs2 = {"input": "whats my name?"}
message2 = chain.invoke(inputs2)
print(f"message2: {message2}")

# save message1 memory
memory.save_context(inputs1, {"output": message1.content})

# 加载 memory
memory.load_memory_variables({})
inputs3 = {"input": "whats my name"}
message3 = chain.invoke(inputs3)
print(f"message3: {message3}")

运行代码

➜ python3 test13.py
message1: content='Hello Wzk! How can I assist you today?' response_metadata={'finish_reason': 'stop', 'logprobs': None}
message2: content="I'm sorry, but I don't know your name. You never told me." response_metadata={'finish_reason': 'stop', 'logprobs': None}
message3: content='Your name is "wzk".' response_metadata={'finish_reason': 'stop', 'logprobs': None}

从会话中我们可以看到，第一次会话大模型是不知道的我们的名字的，但是在代码中，我们通过Memory的方式保留了记忆，此时我们再询问大模型，大模型就知道了我们的名字，并且正确的回答了。