基于pypdf和chromadb构建向量库

前面在langgraph构建智能RAG中加载文档和分块入库使用的均是langchain自带的组件，本文使用纯粹的python库实现同样的功能。本文将使用pypdf和chromadb库。

1.安装库

#pip install pypdf

#pip install chromadb

2.加载pdf文档

使用pypdf加载文件代码如下：

from pypdf import PdfReader

使用《beijing_annual_report_2024》作为示例文件

reader = PdfReader ("../data/beijing_annual_report_2024.pdf")

从该文件中提取文本，并删除前后空格。加载后一页为一项

pdf_texts= $p.extract_text().strip() for p in reader.pages$

剔除空页，因为检索系统不能接受空页

pdf_texts= $text for text in pdf_texts if text$

输出加载后内容：

打印文档第一页作为输出示例

print(pdf_texts $0$ )

输出内容如下：

各位代表：

现在，......支持雄安新区"三校一院"交钥匙项目开学开诊，公共服务共建共享不断深化。

3.对文本分块

仍使用递归文本分块器对文本分块，具体代码如下：

from langchain_text_splitters import RecursiveCharacterTextSplitter

character_splitter = RecursiveCharacterTextSplitter(

#separators= $"\\n\\n", "\\n", ". ", " ", ""$ ,英文文档的分割符示例

separators= $"\\n\\n", "\\n", "。", "，", "、", "；", ""$ ,

chunk_size=1000,

chunk_overlap=100

)

#把所有文本拼接后再分割

character_split_texts = character_splitter.split_text('\n\n'.join(pdf_texts))

查看分段效果：

['各位......健全重大国',

'升，......中关村科技园挂牌运营'，

'......北京向天津、河']

4.创建向量库

chromadb与mongodb类似，使用集合表示一个表，具体代码如下：

import chromadb

from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction

embedding_function = SentenceTransformerEmbeddingFunction(model_name='../models/text2vec-base-chinese')

chroma_client = chromadb.PersistentClient(path="../data/chroma")

chroma_collection = chroma_client.get_or_create_collection("beijing_annual_report_2024", embedding_function=embedding_function)

使用持久化存储，刚创建时为空。

5.分块数据入库

把第3步处理后的所有分块向量化并插入向量库以备检索，具体代码如下：

ids = $str(i) for i in range(len(character_split_texts))$

chroma_collection.add(ids=ids, documents=character_split_texts)

6.相似性检索

使用集合的query方法进行检索，返回2个匹配的分块，具体代码如下：

query = "促进就业的举措有哪些？"

查询Chorma来获取结果，请求5个结果

results = chroma_collection.query(query_texts= $query$ , n_results=2)

查看检索结果：

results

输出如下：

{'ids': $\['8', '36'$ ],

'embeddings': None,

'documents': [['四是加大民生服务保障力度，......深入推进',

'促进 12 万城镇就业困难人员实现就业、......。']],

'uris': None,

'included': $'metadatas', 'documents', 'distances'$ ,

'data': None,

'metadatas': $\[None, None$ ],

'distances': $\[0.4134409427642822, 0.41738224029541016$ ]}

7.工具封装

把第6步的检索功能封装到工具中，工具返回整合后的内容和原始内容供智能体调用，具体代码如下：

from langchain.tools import tool

@tool(response_format="content_and_artifact")#返回原始内容和整合后内容

def retrieve_context(query: str):

"""Retrieve information to help answer a query."""

retrieved_docs = chroma_collection.query(query_texts= $query$ , n_results=2)

serialized = "\n\n".join(

(f"Source: {metadata}\nContent: {content}")

for metadata, content in zip(retrieved_docs $'metadatas'$ $0$ , retrieved_docs $'documents'$ $0$ )

)

#serialized是整合后内容，retrieved_docs为原始内容，原始内容包含元数据

return serialized, retrieved_docs

测试一下工具调用结果：

retrieve_context.invoke("促进就业的举措有哪些？")

输出如下：

#以下格式是输入到大模型中提供上线文的标准格式

'Source: None\nContent: 四是加大民生服务保障力度，......市深入推进\n\nSource: None\nContent: 少于 26 万人......体育赛事。'

8.分析

基于pypdf加载文件后元数据内容不够丰富，仅有公共的元数据，并且没有文件路径，每页没有元数据。

查看元数据：

reader.metadata

输出如下，输出中没有文件名

{'/Author': 'sky',

'/Comments': '',

'/Company': '',

'/CreationDate': "D:20240208203412+12'34'",

'/Creator': 'WPS 文字',

'/Keywords': '',

'/ModDate': "D:20240208203412+12'34'",

'/Producer': '',

'/SourceModified': "D:20240208203412+12'34'",

'/Subject': '',

'/Title': '',

'/Trapped': '/False'}

查看某页内容：

reader.pages $3$

输出如下：

{'/Contents': {'/Filter': '/FlateDecode'},

'/MediaBox': $0, 0, 595.3, 841.9$ ,

'/Parent': {'/Count': 13,

'/Kids': [IndirectObject(6, 0, 140335219092880),

IndirectObject(19, 0, 140335219092880),

IndirectObject(21, 0, 140335219092880),

IndirectObject(23, 0, 140335219092880),

IndirectObject(25, 0, 140335219092880),

IndirectObject(27, 0, 140335219092880),

IndirectObject(29, 0, 140335219092880),

IndirectObject(31, 0, 140335219092880),

IndirectObject(33, 0, 140335219092880),

IndirectObject(35, 0, 140335219092880),

IndirectObject(37, 0, 140335219092880),

IndirectObject(39, 0, 140335219092880),

IndirectObject(41, 0, 140335219092880)],

'/Type': '/Pages'},

'/Resources': {'/ExtGState': {'/GS13': {'/AIS': False,

'/BM': '/Normal',

'/CA': 1,

'/Type': '/ExtGState',

'/ca': 1}},

'/Font': {'/FT14': {'/BaseFont': '/ZLXSHE+Calibri',

'/DescendantFonts': $IndirectObject(16, 0, 140335219092880)$ ,

'/Encoding': '/Identity-H',

'/Subtype': '/Type0',

'/ToUnicode': {'/Filter': '/FlateDecode'},

'/Type': '/Font'},

'/FT8': {'/BaseFont': '/LNUHNF+SimSun',

'/DescendantFonts': $IndirectObject(10, 0, 140335219092880)$ ,

'/Encoding': '/Identity-H',

'/Subtype': '/Type0',

'/ToUnicode': {'/Filter': '/FlateDecode'},

'/Type': '/Font'}}},

'/Type': '/Page'}

基于此提供的上下文内容不够丰富，如果不能满足要求，则不能使用pypdf，仍需使用langchain社区中的文件加载器。