AGI｜基于LangChain实现的三种高级RAG检索方法

一、前言

RAG(Retrieval-Augmented Generation)检索增强生成，是现如今基于企业私域知识的问答应用所使用的主流技术之一。相较于重新训练基于私域知识的大模型来说，RAG没有额外的预训练成本，且回答效果与之相当。

但在实际应用场景中，RAG所面临最大的问题是LLM的上下文长度限制。企业私域知识文本的数量十分庞大，不可能将其全部放在模型的prompt中，即使现在各类模型已经将上下文token从年初的2k、4k扩充到了128k、192k，但是这可能也就是一份合同、一份标书的长度。因此，如何减少传递给模型的内容数量且同时提高内容质量，是提升基于RAG的AI应用回答准确度的一个重要方法。

本篇文章将基于LangChain实现三种高级检索方法，句子窗口检索和自动合并检索旨在改善RAG流程的召回过程中存在的信息残缺的问题，而多路召回检索则保证了在多个文档中检索召回的准确性。

二、先验知识

●RAG简要流程

加载文档——切分划片——嵌入为向量表示——存入数据库

向量化问题——向量召回文档——合并放入Prompt——LLM生成答案

三、句子窗口检索

（一）概念

在文档进行切片工作后，文档被分为若干个Langchain自定义的Document对象，该对象有两个属性，一是page_content即该切片的文本内容，二是meta_data即有关该切片的一些信息和可自定义封装进入的信息。

句子窗口检索方法，将每个切片的相邻切片的内容封装在切片的meta_data中。在检索和召回过程时，根据命中文档的meta_data可获得此段落的上下文信息，并将其封装进入命中文档的page_content中。组合完成的文档列表即可作为prompt交付给大模型生成。

在实际问答任务中，我们建议使用切片器将文档切分为较短的分片，或使用依据标点符号进行切分的切片器。保证整片文档拥有较细的颗粒度。同时在封装和召回阶段，适当扩大窗口大小，保证召回段落的完整性。

（二）BERT

（1）元数据封装

def metadata_format(self, ordered_text, **kwargs):
count = kwargs.get("split_count", 1)
fori, documentinenumerate(ordered_text):
ifi > 0:
document.metadata['previous_page'] = ordered_text[i-count].page_content
else:
document.metadata['previous_page'] = ''ifi < len(ordered_text) - 1:
document.metadata['next_page'] = ordered_text[i+count].page_content
else:
document.metadata['next_page'] = ''
returnordered_text

（2）数据重构

def search_and_format(self, databases, query, **kwargs):
top_documents = []
fordb in databases:
top_documents.append(db.similarity_search_with_score(query))
docs = []
fordoc, _ in top_documents:
doc.page_content = doc.metadata.get("previous_page") + doc.page_content + doc.metadata.get("next_page")
docs.append(doc)
returndocs

（3）调用示例伪代码

#load document
......#split
......
#use smartvision sdk to format
sentence_window_retrival = SentenceWindow()
formatted_documents = sentence_window_retrival.metadata_format(documents, split_count=2)#embedding 
......#load inlocalvector db
......#use smartvision sdk to dosearch and multiple recall
databases = [db]
query = "烟草专卖品的运输"
top_documents = sentence_window_retrival.search_and_format(databases, query)
print(top_documents)

四、自动合并检索

（一）概念

自动合并检索方法，实现方法源自Llamaindex所封装的自动合并检索，但RAG全流程需要制定一套准确的规范，因此在用户文档完成读取和切片工作后，所得到的Langchain格式的Document对象需转化为Llamaindex定义的Document对象，便可通过Llamaindex的自定义算法自动划分整个切片列表的子节点和父节点，最后鉴于规范再重新转化为Langchain格式的Document对象，并将父节点信息、深度信息等封装进每个节点。

在检索阶段，召回最相关的若干个节点，遍历这些节点和附加信息，如若超过K个节点同时属于同一个节点（这里的K为用户自定义阈值，通常为一个节点所有子节点的半数）则执行合并该父节点下属所有子节点，即返回整个父节点内容。这使我们能够将可能不同的较小上下文合并到一个可能有助于综合的更大上下文中。

（二）代码实现和调用

（1）元数据封装

