如何使用LangChain PGVector()添加索引并在RAG链中通过索引检索

一、前置准备 #

• 安装依赖包：pip install langchain langchain-community psycopg2-binary pgvector
• 配置PostgreSQL数据库并启用pgvector扩展（需在数据库中执行SQL：CREATE EXTENSION IF NOT EXISTS vector;）
• 准备嵌入模型（如OpenAIEmbeddings、HuggingFaceEmbeddings等）

二、初始化PGVector实例 #

先创建PGVector连接实例，对应你提供的代码示例：

from langchain_community.vectorstores import PGVector
from langchain_openai import OpenAIEmbeddings

# 初始化嵌入模型，可替换为其他类型
embeddings = OpenAIEmbeddings(api_key="你的API密钥")
CONNECTION_STRING = "postgresql+psycopg2://用户名:密码@主机地址:端口/数据库名"
case_id = "你的案例集合名称"

# 建立PGVector连接
db = PGVector(embedding_function=embeddings, collection_name=case_id, connection_string=CONNECTION_STRING)

三、添加数据并构建索引 #

方式1：直接添加文本列表 #

texts = [
    "RAG是检索增强生成的缩写，用于提升大模型生成内容的准确性",
    "PGVector是PostgreSQL的向量扩展，用于存储和检索向量数据",
    "LangChain提供PGVector集成，简化向量数据库操作流程"
]

# 自动生成向量并构建索引
db.add_texts(texts)

方式2：添加带元数据的文档 #

from langchain_core.documents import Document

docs = [
    Document(
        page_content="RAG是检索增强生成的缩写，用于提升大模型生成内容的准确性",
        metadata={"source": "rag_guide.md", "category": "概念"}
    ),
    Document(
        page_content="PGVector是PostgreSQL的向量扩展，用于存储和检索向量数据",
        metadata={"source": "pgvector_docs.md", "category": "工具"}
    )
]

# 添加文档并生成索引
db.add_documents(docs)

四、通过索引进行检索 #

基础相似性检索 #

query = "什么是RAG？"
# 检索最相似的3条文档
retrieved_docs = db.similarity_search(query, k=3)

# 输出结果
for doc in retrieved_docs:
    print(f"内容: {doc.page_content}\n元数据: {doc.metadata}\n")

带相似度分数的检索 #

retrieved_docs_with_score = db.similarity_search_with_score(query, k=3)
for doc, score in retrieved_docs_with_score:
    print(f"内容: {doc.page_content}\n相似度分数: {score}\n元数据: {doc.metadata}\n")

带元数据过滤的检索 #

# 仅检索category为"概念"的文档
retrieved_docs_filtered = db.similarity_search(
    query,
    k=2,
    filter={"category": "概念"}
)

五、整合到RAG链中 #

将PGVector检索器与大模型结合构建完整RAG链：

from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

# 定义RAG提示模板
template = """基于以下上下文信息回答用户问题：
{context}

用户问题：{question}
"""
prompt = ChatPromptTemplate.from_template(template)

# 初始化大模型
llm = ChatOpenAI(api_key="你的API密钥", model="gpt-3.5-turbo")

# 构建RAG链
rag_chain = (
    {"context": db.as_retriever(k=3), "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

# 执行RAG查询
response = rag_chain.invoke("RAG如何提升大模型性能？")
print(response)

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