如何在LangChain中结合PGVector使用VectorStoreRetrieverMemory？

前置准备 #

• 确保PostgreSQL数据库已安装并启用pgvector扩展，执行SQL命令：

  CREATE EXTENSION IF NOT EXISTS vector;
  

• 安装Python依赖包：

  pip install langchain langchain-community langchain-openai psycopg2-binary pgvector
  

  （注：若使用开源Embedding模型，替换langchain-openai为对应包，如sentence-transformers）

1. 初始化PGVector向量存储 #

from langchain_community.vectorstores import PGVector
from langchain_openai import OpenAIEmbeddings

# 数据库连接字符串，替换为你的实际配置
CONNECTION_STRING = "postgresql+psycopg2://username:password@localhost:5432/your_db"

# 初始化Embedding模型（这里用OpenAI，也可替换为HuggingFaceEmbeddings等）
embeddings = OpenAIEmbeddings(api_key="your_openai_api_key")

# 创建PGVector存储实例
vector_store = PGVector(
    collection_name="chat_memory",  # 自定义集合名称
    connection_string=CONNECTION_STRING,
    embedding_function=embeddings
)

2. 构建VectorStoreRetrieverMemory #

from langchain.memory import VectorStoreRetrieverMemory

# 创建检索器，设置每次返回的相关记忆条数
retriever = vector_store.as_retriever(search_kwargs={"k": 3})

# 初始化向量存储检索记忆
memory = VectorStoreRetrieverMemory(retriever=retriever)

3. 在对话链中使用Memory #

from langchain.chains import LLMChain
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI

# 定义对话Prompt模板
prompt = PromptTemplate(
    input_variables=["history", "input"],
    template="""根据以下历史对话信息回答当前问题：
历史对话：{history}
当前问题：{input}
回答："""
)

# 初始化LLM
llm = ChatOpenAI(temperature=0, model="gpt-3.5-turbo")

# 创建带记忆的对话链
chain = LLMChain(llm=llm, prompt=prompt, memory=memory)

# 示例：执行对话并保存记忆
chain.invoke({"input": "我喜欢吃草莓冰淇淋"})
chain.invoke({"input": "我之前提到过喜欢什么口味的冰淇淋？"})

关键注意事项 #

• 若使用开源Embedding模型（如sentence-transformers），需确保输出向量维度与PGVector存储的维度一致；
• 可通过调整retriever.search_kwargs（如filter参数）实现更精准的记忆检索；
• 若需手动管理记忆，可调用memory.save_context({"input": "用户输入"}, {"output": "AI回答"})直接写入。

内容的提问来源于stack exchange，提问作者Andrea Bency