高级RAG策略学习（四）——上下文窗口增强检索RAG

上下文窗口RAG系统：原理与代码实现详解

引言

检索增强生成（RAG）系统通过结合信息检索和生成模型，显著提升了问答系统的准确性和相关性。然而，传统的RAG系统在检索时往往只关注单个文档块，可能会丢失重要的上下文信息。本文将详细介绍一种基于上下文窗口的RAG系统实现，该系统通过检索相邻文档块来增强上下文信息，从而提供更准确和连贯的回答。

一、技术原理

1.1 传统RAG系统的局限性

传统的RAG系统工作流程如下：

1. 将长文档分割成固定大小的文档块
2. 对每个文档块进行向量化编码
3. 根据查询进行语义相似度检索
4. 将检索到的文档块作为上下文输入生成模型

这种方法的主要问题是：

• 上下文断裂：文档分块可能会切断重要的上下文关系
• 信息不完整：单个文档块可能无法提供足够的背景信息
• 逻辑不连贯：检索到的片段可能缺乏前后逻辑关系

1.2 上下文窗口RAG的解决方案

上下文窗口RAG系统通过以下机制解决上述问题：

1.2.1 索引化分块

• 为每个文档块分配唯一的索引号
• 保持文档块在原文中的顺序关系
• 在元数据中记录块的位置信息

1.2.2 邻居检索策略

• 首先进行标准的语义相似度检索
• 对于每个检索到的相关块，获取其前后N个相邻块
• 将相邻块按原始顺序重新组合

1.2.3 上下文重构

• 处理相邻块之间的重叠部分
• 生成包含完整上下文的文档序列
• 保持原文的逻辑连贯性

1.3 技术优势

1. 上下文完整性：通过检索相邻块，保持了文档的原始逻辑结构
2. 信息丰富度：提供更多的背景信息，有助于生成更准确的回答
3. 灵活配置：可以根据需要调整邻居块的数量和重叠度
4. 对比分析：同时提供基线和增强结果，便于效果评估

二、代码实现详解

2.1 环境配置与依赖

import os
from dotenv import load_dotenv
from langchain.docstore.document import Document
from helper_functions import *
from typing import List
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.vectorstores import FAISS

# 加载环境变量
load_dotenv()

# 配置通义千问API
DASHSCOPE_BASE_URL = os.getenv("DASHSCOPE_BASE_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1")
CHAT_MODEL_NAME = "qwen-plus"
EMBEDDING_MODEL_NAME = "text-embedding-v2"

技术要点：

• 使用通义千问的嵌入模型和聊天模型
• 通过环境变量管理API配置，提高安全性
• 支持中文文本处理，优化本地化体验

2.2 核心功能实现

2.2.1 文档分块与索引

def split_text_to_chunks_with_indices(text: str, chunk_size: int, chunk_overlap: int) -> List[Document]:
    """
    将文本分割为多个文档，每个文档包含指定大小的文本，同时记录每个文档的索引。
    """
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end]
        chunks.append(Document(
            page_content=chunk, 
            metadata={"index": len(chunks), "text": text}
        ))
        start += chunk_size - chunk_overlap
    return chunks

关键特性：

• 索引记录：为每个块分配唯一索引，便于后续检索
• 重叠处理：支持块间重叠，保持上下文连续性
• 元数据保存：在元数据中保存原始文本和索引信息

2.2.2 索引检索机制

def get_chunk_by_index(vectorstore, target_index: int) -> Document:
    """
    根据索引从向量存储中检索文档。
    """
    all_docs = vectorstore.similarity_search("", k=vectorstore.index.ntotal)
    for doc in all_docs:
        if doc.metadata.get('index') == target_index:
            return doc
    return None

实现原理：

• 通过遍历所有文档，匹配元数据中的索引值
• 提供精确的索引检索功能
• 支持快速定位特定位置的文档块

2.2.3 上下文增强检索（核心算法）

def retrieve_with_context_overlap(vectorstore, retriever, query: str, 
                                  num_neighbors: int = 1, chunk_size: int = 200,
                                  chunk_overlap: int = 20) -> List[str]:
    """
    从向量存储中检索文档，根据语义相似度和上下文重叠进行填充。
    """
    # 1. 语义检索
    relevant_chunks = retriever.get_relevant_documents(query)
    result_sequences = []

    for chunk in relevant_chunks:
        current_index = chunk.metadata.get('index')
        if current_index is None:
            continue

        # 2. 确定邻居范围
        start_index = max(0, current_index - num_neighbors)
        end_index = current_index + num_neighbors + 1

        # 3. 检索邻居块
        neighbor_chunks = []
        for i in range(start_index, end_index):
            neighbor_chunk = get_chunk_by_index(vectorstore, i)
            if neighbor_chunk:
                neighbor_chunks.append(neighbor_chunk)

        # 4. 排序和拼接
        neighbor_chunks.sort(key=lambda x: x.metadata.get('index', 0))
        
        # 5. 处理重叠拼接
        concatenated_text = neighbor_chunks[0].page_content
        for i in range(1, len(neighbor_chunks)):
            current_chunk = neighbor_chunks[i].page_content
            overlap_start = max(0, len(concatenated_text) - chunk_overlap)
            concatenated_text = concatenated_text[:overlap_start] + current_chunk

        result_sequences.append(concatenated_text)

    return result_sequences