defauto_merge_format(documents, **kwargs):
ifdocuments isNone:
raiseValueError('documents is required')
formatted_documents = []
doc_text = "\n\n".join([d.page_content ford indocuments])
docs = [Document(text=doc_text)]
node_parser = HierarchicalNodeParser.from_defaults(chunk_sizes=kwargs.get("pc_chunk_size", [2048, 512, 128]),chunk_overlap=kwargs.get("pc_chunk_overlap", 10))
nodes = node_parser.get_nodes_from_documents(docs)
leaf_nodes = get_leaf_nodes(nodes)
root_nodes = get_root_nodes(nodes)
middle_nodes = get_middle_node(nodes, leaf_nodes, root_nodes)
root_context_dict = {}
forroot_node innodes:
root_context_dict[root_node.node_id] = root_node.get_content()fornode innodes:
ifnode.parent_node:
node_id = node.node_id
root_node_id = node.parent_node.node_id
root_node_content = root_context_dict.get(node.parent_node.node_id)
root_node_child_count = 0
forparent_node inroot_nodes + middle_nodes:
ifparent_node.node_id == node.parent_node.node_id:
root_node_child_count = len(parent_node.child_nodes)
break
depth = 2ifnode inmiddle_nodes else3
child_count = len(node.child_nodes) ifnode.child_nodes isnotNoneelse0
document = langchain.schema.Document(page_content=node.get_content(),metadata={"node_id": node_id, "root_node_id": root_node_id, "root_node_content": root_node_content, "root_node_child_count": root_node_child_count, "depth": depth, "child_count": child_count})
formatted_documents.append(document)
returnformatted_documents

（2）数据重构

defsearch_and_format(self, databases, query, **kwargs):
top_documents = []
fordb indatabases:
top_document = db.similarity_search_with_score(query)
top_documents.append(top_document)
leaf_nodes = [doc fordoc, _ intop_documents]
returndo_merge(leaf_nodes, **kwargs)defgroup_nodes_by_depth(nodes, depth):
return[node fornode innodes ifnode.metadata.get("depth") == depth]defprocess_group(nodes, threshold):
grouped_by_root_id = {}
fornode innodes:
root_id = node.metadata.get("root_node_id")
grouped_by_root_id.setdefault(root_id, []).append(node)merge_context = []
forgroup ingrouped_by_root_id.values():
node_count = len(group)
child_count = group[0].metadata.get("root_node_child_count")
ifnode_count / child_count >= threshold:
merge_context.append(langchain.schema.Document(
page_content=group[0].metadata.get("root_node_content")
))
else:
fordocument ingroup:
merge_context.append(document)
returnmerge_contextdefdo_merge(nodes, **kwargs)-> List[langchain.schema.Document]:
threshold = kwargs.get("threshold", 0.5)
leaf_nodes = group_nodes_by_depth(nodes, 3)
middle_nodes = group_nodes_by_depth(nodes, 2)
leaf_merge_context = process_group(leaf_nodes, threshold)
middle_merge_context = process_group(middle_nodes, threshold)
merge_content = leaf_merge_context + middle_merge_context
returnmerge_contentdefget_middle_node(nodes, leaf_nodes, root_nodes):
middle_node = []
fornode innodes:
ifnode notinleaf_nodes andnode notinroot_nodes:
middle_node.append(node)
returnmiddle_node

（3）调用示例伪代码

#load document
......#split
......#use smartvision sdk to format
auto_merge_retrival = AutoMergeRetrieval()
formatted_documents = auto_merge_retrival.metadata_format(documents,
pc_chunk_size=[1024, 128, 32],
pc_chunk_overlap=4)
#embedding 
......#load inlocalvector db
......#use smartvision sdk to dosearch and multiple recall
top_documents = auto_merge_retrival.search_and_format(databases, query, threshold=0.5)
print(top_documents)

五、多路召回检索

（一）概念

多路召回检索方法，在元数据封装环节并未做任何操作，而在检索阶段他允许用户上传多个数据集或不同类型的向量数据库作为检索对象，以适应用户私域知识库文档类型不同，文档数量庞大的问题。从多个数据源检索得到文档列表，而后通过rerank模型对文档与问题的相关性进行评分，筛选出大于一定分值的文档，组合成为prompt。

由此可见，多路召回检索在数据源广而杂的情况下，富有更好的效果。此外，rerank模型虽能进行再次的重排以提高准确性，但是在牺牲速度和效率的前提下进行的，因此需充分考虑这个问题。

（二）代码实现

（1）元数据封装

defmetadata_format(self, ordered_text, **kwargs):
"""
默认rag，不做任何处理
"""
returnordered_text

（2）数据重构

defsearch_and_format(self, databases, query, **kwargs):
top_documents = []
result_data = []
fordb indatabases:
top_document = db.similarity_search_with_score(query)
top_documents.append(top_document)
pairs = [[query, item.page_content] foritem intop_documents]
withtorch.no_grad():
rerank_tokenizer = AutoTokenizer.from_pretrained(RERANK_FILE_PATH)
inputs = rerank_tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
rerank_model = AutoModelForSequenceClassification.from_pretrained(RERANK_FILE_PATH)
scores = rerank_model(**inputs, return_dict=True).logits.view(-1, ).float()
fori, score inenumerate(scores):
data = {
"text": top_documents[i].page_content,
"score": float(score)
}
result_data.append(data)
returnresult_data