一架梯子,一头程序猿,仰望星空!
LangChain教程(Python版本) > 内容正文

LangChain 基于向量存储的记忆(memory)


基于向量数据库的记忆组件

VectorStoreRetrieverMemory将记忆存储在向量数据库中,每次查询历史记忆数据的时候返回前K个相似的内容。

在这里,“文档”指的是以前的对话消息片段。

# 导入基础包
from datetime import datetime
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.llms import OpenAI
from langchain.memory import VectorStoreRetrieverMemory
from langchain.chains import ConversationChain
from langchain.prompts import PromptTemplate

初始化向量数据库

使用不同的向量数据库,初始化步骤有所差异,这里以faiss为例

import faiss

from langchain.docstore import InMemoryDocstore
from langchain.vectorstores import FAISS

embedding_size = 1536 # 向量维度
#定义索引
index = faiss.IndexFlatL2(embedding_size)
# 这里使用openai的嵌入模型计算向量
embedding_fn = OpenAIEmbeddings().embed_query
#定义向量数据库
vectorstore = FAISS(embedding_fn, index, InMemoryDocstore({}), {})

创建VectorStoreRetrieverMemory

通过向量数据库创建VectorStoreRetrieverMemory组件

# 设置向量查询仅返回1个最相似的历史消息
retriever = vectorstore.as_retriever(search_kwargs=dict(k=1))
# 通过向量数据库创建VectorStoreRetrieverMemory记忆组件
memory = VectorStoreRetrieverMemory(retriever=retriever)

# 添加几条历史对话消息,用于测试
memory.save_context({"input": "我最喜欢的食物是披萨"}, {"output": "很高兴知道"})
memory.save_context({"input": "我最喜欢的运动是足球"}, {"output": "..."})
memory.save_context({"input": "我不喜欢凯尔特人"}, {"output": "ok"}) #
# 测试一下,根据提示词内容查询相关对话内容
print(memory.load_memory_variables({"prompt": "我应该看什么运动?"})["history"])
# 返回结果
input: 我最喜欢的运动是足球
output: ...

通过链(Chain)使用记忆组件

下面通过例子学习下如何使用记忆组件,可以设置“verbose=True” 参数,这样会答应跟模型交互的提示词内容。

llm = OpenAI(temperature=0) # 可以是任何有效的LLM,这里选openai
# 提示词模板
_DEFAULT_TEMPLATE = """以下是人类和AI之间友好的对话。AI很健谈,提供了许多来自其上下文的具体细节。如果AI不知道问题的答案,它会真诚地说不知道。

先前对话的相关部分:
{history}

(如果不相关,您无需使用这些信息)

当前对话:
Human: {input}
AI:"""
PROMPT = PromptTemplate(
    input_variables=["history", "input"],
    template=_DEFAULT_TEMPLATE
)
conversation_with_summary = ConversationChain(
    llm=llm,
    prompt=PROMPT,
    # 我们为测试目的设置了非常低的max_token_limit。
    memory=memory,
    verbose=True
)
conversation_with_summary.predict(input="嗨,我的名字叫佩里,有什么事吗?")
> Entering new ConversationChain chain...
    Prompt after formatting:
    以下是人类和AI之间友好的对话。AI很健谈,提供了许多来自其上下文的具体细节。如果AI不知道问题的答案,它会真诚地说不知道。

先前对话的相关部分:
输入:我最喜欢的食物是披萨。
输出:那很好知道。

(如果不相关,则不需要使用这些信息)

当前对话:
人类:嗨,我的名字是佩里,怎么了?
AI:

> Finished chain.
“嗨佩里,我做得很好。你呢?”
# langchain会先找到历史对话中跟自己喜欢运动相关的历史对话,放到背景信息里面
conversation_with_summary.predict(input="我的最爱运动是什么?")
> Entering new ConversationChain chain...
    Prompt after formatting:
    以下是人类和人工智能之间的友好对话。人工智能很健谈,并从其上下文中提供了许多具体细节。如果人工智能不知道问题的答案,它会如实地说它不知道。

先前对话的相关部分:

输入:我最喜欢的运动是足球

输出:。。。

(如果不相关,则无需使用这些信息)

当前对话:

人类:我最喜欢的运动是什么?

AI:

    > Finished chain.

    ' 你早些时候告诉我你最喜欢的运动是足球。'


关联主题