算法流程：

1. 语义检索：基于查询进行初始的相似度检索
2. 邻居扩展：为每个相关块确定前后邻居的范围
3. 批量检索：获取指定范围内的所有邻居块
4. 顺序重排：按原始索引顺序排列邻居块
5. 智能拼接：处理重叠部分，生成连贯的文本序列

2.3 RAG系统封装

2.3.1 主类设计

class RAGMethod:
    """
    RAG方法的主类，封装了文档准备、检索器设置和查询执行。
    """
    def __init__(self, chunk_size: int = 400, chunk_overlap: int = 200):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.docs = self._prepare_docs()
        self.vectorstore, self.retriever = self._prepare_retriever()

设计模式：

• 封装性：将复杂的RAG流程封装在单一类中
• 可配置性：支持自定义块大小和重叠参数
• 模块化：分离文档准备和检索器设置逻辑

2.3.2 检索器准备

def _prepare_retriever(self):
    """
    准备检索器，使用文档和嵌入模型创建向量存储。
    """
    embeddings = DashScopeEmbeddings(
        model=EMBEDDING_MODEL_NAME,
        dashscope_api_key=os.getenv('DASHSCOPE_API_KEY')
    )
    vectorstore = FAISS.from_documents(self.docs, embeddings)
    retriever = vectorstore.as_retriever(search_kwargs={"k": 1})
    return vectorstore, retriever

技术选型：

• FAISS向量存储：高效的相似度搜索和聚类库
• 通义千问嵌入：针对中文优化的嵌入模型
• 灵活检索配置：可调整检索数量和搜索参数

2.4 命令行接口

def parse_args():
    import argparse
    parser = argparse.ArgumentParser(description="Run RAG method on a given PDF and query.")
    parser.add_argument("--query", type=str, default="深度学习何时在AI中变得突出？",
                        help="用于测试检索器的查询")
    parser.add_argument('--chunk_size', type=int, default=400, help="文本块的大小")
    parser.add_argument('--chunk_overlap', type=int, default=200, help="块之间的重叠")
    parser.add_argument('--num_neighbors', type=int, default=1, help="用于上下文的相邻块数量")
    return parser.parse_args()

用户体验：

• 参数化配置：支持命令行参数调整
• 中文界面：本地化的帮助信息
• 默认值设置：提供合理的默认参数

三、优化策略与最佳实践

3.1 窗口大小优化策略

根据不同应用场景选择合适的窗口参数：

3.2 性能优化建议

1. 元数据管理优化：确保窗口元数据不参与向量计算，仅在检索后处理阶段使用
2. 缓存机制：对频繁查询的邻居块进行缓存，减少重复计算
3. 批量处理：对多个查询进行批量检索，提高系统吞吐量

3.3 与其他技术的结合

• 重排序（Reranking）：先筛选相关节点，再补充上下文窗口
• 长上下文模型：配合GPT-4 Turbo等长上下文模型，扩大窗口容量
• 多粒度检索：根据查询复杂度动态调整窗口层级

四、技术路径选择指南

4.1 框架工具 vs 自定义实现

4.2 选择建议

• 选择框架工具：快速搭建RAG原型，标准化需求，团队技术栈统一
• 选择自定义实现：复杂业务逻辑，特殊检索规则，性能要求极高

五、实际应用与效果

5.1 应用场景

1. 长文档问答：适用于技术文档、学术论文等长文本的问答系统
2. 知识库检索：企业内部知识管理和信息检索
3. 教育辅助：在线学习平台的智能答疑系统
4. 法律文档分析：法律条文和案例的智能检索和分析

5.2 性能优势

通过上下文窗口增强，系统能够：

• 提高回答准确性：更完整的上下文信息
• 增强逻辑连贯性：保持原文的逻辑结构
• 减少信息丢失：避免重要信息被分块切断
• 提升用户体验：更自然和连贯的回答

5.3 配置建议

• chunk_size：建议300-500字符，平衡信息密度和处理效率
• chunk_overlap：建议为chunk_size的30-50%，确保上下文连续性
• num_neighbors：建议1-2个邻居，避免引入过多噪声信息

六、总结与展望

6.1 技术创新点

