在现代信息检索和自然语言处理的场景中，检索增强生成 (Retrieval-Augmented Generation, RAG) 系统因其结合了知识库检索和生成模型的优势，成为了一种非常流行的智能问答方法。在这篇博文中，我将展示如何利用PostgreSQL作为向量存储数据库，配合OpenAI嵌入模型和LangChain库，构建一个完整的RAG系统。
RAG 系统简介
RAG 系统的核心理念是：首先从知识库中检索与问题相关的文档或片段，然后通过生成式语言模型（如GPT）生成基于检索结果的答案。这种方法不仅提升了模型的问答准确性，还能够在多种场景中扩展大语言模型的应用。
在本文中，我们将使用：
- LangChain：一个为构建语言模型应用提供丰富工具的框架。
- PostgreSQL：作为存储文本片段及其嵌入向量的数据库。
- OpenAI API：为文档生成嵌入向量并使用 GPT 模型生成答案。
主要实现步骤
1. 环境准备与库的导入
我们首先需要导入必要的库，其中包括用于数据库连接的psycopg2，用于加载网页的bs4，以及LangChain相关的库。
import getpass
import os
import psycopg2
from psycopg2.extras import execute_values
import bs4
from langchain_community.document_loaders import WebBaseLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
import numpy as np
from langchain_openai import ChatOpenAI
from langchain import hub
from langchain_core.output_parsers import StrOutputParser
2. 加载网页内容
我们将网页内容加载为文本，并使用LangChain提供的WebBaseLoader来解析网页，并提取需要的内容。
# 加载网页内容
loader = WebBaseLoader(
web_paths=("https://lilianweng.github.io/posts/2023-06-23-agent/",),
bs_kwargs=dict(
parse_only=bs4.SoupStrainer(class_=("post-content", "post-title", "post-header"))
),
)
docs = loader.load()
3. 文本切分
由于大语言模型处理较长文本时会受到限制，因此我们需要对加载的文本进行切分。在此处，使用RecursiveCharacterTextSplitter按照指定的字符大小将文本切分为多个块。
# 文本切分
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
splits = text_splitter.split_documents(docs)
4. 连接 PostgreSQL 并存储嵌入向量
在这里，我们将利用PostgreSQL作为数据库，存储文本的嵌入向量。我们首先连接数据库，然后利用OpenAI的嵌入模型生成文本向量，并存储到数据库中。
# 连接到 PostgreSQL 数据库
conn = psycopg2.connect(
host="localhost", # 根据需要更改
database="mydb", # 更改为你的数据库名
user="root", # 更改为你的用户名
password=getpass.getpass("请输入你的 PostgreSQL 密码: ")
)
cur = conn.cursor()
# 嵌入并存储向量
embedding_model = OpenAIEmbeddings()
def store_vectors_in_pg(splits):
embeddings = embedding_model.embed_documents([doc.page_content for doc in splits])
data = [
(0, 0, doc.page_content, embedding)
for doc, embedding in zip(splits, embeddings)
]
insert_query = """
INSERT INTO knowledge.vector_data_1 (user_id, file_id, content, featrue)
VALUES %s
"""
execute_values(cur, insert_query, data)
conn.commit()
store_vectors_in_pg(splits)
5. 从数据库检索相似文档
通过查询PostgreSQL中的向量数据，基于余弦相似度查找与用户查询相似的文档。我们利用向量索引（HNSW）来高效检索相似文本片段。
# 检索相似文档
def retrieve_similar_docs(query, k=5):
query_embedding = embedding_model.embed_query(query)
embedding_str = f"'{str(query_embedding)}'"
retrieve_query = f"""
SELECT content, featrue
FROM knowledge.vector_data_1
ORDER BY featrue <-> {embedding_str}
LIMIT %s
"""
cur.execute(retrieve_query, (k,))
results = cur.fetchall()
return [result[0] for result in results]
6. 构建 RAG 链并生成答案
我们使用LangChain的ChatOpenAI模型和Prompt链来完成生成答案的过程。
# 定义 RAG 链
prompt = hub.pull("rlm/rag-prompt")
input_data = {
"context": format_docs(retrieved_docs),
"question": query
}
# 构建 RAG 链
rag_chain = (
prompt
| ChatOpenAI(model="gpt-4o-mini")
| StrOutputParser()
)
# 生成答案
response = rag_chain.invoke(input_data)
print(response)
7. 关闭数据库连接
在程序结束时，别忘了关闭数据库连接。
cur.close()
conn.close()
完整代码实例
# 导入必要的库
import getpass
import os
import psycopg2
from psycopg2.extras import execute_values
import bs4
from langchain_community.document_loaders import WebBaseLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
import numpy as np
from langchain_openai import ChatOpenAI
from langchain import hub
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
# 设置 OpenAI API 密钥
os.environ["OPENAI_API_KEY"] = getpass.getpass("请输入你的 OpenAI API 密钥: ")
# 1. 加载网页内容
loader = WebBaseLoader(
web_paths=("https://lilianweng.github.io/posts/2023-06-23-agent/",),
bs_kwargs=dict(
parse_only=bs4.SoupStrainer(class_=("post-content", "post-title", "post-header"))
),
)
docs = loader.load()
# 2. 切分文本
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
splits = text_splitter.split_documents(docs)
# 3. 连接到 PostgreSQL 数据库
conn = psycopg2.connect(
host="localhost", # 根据需要更改
database="mydb", # 更改为你的数据库名
user="root", # 更改为你的用户名
password=getpass.getpass("请输入你的 PostgreSQL 密码: ")
)
cur = conn.cursor()
# 4. 嵌入并将分割后的文本和向量数据存储到 PostgreSQL
embedding_model = OpenAIEmbeddings()
def store_vectors_in_pg(splits):
# 生成嵌入向量
embeddings = embedding_model.embed_documents([doc.page_content for doc in splits])
# 将 numpy.ndarray 转换为 Python 列表
data = [
(0, 0, doc.page_content, embedding) # 假设 user_id 和 file_id 为 0，实际可以调整
for doc, embedding in zip(splits, embeddings)
]
# 插入数据到 PostgreSQL 表 knowledge.vector_data_1
insert_query = """
INSERT INTO knowledge.vector_data_1 (user_id, file_id, content, featrue)
VALUES %s
"""
execute_values(cur, insert_query, data)
conn.commit()
store_vectors_in_pg(splits)
# 5. 从 PostgreSQL 中检索数据并计算相似度
def retrieve_similar_docs(query, k=5):
# 将查询嵌入为向量
query_embedding = embedding_model.embed_query(query)
# 将 query_embedding 转换为 PostgreSQL 可识别的向量字符串格式
embedding_str = f"'{str(query_embedding)}'" # 将list转为vector字符串
# SQL 查询：通过 HNSW 索引查找最相似的文档
retrieve_query = f"""
SELECT content, featrue
FROM knowledge.vector_data_1
ORDER BY featrue <-> {embedding_str} -- 使用转换后的字符串进行余弦相似度计算
LIMIT %s
"""
cur.execute(retrieve_query, (k,))
results = cur.fetchall()
return [result[0] for result in results] # 返回相似的内容
# 6. 检索并生成答案
query = "任务分解是什么？"
retrieved_docs = retrieve_similar_docs(query)
def format_docs(docs):
return "\n\n".join(docs)
# 定义 RAG 链
prompt = hub.pull("rlm/rag-prompt")
# 构建 dict 输入
input_data = {
"context": format_docs(retrieved_docs),
"question": query
}
# 构建 RAG 链
rag_chain = (
prompt # 提示模板
| ChatOpenAI(model="gpt-4o-mini") # 使用 OpenAI 的 LLM 模型生成答案
| StrOutputParser() # 将输出解析为字符串格式
)
# 生成答案
response = rag_chain.invoke(input_data)
print(response)
# 7. 关闭连接
cur.close()
conn.close()
结论
通过本文，我们展示了如何使用PostgreSQL作为向量存储的数据库，配合OpenAI嵌入模型及LangChain库构建一个简单的RAG系统。这个系统能够高效检索文本片段，并基于检索结果生成回答。RAG 系统在知识问答、信息检索等领域具有广泛的应用前景，尤其是在处理大量结构化或非结构化数据时，结合自然语言处理模型的强大生成能力，可以显著提升用户体验。
希望这篇文章能为你构建自己的RAG系统提供参考！
