利用 PostgreSQL 构建 RAG 系统实现智能问答

在现代信息检索和自然语言处理的场景中，检索增强生成 (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, featrueFROM knowledge.vector_data_1ORDER 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_1insert_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, featrueFROM knowledge.vector_data_1ORDER 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系统提供参考！