1. 索引化分块：为文档块建立有序索引，支持精确的位置检索
2. 邻居扩展策略：智能获取相邻上下文，增强信息完整性
3. 重叠处理算法：优雅处理块间重叠，保持文本连贯性
4. 对比评估框架：同时提供基线和增强结果，便于效果分析

6.2 未来发展方向

1. 动态窗口调整：根据查询复杂度自动调整邻居数量
2. 多模态支持：扩展到图像、表格等多模态内容
3. 实时优化：基于用户反馈动态优化检索策略
4. 分布式部署：支持大规模文档库的分布式检索

6.3 结语

上下文窗口RAG系统通过创新的邻居检索和上下文重构机制，有效解决了传统RAG系统的上下文断裂问题。该系统不仅保持了原有的语义检索能力，还显著提升了回答的完整性和连贯性。通过本文的详细介绍，读者可以深入理解上下文窗口RAG系统的原理和实现，并将其应用到实际的项目中，构建更智能和高效的问答系统。

7. 代码

import os
from dotenv import load_dotenv
from langchain.docstore.document import Document
from helper_functions import *
from typing import List
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.vectorstores import FAISS

# Load environment variables from a .env file
load_dotenv()

# 配置通义千问API
DASHSCOPE_BASE_URL = os.getenv("DASHSCOPE_BASE_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1")
CHAT_MODEL_NAME = "qwen-plus"
EMBEDDING_MODEL_NAME = "text-embedding-v2"


# Function to split text into chunks with metadata of the chunk chronological index
def split_text_to_chunks_with_indices(text: str, chunk_size: int, chunk_overlap: int) -> List[Document]:
    """
    将文本分割为多个文档，每个文档包含指定大小的文本，同时记录每个文档的索引。

    :param text: 要分割的文本
    :param chunk_size: 每个文档的最大字符数
    :param chunk_overlap: 文档之间的重叠字符数
    :return: 包含多个文档的列表，每个文档包含文本和索引元数据
    """
    chunks = []
    start = 0
    while start < len(text):
        """
        循环分割文本，每次提取一个文档的文本内容。
        文档的索引从0开始，每次增加1。
        文档的元数据包含原始文本和索引。
        """
        end = start + chunk_size
        chunk = text[start:end]
        chunks.append(Document(page_content=chunk, metadata={"index": len(chunks), "text": text}))
        start += chunk_size - chunk_overlap
        
    return chunks


# Function to retrieve a chunk from the vectorstore based on its index in the metadata
def get_chunk_by_index(vectorstore, target_index: int) -> Document:
    """
    根据索引从向量存储中检索文档。

    :param vectorstore: 向量存储对象
    :param target_index: 目标文档的索引
    :return: 检索到的文档对象
    """
    all_docs = vectorstore.similarity_search("", k=vectorstore.index.ntotal)
    for doc in all_docs:
        if doc.metadata.get('index') == target_index:
            return doc
    return None


# Function that retrieves from the vectorstore based on semantic similarity and pads each retrieved chunk with its neighboring chunks
def retrieve_with_context_overlap(vectorstore, retriever, query: str, num_neighbors: int = 1, chunk_size: int = 200,
                                  chunk_overlap: int = 20) -> List[str]:
    """
    从向量存储中检索文档，根据语义相似度和上下文重叠进行填充。

    :param vectorstore: 向量存储对象
    :param retriever: 检索器对象
    :param query: 检索查询
    :param num_neighbors: 要检索的邻居文档数量
    :param chunk_size: 每个文档的最大字符数
    :param chunk_overlap: 文档之间的重叠字符数
    :return: 包含填充后的文档序列的列表
    """
    relevant_chunks = retriever.get_relevant_documents(query)
    result_sequences = []

    for chunk in relevant_chunks:
        current_index = chunk.metadata.get('index')
        if current_index is None:
            continue

        # Determine the range of chunks to retrieve
        start_index = max(0, current_index - num_neighbors)
        end_index = current_index + num_neighbors + 1

        # Retrieve all chunks in the range
        neighbor_chunks = []
        for i in range(start_index, end_index):
            neighbor_chunk = get_chunk_by_index(vectorstore, i)
            if neighbor_chunk:
                neighbor_chunks.append(neighbor_chunk)

        # Sort chunks by their index to ensure correct order
        neighbor_chunks.sort(key=lambda x: x.metadata.get('index', 0))

        # Concatenate chunks, accounting for overlap
        concatenated_text = neighbor_chunks[0].page_content
        for i in range(1, len(neighbor_chunks)):
            current_chunk = neighbor_chunks[i].page_content
            overlap_start = max(0, len(concatenated_text) - chunk_overlap)
            concatenated_text = concatenated_text[:overlap_start] + current_chunk

        result_sequences.append(concatenated_text)

    return result_sequences


