Langchain之让LLM拥有记忆

langchain的Memory

如果AI有记忆，我们就不需要手动维护一个储存消息历史的列表

让LLM拥有记忆的方法有很多，我更喜欢使用的方法是以下方案，其优点是灵活度比较高

from langchain.memory import ConversationBufferMemory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI
from langchain_core._api.deprecation import LangChainDeprecationWarning
import os
import dotenv
import warningswarnings.filterwarnings('ignore', category=LangChainDeprecationWarning)dotenv.load_dotenv()
api_key = os.environ.get('API_KEY')llm = ChatOpenAI(model='deepseek-chat',base_url='https://api.deepseek.com/',openai_api_key=api_key
)memory = ConversationBufferMemory(memory_key='history' ,return_messages=True)
prompts = ChatPromptTemplate.from_messages([('system','你是一个ai旅伴，现在需要你回答用户的问题，请不要输出无关的内容'),MessagesPlaceholder(variable_name='history'),('human','{text}')]
)
chain = (RunnablePassthrough.assign(history = lambda x : memory.load_memory_variables({})['history'])| prompts| llm| StrOutputParser()
)flag = Truewhile flag:msg = input('请输入:')d_end = ['再见', '结束', 'end', 'bye', 'quit', 'Bye']for i in d_end:if i in msg:flag = Falseres = chain.invoke({'text':msg})print(res)memory.save_context({'input':msg}, {'output':res})

看起来比较晦涩难懂，不过没关系，我们只需要将最重要的步骤拆机出来：

• 定义记忆

memory = ConversationBufferMemory(memory_key='history' ,return_messages=True) #注意此处的使用方法已经被新版本遗弃，需要忽略掉这个警告

• 定义提示词

prompts = ChatPromptTemplate.from_messages([('system','你是一个ai旅伴，现在需要你回答用户的问题，请不要输出无关的内容'),MessagesPlaceholder(variable_name='history'),('human','{text}')]
)

• 定义链

chain = (RunnablePassthrough.assign(history = lambda x : memory.load_memory_variables({})['history'])| prompts| llm| StrOutputParser()
)

• 保存对话

memory.save_context({'input':msg}, {'output':res})

对于输出的地方，我们要想达到流式输出的效果只需要将管道中的StrOutputParser()去掉，循环遍历res,输出遍历后的chunk.content，然后把print()中的flush参数设置为True

while True:msg = input('请输入:')if '再见'in msg or 'quit'in msg or 'Bye'in msg or 'bye' in msg:print('再见啦，希望我们的下次见面')break# print(mchat(msg=msg))res = chain.stream({'text':msg})full_str = ''for chunk in res:print(chunk.content, end='', flush=True)full_str += chunk.contentprint('\n')memory.save_context({'input':msg},{'output':full_str})