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})