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 warnings

warnings.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 = True

while flag:
    msg = input('请输入:')
    d_end = ['再见', '结束', 'end', 'bye', 'quit', 'Bye']
    for i in d_end:
        if i in msg:
            flag = False
    res = 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.content
    print('\n')
    memory.save_context({'input':msg},{'output':full_str})