# Main class that encapsulates the RAG method
class RAGMethod:
    """
    RAG方法的主类，封装了文档准备、检索器设置和查询执行。
    """
    def __init__(self, chunk_size: int = 400, chunk_overlap: int = 200):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.docs = self._prepare_docs()
        self.vectorstore, self.retriever = self._prepare_retriever()

    def _prepare_docs(self) -> List[Document]:
        content = """
            人工智能（AI）的历史可以追溯到20世纪中叶。"人工智能"这一术语在1956年的达特茅斯会议上被提出，标志着该领域的正式开始。
            
            在20世纪50年代和60年代，AI研究专注于符号方法和问题解决。1955年由艾伦·纽厄尔和赫伯特·西蒙创建的逻辑理论家，通常被认为是第一个AI程序。
            
            20世纪60年代见证了专家系统的发展，这些系统使用预定义规则来解决复杂问题。1965年创建的DENDRAL是最早的专家系统之一，专门用于分析化学化合物。
            
            然而，20世纪70年代带来了第一个"AI寒冬"，这是一个AI研究资金减少和兴趣下降的时期，主要是由于过度承诺的能力和未能兑现的结果。
            
            20世纪80年代随着专家系统在企业中的普及而复苏。日本政府的第五代计算机项目也刺激了全球对AI研究的投资增加。
            
            神经网络在20世纪80年代和90年代获得了突出地位。反向传播算法虽然早期就被发现，但在这个时期被广泛用于训练多层网络。
            
            20世纪90年代末和2000年代标志着机器学习方法的兴起。支持向量机（SVM）和随机森林在各种分类和回归任务中变得流行。
            
            深度学习，一种使用多层神经网络的机器学习子集，在2010年代初开始显示出有希望的结果。突破出现在2012年，当时一个深度神经网络在ImageNet竞赛中显著超越了其他机器学习方法。
            
            从那时起，深度学习已经革命性地改变了许多AI应用，包括图像和语音识别、自然语言处理和游戏。2016年，谷歌的AlphaGo击败了世界冠军围棋选手，这是AI的一个里程碑式成就。
            
            当前的AI时代的特点是深度学习与其他AI技术的整合、更高效和强大硬件的发展，以及围绕AI部署的伦理考虑。
            
            2017年引入的Transformer已成为自然语言处理中的主导架构，使得像GPT（生成式预训练Transformer）这样的模型能够生成类似人类的文本。
            
            随着AI的不断发展，新的挑战和机遇不断涌现。可解释AI、鲁棒和公平的机器学习，以及通用人工智能（AGI）是该领域当前和未来研究的关键领域。
            """
        return split_text_to_chunks_with_indices(content, self.chunk_size, self.chunk_overlap)

    def _prepare_retriever(self):
        """
        准备检索器，使用文档和嵌入模型创建向量存储。

        :return: 向量存储对象和检索器对象
        """
        embeddings = DashScopeEmbeddings(
            model=EMBEDDING_MODEL_NAME,
            dashscope_api_key=os.getenv('DASHSCOPE_API_KEY')
        )
        vectorstore = FAISS.from_documents(self.docs, embeddings)
        retriever = vectorstore.as_retriever(search_kwargs={"k": 1})
        return vectorstore, retriever

    def run(self, query: str, num_neighbors: int = 1):
        """
        运行RAG方法，执行查询并返回基线块和增强块。

        :param query: 检索查询
        :param num_neighbors: 要检索的邻居文档数量
        :return: 基线块和增强块的元组
        """
        baseline_chunk = self.retriever.get_relevant_documents(query)
        enriched_chunks = retrieve_with_context_overlap(self.vectorstore, self.retriever, query, num_neighbors,
                                                        self.chunk_size, self.chunk_overlap)
        return baseline_chunk[0].page_content, enriched_chunks[0]


# Argument parsing function
def parse_args():
    import argparse
    parser = argparse.ArgumentParser(description="Run RAG method on a given PDF and query.")
    parser.add_argument("--query", type=str, default="深度学习何时在AI中变得突出？",
                        help="用于测试检索器的查询（默认：'深度学习何时在AI中变得突出？'）。")
    parser.add_argument('--chunk_size', type=int, default=400, help="文本块的大小。")
    parser.add_argument('--chunk_overlap', type=int, default=200, help="块之间的重叠。")
    parser.add_argument('--num_neighbors', type=int, default=1, help="用于上下文的相邻块数量。")
    return parser.parse_args()


# Main execution
if __name__ == "__main__":
    args = parse_args()

    # Initialize and run the RAG method
    rag_method = RAGMethod(chunk_size=args.chunk_size, chunk_overlap=args.chunk_overlap)
    baseline, enriched = rag_method.run(args.query, num_neighbors=args.num_neighbors)

    print("基线块:")
    print(baseline)

    print("\n增强块:")
    print(enriched)

调用方式:

 python context_enrichment_window_around_chunk.py --query "什么是人工智能？" --chunk_size 300 --num_neighbors 3